feat: add support for csrf check on websocket upgrade

fixes #943

Author: oliemansm

graphql-spring-boot-autoconfigure/src/main/java/graphql/kickstart/autoconfigure/web/servlet/DefaultWsCsrfToken.java

@@ -0,0 +1,20 @@
+package graphql.kickstart.autoconfigure.web.servlet;
+
+import lombok.RequiredArgsConstructor;
+
+@RequiredArgsConstructor
+class DefaultWsCsrfToken implements WsCsrfToken {
+
+  private final String token;
+  private final String parameterName;
+
+  @Override
+  public String getToken() {
+    return token;
+  }
+
+  @Override
+  public String getParameterName() {
+    return parameterName;
+  }
+}

graphql-spring-boot-autoconfigure/src/main/java/graphql/kickstart/autoconfigure/web/servlet/GraphQLWebsocketAutoConfiguration.java

@@ -26,6 +26,7 @@
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Conditional;
+import org.springframework.security.web.csrf.HttpSessionCsrfTokenRepository;
 import org.springframework.web.servlet.DispatcherServlet;
 import org.springframework.web.socket.server.standard.ServerEndpointExporter;
 import org.springframework.web.socket.server.standard.ServerEndpointRegistration;
@@ -78,10 +79,25 @@ private Optional<SubscriptionConnectionListener> keepAliveListener() {
     return Optional.empty();
   }
 
+  @Bean
+  public WsCsrfFilter wsCsrfFilter(
+      @Autowired(required = false) WsCsrfTokenRepository csrfTokenRepository) {
+    return new WsCsrfFilter(csrfTokenRepository);
+  }
+
+  @Bean
+  @ConditionalOnMissingBean
+  @ConditionalOnClass(HttpSessionCsrfTokenRepository.class)
+  public WsCsrfTokenRepository wsCsrfTokenRepository() {
+    return new WsSessionCsrfTokenRepository();
+  }
+
   @Bean
   @ConditionalOnClass(ServerContainer.class)
