fix(#516): add origin check to websockets

Author: oliemansm

graphql-java-servlet/src/main/java/graphql/kickstart/servlet/GraphQLConfiguration.java

@@ -36,6 +36,7 @@ public class GraphQLConfiguration {
   private final ContextSetting contextSetting;
   private final GraphQLResponseCacheManager responseCacheManager;
   @Getter private final Executor asyncExecutor;
+  @Getter private final List<String> allowedOrigins;
   private HttpRequestHandler requestHandler;
 
   private GraphQLConfiguration(
@@ -49,9 +50,11 @@ private GraphQLConfiguration(
       ContextSetting contextSetting,
       Supplier<BatchInputPreProcessor> batchInputPreProcessor,
       GraphQLResponseCacheManager responseCacheManager,
-      Executor asyncExecutor) {
+      Executor asyncExecutor,
+      List<String> allowedOrigins) {
     this.invocationInputFactory = invocationInputFactory;
     this.asyncExecutor = asyncExecutor;
+    this.allowedOrigins = allowedOrigins;
     this.graphQLInvoker = graphQLInvoker != null ? graphQLInvoker : queryInvoker.toGraphQLInvoker();
     this.objectMapper = objectMapper;
     this.listeners = listeners;
@@ -148,6 +151,7 @@ public static class Builder {
     private int asyncMaxPoolSize = 200;
     private Executor asyncExecutor;
     private AsyncTaskDecorator asyncTaskDecorator;
+    private List<String> allowedOrigins = new ArrayList<>();
 
     private Builder(GraphQLInvocationInputFactory.Builder invocationInputFactoryBuilder) {
       this.invocationInputFactoryBuilder = invocationInputFactoryBuilder;
@@ -249,6 +253,13 @@ public Builder with(AsyncTaskDecorator asyncTaskDecorator) {
       return this;
     }
 
+    public Builder allowedOrigins(List<String> allowedOrigins) {
+      if (allowedOrigins != null) {
+        this.allowedOrigins.addAll(allowedOrigins);
+      }
+      return this;
+    }
+
     private Executor getAsyncExecutor() {
       if (asyncExecutor != null) {
         return asyncExecutor;
@@ -279,7 +290,8 @@ public GraphQLConfiguration build() {
           contextSetting,
           batchInputPreProcessorSupplier,
           responseCacheManager,
-          getAsyncTaskExecutor());
+          getAsyncTaskExecutor(),
+          allowedOrigins);
     }
   }
 }

graphql-java-servlet/src/main/java/graphql/kickstart/servlet/GraphQLWebsocketServlet.java

@@ -1,6 +1,7 @@
 package graphql.kickstart.servlet;
 
 import static java.util.Arrays.asList;
+import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 import static java.util.stream.Collectors.toList;
 
@@ -65,6 +66,7 @@ public class GraphQLWebsocketServlet extends Endpoint {
   private final AtomicBoolean isShuttingDown = new AtomicBoolean(false);
   private final AtomicBoolean isShutDown = new AtomicBoolean(false);
   private final Object cacheLock = new Object();
+  private final List<String> allowedOrigins;
 
   public GraphQLWebsocketServlet(GraphQLConfiguration configuration) {
     this(configuration, null);
@@ -77,21 +79,23 @@ public GraphQLWebsocketServlet(
         configuration.getGraphQLInvoker(),
         configuration.getInvocationInputFactory(),
         configuration.getObjectMapper(),
-        connectionListeners);
+        connectionListeners,
+        configuration.getAllowedOrigins());
   }
 
   public GraphQLWebsocketServlet(
       GraphQLInvoker graphQLInvoker,
       GraphQLSubscriptionInvocationInputFactory invocationInputFactory,
       GraphQLObjectMapper graphQLObjectMapper) {
-    this(graphQLInvoker, invocationInputFactory, graphQLObjectMapper, null);
+    this(graphQLInvoker, invocationInputFactory, graphQLObjectMapper, null, emptyList());
   }
 
   public GraphQLWebsocketServlet(
       GraphQLInvoker graphQLInvoker,
       GraphQLSubscriptionInvocationInputFactory invocationInputFactory,
       GraphQLObjectMapper graphQLObjectMapper,
-      Collection<SubscriptionConnectionListener> connectionListeners) {
+      Collection<SubscriptionConnectionListener> connectionListeners,
+      List<String> allowedOrigins) {
     List<ApolloSubscriptionConnectionListener> listeners = new ArrayList<>();
     if (connectionListeners != null) {
       connectionListeners.stream()
@@ -114,12 +118,10 @@ public GraphQLWebsocketServlet(
                 Stream.of(fallbackSubscriptionProtocolFactory))
             .map(SubscriptionProtocolFactory::getProtocol)
             .collect(toList());
+    this.allowedOrigins = allowedOrigins;
   }
 
   public GraphQLWebsocketServlet(
-      GraphQLInvoker graphQLInvoker,
-      GraphQLSubscriptionInvocationInputFactory invocationInputFactory,
-      GraphQLObjectMapper graphQLObjectMapper,
       List<SubscriptionProtocolFactory> subscriptionProtocolFactory,
       SubscriptionProtocolFactory fallbackSubscriptionProtocolFactory) {
 
@@ -132,6 +134,8 @@ public GraphQLWebsocketServlet(
                 Stream.of(fallbackSubscriptionProtocolFactory))
             .map(SubscriptionProtocolFactory::getProtocol)
             .collect(toList());
+
+    this.allowedOrigins = emptyList();
   }
 
   @Override
@@ -202,6 +206,26 @@ private void closeUnexpectedly(Session session, Throwable t) {
     }
   }
 
+  public boolean checkOrigin(String originHeaderValue) {
+    if (originHeaderValue == null || originHeaderValue.isBlank()) {
+      return allowedOrigins.isEmpty();
+    }
+    String originToCheck = trimTrailingSlash(originHeaderValue);
+    if (!allowedOrigins.isEmpty()) {
+      if (allowedOrigins.contains("*")) {
+        return true;
+      }
+      return allowedOrigins.stream()
+          .map(this::trimTrailingSlash)
+          .anyMatch(originToCheck::equalsIgnoreCase);
+    }
+    return true;
+  }
+
+  private String trimTrailingSlash(String origin) {
+    return (origin.endsWith("/") ? origin.substring(0, origin.length() - 1) : origin);
+  }
+
   public void modifyHandshake(
       ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {
     sec.getUserProperties().put(HANDSHAKE_REQUEST_KEY, request);