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);
+    }
+}