Add WebSocket Support

Fixes gh-35

spring-session/src/main/java/org/springframework/session/web/socket/config/annotation/AbstractSessionWebSocketMessageBrokerConfigurer.java

@@ -0,0 +1,145 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.session.web.socket.config.annotation;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.context.annotation.Bean;
+import org.springframework.messaging.simp.config.ChannelRegistration;
+import org.springframework.session.ExpiringSession;
+import org.springframework.session.Session;
+import org.springframework.session.SessionRepository;
+import org.springframework.session.web.socket.handler.WebSocketConnectHandlerDecoratorFactory;
+import org.springframework.session.web.socket.handler.WebSocketRegistryListener;
+import org.springframework.session.web.socket.server.SessionRepositoryMessageInterceptor;
+import org.springframework.web.socket.config.annotation.AbstractWebSocketMessageBrokerConfigurer;
+import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
+import org.springframework.web.socket.config.annotation.StompWebSocketEndpointRegistration;
+import org.springframework.web.socket.config.annotation.WebSocketTransportRegistration;
+import org.springframework.web.socket.server.HandshakeInterceptor;
+
+/**
+ * Eases configuration of Web Socket and Spring Session integration.
+ *
+ * <p>
+ * The configuration:
+ * </p>
+ * <ul>
+ * <li>Ensures the the {@link Session} is kept alive on incoming web socket
+ * messages.</li>
+ * <li>Ensures that Web Socket Sessions are destroyed when a {@link Session} is
+ * terminated</li>
+ * </ul>
+ *
+ * <p>Example usage</p>
+ *
+ * <code>
+ * @Configuration
+ * @EnableScheduling
+ * @EnableWebSocketMessageBroker
+ * public class WebSocketConfig<S extends ExpiringSession> extends AbstractSessionWebSocketMessageBrokerConfigurer<S> {
+ *
+ *     @Override
+ *     protected void configureStompEndpoints(StompEndpointRegistry registry) {
+ *         registry.addEndpoint("/messages")
+ *                 .withSockJS();
+ *     }
+ *
+ *     @Override
+ *     public void configureMessageBroker(MessageBrokerRegistry registry) {
+ *         registry.enableSimpleBroker("/queue/", "/topic/");
+ *         registry.setApplicationDestinationPrefixes("/app");
+ *     }
+ * }
+ * </code>
+ *
+ * @author Rob Winch
+ * @since 1.0
+ *
+ * @param <S>
+ *            the type of ExpiringSession
+ */
+public abstract class AbstractSessionWebSocketMessageBrokerConfigurer<S extends ExpiringSession> extends AbstractWebSocketMessageBrokerConfigurer {
+
+    @Autowired
+    @SuppressWarnings("rawtypes")
+    private SessionRepository sessionRepository;
+
+    @Autowired
+    private ApplicationEventPublisher eventPublisher;
+
+    @Override
+    public void configureClientInboundChannel(ChannelRegistration registration) {
+        registration.setInterceptors(sessionRepositoryInterceptor());
+    }
+
+    @Override
+    public final void registerStompEndpoints(StompEndpointRegistry registry) {
+        configureStompEndpoints(new SessionStompEndpointRegistry(registry, sessionRepositoryInterceptor()));
+    }
+
+
+    /**
+     * Register STOMP endpoints mapping each to a specific URL and (optionally)
+     * enabling and configuring SockJS fallback options with a
+     * {@link SessionRepositoryMessageInterceptor} automatically added as an
+     * interceptor.
+     *
+     * @param registry
+     *            the {@link StompEndpointRegistry} which automatically has a
+     *            {@link SessionRepositoryMessageInterceptor} added to it.
+     */
+    protected abstract void configureStompEndpoints(StompEndpointRegistry registry);
+
+    @Override
+    public void configureWebSocketTransport(
+            WebSocketTransportRegistration registration) {
+        registration.addDecoratorFactory(wsConnectHandlerDecoratorFactory());
+    }
+
+    @Bean
+    public WebSocketRegistryListener webSocketRegistryListener() {
+        return new WebSocketRegistryListener();
+    }
+
+    @Bean
+    public WebSocketConnectHandlerDecoratorFactory wsConnectHandlerDecoratorFactory() {
+        return new WebSocketConnectHandlerDecoratorFactory(eventPublisher);
+    }
+
+    @Bean
+    @SuppressWarnings("unchecked")
+    public SessionRepositoryMessageInterceptor<S> sessionRepositoryInterceptor() {
+        return new SessionRepositoryMessageInterceptor<S>(sessionRepository);
+    }
+
+    static class SessionStompEndpointRegistry implements StompEndpointRegistry {
+        private final StompEndpointRegistry registry;
+        private final HandshakeInterceptor interceptor;
+
+        public SessionStompEndpointRegistry(StompEndpointRegistry registry,
+                HandshakeInterceptor interceptor) {
+            this.registry = registry;
+            this.interceptor = interceptor;
+        }
+
+        public StompWebSocketEndpointRegistration addEndpoint(String... paths) {
+            StompWebSocketEndpointRegistration endpoints = registry.addEndpoint(paths);
+            endpoints.addInterceptors(interceptor);
+            return endpoints;
+        }
+    }
+}

