From 9074828478cd964e3e1e6a8f0e02739cc88fe055 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Fri, 18 May 2018 16:43:17 -0400
Subject: [PATCH] Add attributes to WebFlux WebSocketSession

Issue: SPR-16212
---
 .../web/reactive/socket/WebSocketSession.java         |  9 +++++++++
 .../socket/adapter/AbstractWebSocketSession.java      | 11 ++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/WebSocketSession.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/WebSocketSession.java
index 10fa28176d9..890067dcb62 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/WebSocketSession.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/WebSocketSession.java
@@ -15,6 +15,8 @@
  */
 package org.springframework.web.reactive.socket;
 
+import java.util.Collections;
+import java.util.Map;
 import java.util.function.Function;
 
 import org.reactivestreams.Publisher;
@@ -52,6 +54,13 @@ public interface WebSocketSession {
 	 */
 	DataBufferFactory bufferFactory();
 
+	/**
+	 * Return the map with attributes associated with the WebSocket session.
+	 * @return a Map with the session attributes (never {@code null})
+	 * @since 5.1
+	 */
+	Map<String, Object> getAttributes();
+
 	/**
 	 * Provides access to the stream of inbound messages.
 	 * <p>This stream receives a completion or error signal when the connection
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/AbstractWebSocketSession.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/AbstractWebSocketSession.java
index 46408b61f89..72dc7616cf0 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/AbstractWebSocketSession.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/AbstractWebSocketSession.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,8 @@
 package org.springframework.web.reactive.socket.adapter;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 
 import org.reactivestreams.Publisher;
@@ -48,6 +50,8 @@ public abstract class AbstractWebSocketSession<T> implements WebSocketSession {
 
 	private final DataBufferFactory bufferFactory;
 
+	private final Map<String, Object> attributes = new ConcurrentHashMap<>();
+
 
 	/**
 	 * Create a new instance and associate the given attributes with it.
@@ -86,6 +90,11 @@ public abstract class AbstractWebSocketSession<T> implements WebSocketSession {
 		return this.bufferFactory;
 	}
 
+	@Override
+	public Map<String, Object> getAttributes() {
+		return this.attributes;
+	}
+
 	@Override
 	public abstract Flux<WebSocketMessage> receive();