Remove session on 4xx response from WebSocket handshake #25608

spring-projects · Aug 24, 2020 · f0fe2f1 · f0fe2f1
1 parent 1891f8a
commit f0fe2f1
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/.../java/org/springframework/web/socket/sockjs/transport/TransportHandlingSockJsService.java b/.../java/org/springframework/web/socket/sockjs/transport/TransportHandlingSockJsService.java
@@ -34,6 +34,7 @@
 import org.springframework.http.HttpStatus;
 import org.springframework.http.server.ServerHttpRequest;
 import org.springframework.http.server.ServerHttpResponse;
+import org.springframework.http.server.ServletServerHttpResponse;
 import org.springframework.lang.Nullable;
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.util.Assert;
@@ -270,6 +271,7 @@ else if (transportType.supportsCors()) {
 			}
 
 			SockJsSession session = this.sessions.get(sessionId);
+			boolean newSession = false;
 			if (session == null) {
 				if (transportHandler instanceof SockJsSessionFactory) {
 					Map<String, Object> attributes = new HashMap<>();
@@ -278,6 +280,7 @@ else if (transportType.supportsCors()) {
 					}
 					SockJsSessionFactory sessionFactory = (SockJsSessionFactory) transportHandler;
 					session = createSockJsSession(sessionId, sessionFactory, handler, attributes);
+					newSession = true;
 				}
 				else {
 					response.setStatusCode(HttpStatus.NOT_FOUND);
@@ -311,6 +314,14 @@ else if (transportType.supportsCors()) {
 			}
 
 			transportHandler.handleRequest(request, response, handler, session);
+
+			if (newSession && (response instanceof ServletServerHttpResponse)) {
+				int status = ((ServletServerHttpResponse) response).getServletResponse().getStatus();
+				if (HttpStatus.valueOf(status).is4xxClientError()) {
+					sessions.remove(sessionId);
+				}
+			}
+
 			chain.applyAfterHandshake(request, response, null);
 		}
 		catch (SockJsException ex) {