spring-session/src/main/java/org/springframework/session/web/socket/events/SessionConnectEvent.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.session.web.socket.events;
+
+import org.springframework.context.ApplicationEvent;
+import org.springframework.session.web.socket.handler.WebSocketConnectHandlerDecoratorFactory;
+import org.springframework.session.web.socket.handler.WebSocketRegistryListener;
+import org.springframework.web.socket.WebSocketSession;
+
+/**
+ * Similar to Spring
+ * {@link org.springframework.web.socket.messaging.SessionConnectEvent} except
+ * that it provides access to the {@link WebSocketSession} to allow mapping the
+ * Spring Session to the {@link WebSocketSession}.
+ *
+ * @author Rob Winch
+ * @since 1.0
+ * @see WebSocketRegistryListener
+ * @see WebSocketConnectHandlerDecoratorFactory
+ */
+@SuppressWarnings("serial")
+public class SessionConnectEvent extends ApplicationEvent {
+
+    private final WebSocketSession webSocketSession;
+
+    public SessionConnectEvent(Object source, WebSocketSession webSocketSession) {
+        super(source);
+        this.webSocketSession = webSocketSession;
+    }
+
+    public WebSocketSession getWebSocketSession() {
+        return webSocketSession;
+    }
+}

spring-session/src/main/java/org/springframework/session/web/socket/handler/WebSocketConnectHandlerDecoratorFactory.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.session.web.socket.handler;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.context.ApplicationEvent;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.session.Session;
+import org.springframework.session.web.socket.events.SessionConnectEvent;
+import org.springframework.util.Assert;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.WebSocketSession;
+import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
+import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;
+
+/**
+ * Ensures that a {@link SessionConnectEvent} is published in
+ * {@link WebSocketHandler#afterConnectionEstablished(WebSocketSession)}. This
+ * is necessary so that the {@link WebSocketSession} can be mapped to the
+ * corresponding Spring {@link Session} to terminate any
+ * {@link WebSocketSession} associated with a Spring {@link Session} that was
+ * destroyed.
+ *
+ * @author Rob Winch
+ * @since 1.0
+ *
+ * @see WebSocketRegistryListener
+ */
+public final class WebSocketConnectHandlerDecoratorFactory implements WebSocketHandlerDecoratorFactory {
+
+    private static final Log logger = LogFactory.getLog(WebSocketConnectHandlerDecoratorFactory.class);
+
+    private final ApplicationEventPublisher eventPublisher;
+
+    /**
+     * Creates a new instance
+     *
+     * @param eventPublisher the {@link ApplicationEventPublisher} to use. Cannot be null.
+     */
+    public WebSocketConnectHandlerDecoratorFactory(
+            ApplicationEventPublisher eventPublisher) {
+        Assert.notNull(eventPublisher, "eventPublisher cannot be null");
+        this.eventPublisher = eventPublisher;
+    }
+
+    @Override
+    public WebSocketHandler decorate(WebSocketHandler handler) {
+        return new SessionWebSocketHandler(handler);
+    }
+
+    private final class SessionWebSocketHandler extends WebSocketHandlerDecorator {
+
+        public SessionWebSocketHandler(WebSocketHandler delegate) {
+            super(delegate);
+        }
+
+        @Override
+        public void afterConnectionEstablished(WebSocketSession wsSession)
+                throws Exception {
+            super.afterConnectionEstablished(wsSession);
+
+            publishEvent(new SessionConnectEvent(this,wsSession));
+        }
+
+        private void publishEvent(ApplicationEvent event) {
+            try {
+                eventPublisher.publishEvent(event);
+            }
+            catch (Throwable ex) {
+                logger.error("Error publishing " + event + ".", ex);
+            }
+        }
+    }
+}