-  public ServerEndpointRegistration serverEndpointRegistration(GraphQLWebsocketServlet servlet) {
-    return new GraphQLWsServerEndpointRegistration(websocketProperties.getPath(), servlet);
+  public ServerEndpointRegistration serverEndpointRegistration(
+      GraphQLWebsocketServlet servlet, WsCsrfFilter csrfFilter) {
+    return new GraphQLWsServerEndpointRegistration(
+        websocketProperties.getPath(), servlet, csrfFilter);
   }
 
   @Bean

graphql-spring-boot-autoconfigure/src/main/java/graphql/kickstart/autoconfigure/web/servlet/GraphQLWsServerEndpointRegistration.java

@@ -7,15 +7,20 @@
 import org.springframework.context.Lifecycle;
 import org.springframework.web.socket.server.standard.ServerEndpointRegistration;
 
-/** @author Andrew Potter */
+/**
+ * @author Andrew Potter
+ */
 public class GraphQLWsServerEndpointRegistration extends ServerEndpointRegistration
     implements Lifecycle {
 
   private final GraphQLWebsocketServlet servlet;
+  private final WsCsrfFilter csrfFilter;
 
-  public GraphQLWsServerEndpointRegistration(String path, GraphQLWebsocketServlet servlet) {
+  public GraphQLWsServerEndpointRegistration(
+      String path, GraphQLWebsocketServlet servlet, WsCsrfFilter csrfFilter) {
     super(path, servlet);
     this.servlet = servlet;
+    this.csrfFilter = csrfFilter;
   }
 
   @Override
@@ -27,6 +32,7 @@ public boolean checkOrigin(String originHeaderValue) {
   public void modifyHandshake(
       ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {
     super.modifyHandshake(sec, request, response);
+    csrfFilter.doFilter(request);
     servlet.modifyHandshake(sec, request, response);
   }
 

graphql-spring-boot-autoconfigure/src/main/java/graphql/kickstart/autoconfigure/web/servlet/WsCsrfFilter.java

@@ -0,0 +1,35 @@
+package graphql.kickstart.autoconfigure.web.servlet;
+
+import static org.springframework.util.CollectionUtils.firstElement;
+
+import jakarta.websocket.server.HandshakeRequest;
+import java.util.Objects;
+import lombok.RequiredArgsConstructor;
+
+@RequiredArgsConstructor
+class WsCsrfFilter {
+
+  private final WsCsrfTokenRepository tokenRepository;
+
+  void doFilter(HandshakeRequest request) {
+    if (tokenRepository != null) {
+      WsCsrfToken csrfToken = tokenRepository.loadToken(request);
+      boolean missingToken = csrfToken == null;
+      if (missingToken) {
+        csrfToken = tokenRepository.generateToken(request);
+        tokenRepository.saveToken(csrfToken, request);
+      }
+
+      String actualToken =
+          firstElement(request.getParameterMap().get(csrfToken.getParameterName()));
+      if (!Objects.equals(csrfToken.getToken(), actualToken)) {
+        throw new IllegalStateException(
+            "Invalid CSRF Token '"
+                + actualToken
+                + "' was found on the request parameter '"
+                + csrfToken.getParameterName()
+                + "'.");
+      }
+    }
+  }
+}

graphql-spring-boot-autoconfigure/src/main/java/graphql/kickstart/autoconfigure/web/servlet/WsCsrfToken.java

@@ -0,0 +1,10 @@
+package graphql.kickstart.autoconfigure.web.servlet;
+
+import java.io.Serializable;
+
+public interface WsCsrfToken extends Serializable {
+
+    String getToken();
+
+    String getParameterName();
+}

graphql-spring-boot-autoconfigure/src/main/java/graphql/kickstart/autoconfigure/web/servlet/WsCsrfTokenRepository.java

@@ -0,0 +1,12 @@
+package graphql.kickstart.autoconfigure.web.servlet;
+
+import jakarta.websocket.server.HandshakeRequest;
+
+public interface WsCsrfTokenRepository {
+
+  WsCsrfToken loadToken(HandshakeRequest request);
+
+  WsCsrfToken generateToken(HandshakeRequest request);
+
+  void saveToken(WsCsrfToken csrfToken, HandshakeRequest request);
+}

graphql-spring-boot-autoconfigure/src/main/java/graphql/kickstart/autoconfigure/web/servlet/WsSessionCsrfTokenRepository.java

@@ -0,0 +1,42 @@
+package graphql.kickstart.autoconfigure.web.servlet;
+
+import jakarta.servlet.http.HttpSession;
+import jakarta.websocket.server.HandshakeRequest;
+import java.util.UUID;
+import org.springframework.security.web.csrf.HttpSessionCsrfTokenRepository;
+
+class WsSessionCsrfTokenRepository implements WsCsrfTokenRepository {
+
+  private static final String DEFAULT_CSRF_PARAMETER_NAME = "_csrf";
+
+  private static final String DEFAULT_CSRF_TOKEN_ATTR_NAME =
+      HttpSessionCsrfTokenRepository.class.getName().concat(".CSRF_TOKEN");
+
+  private String sessionAttributeName = DEFAULT_CSRF_TOKEN_ATTR_NAME;
+
+  @Override
+  public void saveToken(WsCsrfToken token, HandshakeRequest request) {
+    HttpSession session = (HttpSession) request.getHttpSession();
+    if (session != null) {
+      if (token == null) {
+        session.removeAttribute(this.sessionAttributeName);
+      } else {
+        session.setAttribute(this.sessionAttributeName, token);
+      }
+    }
+  }
+
+  @Override
+  public WsCsrfToken loadToken(HandshakeRequest request) {
+    HttpSession session = (HttpSession) request.getHttpSession();
+    if (session == null) {
+      return null;
+    }
+    return (WsCsrfToken) session.getAttribute(this.sessionAttributeName);
+  }
+
+  @Override
+  public WsCsrfToken generateToken(HandshakeRequest request) {
+    return new DefaultWsCsrfToken(UUID.randomUUID().toString(), DEFAULT_CSRF_PARAMETER_NAME);
+  }
+}