HttpSessionPuntingHandshakeInterceptor.java

import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import lombok.val;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.lang.Nullable;
import org.springframework.security.core.session.SessionInformation;
import org.springframework.security.core.session.SessionRegistry;
import org.springframework.session.SessionRepository;
import org.springframework.util.Assert;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

import javax.servlet.http.HttpSession;
import java.util.Collection;
import java.util.Collections;


/**
 * On handshake success this interceptor forcibly deletes all other http sessions belonging to the current user. If this
 * interceptor is used with a redis session store (properly configured to broadcast keyspace events) then all websocket
 * connections belonging to the current user but opened in a different session will be closed.
 *
 * This differs from default framework behavior, where websocket connections will linger in a potentially open state
 * until another http request is made or the cron task deletes the expired session.
 *
 * If max session concurrency is set to n=1, this behavior could be used to enforce near real-time termination of
 * websocket connections in the event of a concurrency violation.
 */
@Slf4j
public class HttpSessionPuntingHandshakeInterceptor extends HttpSessionHandshakeInterceptor implements ApplicationContextAware {
    @Setter
    private ApplicationContext applicationContext;

    public HttpSessionPuntingHandshakeInterceptor() {
        this(Collections.emptyList());
    }

    public HttpSessionPuntingHandshakeInterceptor(Collection<String> attributeNames) {
        super(attributeNames);
        setCreateSession(true);
    }


    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, @Nullable Exception exception) {
        if (exception != null) {
            return;
        }
        val httpSessionId = getHttpSessionId(request);
        val principal = request.getPrincipal();
        if (httpSessionId == null || principal == null) {
            return;
        }
        applicationContext.getBean(SessionRegistry.class).getAllSessions(principal.getName(), true)
            .stream()
            .filter(sessionInformation -> !sessionInformation.getSessionId().equals(httpSessionId))
            .forEach(this::puntSession);
    }

    protected void puntSession(SessionInformation sessionInformation) {
        val id = sessionInformation.getSessionId();
        log.debug("HttpSession(id={}) punt requested", id);
        try {
            applicationContext.getBean(SessionRepository.class).deleteById(id);
            log.debug("HttpSession(id={}) punt succeeded", id);
        } catch (Exception e) {
            log.warn(String.format("HttpSession(id=%s) punt failed", id), e);
        }
    }

    protected String getHttpSessionId(ServerHttpRequest request) {
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest serverRequest = (ServletServerHttpRequest) request;
            HttpSession session = serverRequest.getServletRequest().getSession(false);
            Assert.notNull(session, "session should already have been created here");
            return session.getId();
        }
        return null;
    }
}