spring-session/src/main/java/org/springframework/session/web/socket/handler/WebSocketRegistryListener.java

@@ -0,0 +1,124 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.session.web.socket.handler;
+
+import java.io.IOException;
+import java.security.Principal;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.context.ApplicationEvent;
+import org.springframework.context.ApplicationListener;
+import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
+import org.springframework.session.events.SessionDestroyedEvent;
+import org.springframework.session.web.socket.events.SessionConnectEvent;
+import org.springframework.session.web.socket.server.SessionRepositoryMessageInterceptor;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.WebSocketSession;
+import org.springframework.web.socket.messaging.SessionDisconnectEvent;
+
+/**
+ * <p>
+ * Keeps track of mapping the Spring Session ID to the {@link WebSocketSession}
+ * and ensuring when a {@link SessionDestroyedEvent} is fired that the
+ * {@link WebSocketSession} is closed.
+ * </p>
+ *
+ *
+ * @author Rob Winch
+ * @since 1.0
+ */
+public final class WebSocketRegistryListener implements ApplicationListener<ApplicationEvent> {
+
+    private static final Log logger = LogFactory.getLog(WebSocketRegistryListener.class);
+
+    static final CloseStatus SESSION_EXPIRED_STATUS = new CloseStatus(CloseStatus.POLICY_VIOLATION.getCode(),
+            "This connection was established under an authenticated HTTP Session that has expired");
+
+    private final ConcurrentHashMap<String,Map<String,WebSocketSession>> httpSessionIdToWsSessions = new ConcurrentHashMap<String,Map<String,WebSocketSession>>();
+
+    @Override
+    public void onApplicationEvent(ApplicationEvent event) {
+        if(event instanceof SessionDestroyedEvent) {
+            SessionDestroyedEvent e = (SessionDestroyedEvent) event;
+            closeWsSessions(e.getSessionId());
+        } else if(event instanceof SessionConnectEvent) {
+            SessionConnectEvent e = (SessionConnectEvent) event;
+            afterConnectionEstablished(e.getWebSocketSession());
+        } else if(event instanceof SessionDisconnectEvent) {
+            SessionDisconnectEvent e = (SessionDisconnectEvent) event;
+            Map<String, Object> sessionAttributes = SimpMessageHeaderAccessor.getSessionAttributes(e.getMessage().getHeaders());
+            String httpSessionId =  sessionAttributes == null ? null : SessionRepositoryMessageInterceptor.getSessionId(sessionAttributes);
+            afterConnectionClosed(httpSessionId, e.getSessionId());
+        }
+    }
+
+    private void afterConnectionEstablished(WebSocketSession wsSession) {
+        Principal principal = wsSession.getPrincipal();
+        if(principal == null) {
+            return;
+        }
+
+        String id = getHttpSessionId(wsSession);
+        registerWsSession(id, wsSession);
+    }
+
+    private String getHttpSessionId(WebSocketSession wsSession) {
+        Map<String, Object> attributes = wsSession.getAttributes();
+        return SessionRepositoryMessageInterceptor.getSessionId(attributes);
+    }
+
+    private void afterConnectionClosed(String httpSessionId, String wsSessionId) {
+
+        Map<String,WebSocketSession> sessions = httpSessionIdToWsSessions.get(httpSessionId);
+        if(sessions != null) {
+            boolean result = sessions.remove(wsSessionId) != null;
+            if(logger.isDebugEnabled()) {
+                logger.debug("Removal of " + wsSessionId + " was " + result);
+            }
+        }
+    }
+
+    private void registerWsSession(String sessionId, WebSocketSession wsSession) {
+        Map<String,WebSocketSession> sessions = httpSessionIdToWsSessions.get(sessionId);
+        if(sessions == null) {
+            sessions =
+                    new ConcurrentHashMap<String,WebSocketSession>();
+            httpSessionIdToWsSessions.putIfAbsent(sessionId, sessions);
+            sessions = httpSessionIdToWsSessions.get(sessionId);
+        }
+        sessions.put(wsSession.getId(), wsSession);
+    }
+
+    private void closeWsSessions(String sessionId) {
+        Map<String,WebSocketSession> sessionsToClose = httpSessionIdToWsSessions.remove(sessionId);
+        if(sessionsToClose == null) {
+            return;
+        }
+        if(logger.isDebugEnabled()) {
+            logger.debug("Closing WebSocket connections associated to expired HTTP Session " + sessionId);
+        }
+        for(WebSocketSession toClose : sessionsToClose.values()) {
+            try {
+                toClose.close(SESSION_EXPIRED_STATUS);
+            } catch (IOException e) {
+                logger.debug("Failed to close WebSocketSession (this is nothing to worry about but for debugging only)",e);
+            }
+        }
+    }
+}

spring-session/src/main/java/org/springframework/session/web/socket/server/SessionRepositoryMessageInterceptor.java

@@ -0,0 +1,161 @@
+/*
+ * Copyright 2002-2014 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.session.web.socket.server;
+
+import java.util.EnumSet;
+import java.util.Map;
+import java.util.Set;
+
+import javax.servlet.http.HttpSession;
+
+import org.springframework.http.server.ServerHttpRequest;
+import org.springframework.http.server.ServerHttpResponse;
+import org.springframework.http.server.ServletServerHttpRequest;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
+import org.springframework.messaging.simp.SimpMessageType;
+import org.springframework.messaging.support.ChannelInterceptor;
+import org.springframework.messaging.support.ChannelInterceptorAdapter;
+import org.springframework.session.ExpiringSession;
+import org.springframework.session.Session;
+import org.springframework.session.SessionRepository;
+import org.springframework.util.Assert;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.server.HandshakeInterceptor;
+
+/**
+ * <p>
+ * Acts as a {@link ChannelInterceptor} and a {@link HandshakeInterceptor} to
+ * ensure the {@link ExpiringSession#getLastAccessedTime()} is up to date.
+ * </p>
+ * <ul>
+ * <li>
+ * Associates the {@link Session#getId()} with the WebSocket Session
+ * attributes when the handshake is performed. This is later used when
+ * intercepting messages to ensure the
+ * {@link ExpiringSession#getLastAccessedTime()} is updated.
+ * </li>
+ * </li>
+ * <li>
+ * Intercepts {@link Message}'s that are have {@link SimpMessageType} that
+ * corresponds to {@link #setMatchingMessageTypes(Set)} and updates the last
+ * accessed time of the {@link Session}. If the {@link Session} is expired, the
+ * {@link Message} is prevented from proceeding.</li>
+ * </ul>
+ *
+ * <p>
+ * In order to work {@link SessionRepositoryMessageInterceptor} must be
+ * registered as a {@link ChannelInterceptor} and a {@link HandshakeInterceptor}
+ * .
+ * </p>
+ *
+ * @author Rob Winch
+ * @since 1.0
+ */
+public final class SessionRepositoryMessageInterceptor<S extends ExpiringSession> extends ChannelInterceptorAdapter
+    implements HandshakeInterceptor {
+
+    private static final String SPRING_SESSION_ID_ATTR_NAME = "SPRING.SESSION.ID";
+
+    private final SessionRepository<S> sessionRepository;
+
+    private Set<SimpMessageType> matchingMessageTypes;
+
+    /**
+     * Creates a new instance
+     *
+     * @param sessionRepository the {@link SessionRepository} to use. Cannot be null.
+     */
+    public SessionRepositoryMessageInterceptor(SessionRepository<S> sessionRepository) {
+        Assert.notNull(sessionRepository, "sessionRepository cannot be null");
+        this.sessionRepository = sessionRepository;
+        this.matchingMessageTypes = EnumSet.of(SimpMessageType.CONNECT, SimpMessageType.MESSAGE, SimpMessageType.SUBSCRIBE, SimpMessageType.UNSUBSCRIBE);
+    }
+
+    /**
+     * <p>
+     * Sets the {@link SimpMessageType} to match on. If the {@link Message}
+     * matches, then {@link #preSend(Message, MessageChannel)} ensures the
+     * {@link Session} is not expired and updates the
+     * {@link ExpiringSession#getLastAccessedTime()}
+     * </p>
+     *
+     * <p>
+     * The default is: SimpMessageType.CONNECT, SimpMessageType.MESSAGE,
+     * SimpMessageType.SUBSCRIBE, SimpMessageType.UNSUBSCRIBE.
+     * </p>
+     *
+     * @param matchingMessageTypes
+     *            the {@link SimpMessageType} to match on in
+     *            {@link #preSend(Message, MessageChannel)}, else the
+     *            {@link Message} is continued without accessing or updating the
+     *            {@link Session}
+     */
+    public void setMatchingMessageTypes(Set<SimpMessageType> matchingMessageTypes) {
+        Assert.notEmpty(matchingMessageTypes,"matchingMessageTypes cannot be null or empty");
+        this.matchingMessageTypes = matchingMessageTypes;
+    }
+
+    @Override
+    public Message<?> preSend(Message<?> message, MessageChannel channel) {
+        if(message == null) {
+            return message;
+        }
+        SimpMessageType messageType = SimpMessageHeaderAccessor.getMessageType(message.getHeaders());
+        if(!this.matchingMessageTypes.contains(messageType)) {
+            return super.preSend(message, channel);
+        }
+        Map<String, Object> sessionHeaders = SimpMessageHeaderAccessor.getSessionAttributes(message.getHeaders());
+        String sessionId = sessionHeaders == null ? null : (String) sessionHeaders.get(SPRING_SESSION_ID_ATTR_NAME);
+        if (sessionId != null) {
+            S session = sessionRepository.getSession(sessionId);
+            if (session != null) {
+                // update the last accessed time
+                sessionRepository.save(session);
+            }
+        }
+        return super.preSend(message, channel);
+    }
+
+    @Override
+    public boolean beforeHandshake(ServerHttpRequest request,
+            ServerHttpResponse response, WebSocketHandler wsHandler,
+            Map<String, Object> attributes) throws Exception {
+        if (request instanceof ServletServerHttpRequest) {
+            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
+            HttpSession session = servletRequest.getServletRequest().getSession(false);
+            if (session != null) {
+                setSessionId(attributes, session.getId());
+            }
+        }
+        return true;
+    }
+
+    @Override
+    public void afterHandshake(ServerHttpRequest request,
+            ServerHttpResponse response, WebSocketHandler wsHandler,
+            Exception exception) {
+    }
+
+    public static String getSessionId(Map<String, Object> attributes) {
+        return (String) attributes.get(SPRING_SESSION_ID_ATTR_NAME);
+    }
+
+    public static void setSessionId(Map<String, Object> attributes, String sessionId) {
+        attributes.put(SPRING_SESSION_ID_ATTR_NAME, sessionId);
+    }
+}

spring-session/src/test/java/org/springframework/session/web/socket/handler/WebSocketConnectHandlerDecoratorFactoryTests.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.session.web.socket.handler;
+
+import static org.fest.assertions.Assertions.assertThat;
+import static org.mockito.Mockito.*;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+import org.springframework.context.ApplicationEvent;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.session.web.socket.events.SessionConnectEvent;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.WebSocketSession;
+
+@RunWith(MockitoJUnitRunner.class)
+public class WebSocketConnectHandlerDecoratorFactoryTests {
+    @Mock
+    ApplicationEventPublisher eventPublisher;
+    @Mock
+    WebSocketHandler delegate;
+    @Mock
+    WebSocketSession session;
+    @Captor
+    ArgumentCaptor<SessionConnectEvent> event;
+
+    WebSocketConnectHandlerDecoratorFactory factory;
+
+    @Before
+    public void setup() {
+        factory = new WebSocketConnectHandlerDecoratorFactory(eventPublisher);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void constructorNullEventPublisher() {
+        new WebSocketConnectHandlerDecoratorFactory(null);
+    }
+
+    @Test
+    public void decorateAfterConnectionEstablished() throws Exception {
+        WebSocketHandler decorated = factory.decorate(delegate);
+
+        decorated.afterConnectionEstablished(session);
+
+        verify(eventPublisher).publishEvent(event.capture());
+        assertThat(event.getValue().getWebSocketSession()).isSameAs(session);
+    }
+
+    @Test
+    public void decorateAfterConnectionEstablishedEventError() throws Exception {
+        WebSocketHandler decorated = factory.decorate(delegate);
+        doThrow(new IllegalStateException("Test throw on publishEvent")).when(eventPublisher).publishEvent(any(ApplicationEvent.class));
+
+        decorated.afterConnectionEstablished(session);
+
+        verify(eventPublisher).publishEvent(any(SessionConnectEvent.class));
+    }
+}

spring-session/src/test/java/org/springframework/session/web/socket/handler/WebSocketRegistryListenerTests.java

@@ -0,0 +1,130 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.session.web.socket.handler;
+
+import static org.mockito.Mockito.*;
+
+import java.security.Principal;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
+import org.springframework.session.events.SessionDestroyedEvent;
+import org.springframework.session.web.socket.events.SessionConnectEvent;
+import org.springframework.session.web.socket.server.SessionRepositoryMessageInterceptor;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.WebSocketSession;
+import org.springframework.web.socket.messaging.SessionDisconnectEvent;
+
+@RunWith(MockitoJUnitRunner.class)
+public class WebSocketRegistryListenerTests {
+    @Mock
+    WebSocketSession wsSession;
+    @Mock
+    WebSocketSession wsSession2;
+    @Mock
+    Message<byte[]> message;
+    @Mock
+    Principal principal;
+
+    SessionConnectEvent connect;
+
+    SessionConnectEvent connect2;
+
+    SessionDisconnectEvent disconnect;
+
+    SessionDestroyedEvent destroyed;
+
+    Map<String, Object> attributes;
+
+    String sessionId;
+
+    WebSocketRegistryListener listener;
+
+
+    @Before
+    public void setup() {
+        sessionId = "session-id";
+        attributes = new HashMap<>();
+        SessionRepositoryMessageInterceptor.setSessionId(attributes, sessionId);
+
+        when(wsSession.getAttributes()).thenReturn(attributes);
+        when(wsSession.getPrincipal()).thenReturn(principal);
+        when(wsSession.getId()).thenReturn("wsSession-id");
+
+        when(wsSession2.getAttributes()).thenReturn(attributes);
+        when(wsSession2.getPrincipal()).thenReturn(principal);
+        when(wsSession2.getId()).thenReturn("wsSession-id2");
+
+        Map<String,Object> headers = new HashMap<>();
+        headers.put(SimpMessageHeaderAccessor.SESSION_ATTRIBUTES, attributes);
+        when(message.getHeaders()).thenReturn(new MessageHeaders(headers));
+
+        listener = new WebSocketRegistryListener();
+        connect = new SessionConnectEvent(listener,wsSession);
+        connect2 = new SessionConnectEvent(listener,wsSession2);
+        disconnect = new SessionDisconnectEvent(listener, message, wsSession.getId(), CloseStatus.NORMAL);
+        destroyed = new SessionDestroyedEvent(listener, sessionId);
+    }
+
+    @Test
+    public void onApplicationEventConnectSessionDestroyed() throws Exception {
+        listener.onApplicationEvent(connect);
+
+        listener.onApplicationEvent(destroyed);
+
+        verify(wsSession).close(WebSocketRegistryListener.SESSION_EXPIRED_STATUS);
+    }
+
+    @Test
+    public void onApplicationEventConnectSessionDestroyedNullPrincipal() throws Exception {
+        when(wsSession.getPrincipal()).thenReturn(null);
+        listener.onApplicationEvent(connect);
+
+        listener.onApplicationEvent(destroyed);
+
+        verify(wsSession,times(0)).close(any(CloseStatus.class));
+    }
+
+    @Test
+    public void onApplicationEventConnectDisonnect() throws Exception {
+        listener.onApplicationEvent(connect);
+        listener.onApplicationEvent(disconnect);
+
+        listener.onApplicationEvent(destroyed);
+
+        verify(wsSession,times(0)).close(any(CloseStatus.class));
+    }
+
+    @Test
+    public void onApplicationEventConnectConnectDisonnect() throws Exception {
+        listener.onApplicationEvent(connect);
+        listener.onApplicationEvent(connect2);
+        listener.onApplicationEvent(disconnect);
+
+        listener.onApplicationEvent(destroyed);
+
+        verify(wsSession2).close(WebSocketRegistryListener.SESSION_EXPIRED_STATUS);
+        verify(wsSession,times(0)).close(any(CloseStatus.class));
+    }
+}

spring-session/src/test/java/org/springframework/session/web/socket/server/SessionRepositoryMessageInterceptorTests.java

@@ -0,0 +1,256 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.session.web.socket.server;
+
+import static org.fest.assertions.Assertions.assertThat;
+import static org.mockito.Matchers.*;
+import static org.mockito.Mockito.*;
+
+import java.util.Collections;
+import java.util.EnumSet;
+import java.util.HashMap;
+import java.util.Map;
+
+import javax.servlet.http.HttpSession;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+import org.springframework.http.server.ServletServerHttpRequest;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
+import org.springframework.messaging.simp.SimpMessageType;
+import org.springframework.messaging.support.MessageBuilder;
+import org.springframework.mock.web.MockHttpServletRequest;
+import org.springframework.session.ExpiringSession;
+import org.springframework.session.SessionRepository;
+
+@RunWith(MockitoJUnitRunner.class)
+public class SessionRepositoryMessageInterceptorTests {
+    @Mock
+    SessionRepository<ExpiringSession> sessionRepository;
+    @Mock
+    MessageChannel channel;
+    @Mock
+    ExpiringSession session;
+
+    Message<?> createMessage;
+
+    SimpMessageHeaderAccessor headers;
+
+    SessionRepositoryMessageInterceptor<ExpiringSession> interceptor;
+
+    @Before
+    public void setup() {
+        interceptor = new SessionRepositoryMessageInterceptor<ExpiringSession>(sessionRepository);
+        headers = SimpMessageHeaderAccessor.create();
+        headers.setSessionId("session");
+        headers.setSessionAttributes(new HashMap<String,Object>());
+        setMessageType(SimpMessageType.MESSAGE);
+        String sessionId = "http-session";
+        setSessionId(sessionId);
+        when(sessionRepository.getSession(sessionId)).thenReturn(session);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void preSendconstructorNullRepository() {
+        new SessionRepositoryMessageInterceptor<ExpiringSession>(null);
+    }
+
+    @Test
+    public void preSendNullMessage() {
+        assertThat(interceptor.preSend(null, channel)).isNull();
+    }
+
+    @Test
+    public void preSendConnectAckDoesNotInvokeSessionRepository() {
+        setMessageType(SimpMessageType.CONNECT_ACK);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verifyZeroInteractions(sessionRepository);
+    }
+
+    @Test
+    public void preSendHeartbeatDoesNotInvokeSessionRepository() {
+        setMessageType(SimpMessageType.HEARTBEAT);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verifyZeroInteractions(sessionRepository);
+    }
+
+    @Test
+    public void preSendDisconnectDoesNotInvokeSessionRepository() {
+        setMessageType(SimpMessageType.DISCONNECT);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verifyZeroInteractions(sessionRepository);
+    }
+
+    @Test
+    public void preSendOtherDoesNotInvokeSessionRepository() {
+        setMessageType(SimpMessageType.OTHER);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verifyZeroInteractions(sessionRepository);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void setMatchingMessageTypesNull() {
+        interceptor.setMatchingMessageTypes(null);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void setMatchingMessageTypesEmpty() {
+        interceptor.setMatchingMessageTypes(Collections.<SimpMessageType>emptySet());
+    }
+
+    @Test
+    public void preSendSetMatchingMessageTypes() {
+        interceptor.setMatchingMessageTypes(EnumSet.of(SimpMessageType.DISCONNECT));
+        setMessageType(SimpMessageType.DISCONNECT);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verify(sessionRepository).getSession(anyString());
+        verify(sessionRepository).save(session);
+    }
+
+    @Test
+    public void preSendConnectUpdatesLastUpdateTime() {
+        setMessageType(SimpMessageType.CONNECT);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verify(sessionRepository).getSession(anyString());
+        verify(sessionRepository).save(session);
+    }
+
+    @Test
+    public void preSendMessageUpdatesLastUpdateTime() {
+        setMessageType(SimpMessageType.MESSAGE);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verify(sessionRepository).getSession(anyString());
+        verify(sessionRepository).save(session);
+    }
+
+    @Test
+    public void preSendSubscribeUpdatesLastUpdateTime() {
+        setMessageType(SimpMessageType.SUBSCRIBE);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verify(sessionRepository).getSession(anyString());
+        verify(sessionRepository).save(session);
+    }
+
+    @Test
+    public void preSendUnsubscribeUpdatesLastUpdateTime() {
+        setMessageType(SimpMessageType.UNSUBSCRIBE);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verify(sessionRepository).getSession(anyString());
+        verify(sessionRepository).save(session);
+    }
+
+    // This will updated when SPR-12288 is resolved
+    @Test
+    public void preSendExpiredSession() {
+        setSessionId("expired");
+
+        interceptor.preSend(createMessage(), channel);
+
+        verify(sessionRepository,times(0)).save(any(ExpiringSession.class));
+    }
+
+    @Test
+    public void preSendNullSessionId() {
+        setSessionId(null);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verifyZeroInteractions(sessionRepository);
+    }
+
+    @Test
+    public void preSendNullSessionAttributes() {
+        headers.setSessionAttributes(null);
+
+        assertThat(interceptor.preSend(createMessage(), channel)).isSameAs(createMessage);
+
+        verifyZeroInteractions(sessionRepository);
+    }
+
+    @Test
+    public void beforeHandshakeNotServletServerHttpRequest() throws Exception {
+        assertThat(interceptor.beforeHandshake(null,null,null,null)).isTrue();
+
+        verifyZeroInteractions(sessionRepository);
+    }
+
+    @Test
+    public void beforeHandshakeNullSession() throws Exception {
+        ServletServerHttpRequest request = new ServletServerHttpRequest(new MockHttpServletRequest());
+        assertThat(interceptor.beforeHandshake(request,null,null,null)).isTrue();
+
+        verifyZeroInteractions(sessionRepository);
+    }
+
+    @Test
+    public void beforeHandshakeSession() throws Exception {
+        MockHttpServletRequest httpRequest = new MockHttpServletRequest();
+        HttpSession httpSession = httpRequest.getSession();
+        ServletServerHttpRequest request = new ServletServerHttpRequest(httpRequest);
+        Map<String,Object> attributes = new HashMap<String,Object>();
+
+        assertThat(interceptor.beforeHandshake(request,null,null,attributes)).isTrue();
+
+        assertThat(attributes.size()).isEqualTo(1);
+        assertThat(SessionRepositoryMessageInterceptor.getSessionId(attributes)).isEqualTo(httpSession.getId());
+    }
+
+    /**
+     * At the moment there is no need for afterHandshake to do anything.
+     */
+    @Test
+    public void afterHandshakeDoesNothing() {
+        interceptor.afterHandshake(null,null,null,null);
+
+        verifyZeroInteractions(sessionRepository);
+    }
+
+    private void setSessionId(String id) {
+        SessionRepositoryMessageInterceptor.setSessionId(headers.getSessionAttributes(), id);
+    }
+
+    private Message<?> createMessage() {
+        createMessage = MessageBuilder.createMessage("", headers.getMessageHeaders());
+        return createMessage;
+    }
+
+    private void setMessageType(SimpMessageType type) {
+        headers.setHeader(SimpMessageHeaderAccessor.MESSAGE_TYPE_HEADER, type);
+    }
+}