Chapter 34: Configure subscription Pre-Auth spring security (#7)

Author: philip-jvm

src/main/java/com/learn/graphql/config/security/AuthenticationConnectionListener.java

@@ -0,0 +1,81 @@
+package com.learn.graphql.config.security;
+
+import graphql.kickstart.execution.subscriptions.SubscriptionSession;
+import graphql.kickstart.execution.subscriptions.apollo.ApolloSubscriptionConnectionListener;
+import graphql.kickstart.execution.subscriptions.apollo.OperationMessage;
+import java.util.Map;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.web.authentication.preauth.PreAuthenticatedAuthenticationToken;
+import org.springframework.stereotype.Component;
+
+@Slf4j
+@Component
+public class AuthenticationConnectionListener implements ApolloSubscriptionConnectionListener {
+
+  public static final String AUTHENTICATION = "AUTHENTICATION";
+
+  /**
+   * Chapter 34: Subscriptions Spring Security Pre-Auth. When using pre-auth, you must ensure that
+   * all the graphql requests have been previously authorized/authenticated by an upstream service.
+   * For example, all ingress traffic to this graphql server must bypass an upstream proxy node that
+   * will validate the request's JWT token. This code alone performs no authorization. Read more
+   * about Pre-auth before using this.
+   */
+  @Override
+  public void onConnect(SubscriptionSession session, OperationMessage message) {
+    log.info("onConnect with payload {}", message.getPayload());
+
+    var payload = (Map<String, String>) message.getPayload();
+
+    // Get the user id, roles (or JWT etc) and perform authentication / rejection here
+    var userId = payload.get(GraphQLSecurityConfig.USER_ID_PRE_AUTH_HEADER);
+    var userRoles = payload.get(GraphQLSecurityConfig.USER_ROLES_PRE_AUTH_HEADER);
+    var grantedAuthorities = GrantedAuthorityFactory.getAuthoritiesFrom(userRoles);
+
+    /**
+
+     Q: Why do not set the token/Authentication inside Spring Security SecurityContextHolder here?
+
+     If the start frame is not sent directly with the connection_init then the two frames may be serviced on different threads.
+     The thread servicing the connection_init frame will check the websocket for any further inbound frames,
+     if false the thread will move onto another websocket. Another thread is then free to service the following start frame.
+     In this case, that thread not have the security context of the correct session/thread.
+
+     Same scenario happens for onStop. (Message can be executed on different thread).
+
+     This seems to be why some users are reporting intermittent failures with spring security.
+     E.g. https://github.com/graphql-java-kickstart/graphql-java-servlet/discussions/134#discussioncomment-225980
+
+     With the NIO connector, a small number of threads will check sessions for new frames.
+     If a session has a frame available, the session will be passed to another thread pool which will read frame, execute it, check for another frame, execute it (loop).
+     The session will be released when there are no further frames available. With this, we know that at most one thread will concurrently access one socket,
+     therefore frames will be read sequentially. We can therefore extract the auth credentials from onConnect and add them to the session.getUserProperties().
+     These properties are available in the onStart and onStop callbacks. Inside these callbacks, we can add the token to the SecurityContextHolder if we decide to use method level security,
+     or simply access the credentials inside the subscription resolver via DataFetchingEnvironment.
+
+     */
+
+    var token = new PreAuthenticatedAuthenticationToken(userId, null, grantedAuthorities);
+    session.getUserProperties().put(AUTHENTICATION, token);
+  }
+
+  @Override
+  public void onStart(SubscriptionSession session, OperationMessage message) {
+    log.info("onStart with payload {}", message.getPayload());
+    var authentication = (Authentication) session.getUserProperties().get(AUTHENTICATION);
+    SecurityContextHolder.getContext().setAuthentication(authentication);
+  }
+
+  @Override
+  public void onStop(SubscriptionSession session, OperationMessage message) {
+    log.info("onStop with payload {}", message.getPayload());
+  }
+
+  @Override
+  public void onTerminate(SubscriptionSession session, OperationMessage message) {
+    log.info("onTerminate with payload {}", message.getPayload());
+  }
+
+}

src/main/java/com/learn/graphql/config/security/GrantedAuthoritiesAuthenticationDetailsSource.java

@@ -2,14 +2,8 @@
 
 import static com.learn.graphql.config.security.GraphQLSecurityConfig.USER_ROLES_PRE_AUTH_HEADER;
 
-import java.util.List;
-import java.util.Set;
-import java.util.stream.Collectors;
 import javax.servlet.http.HttpServletRequest;
-import org.apache.commons.lang3.StringUtils;
 import org.springframework.security.authentication.AuthenticationDetailsSource;
-import org.springframework.security.core.GrantedAuthority;
-import org.springframework.security.core.authority.SimpleGrantedAuthority;
 import org.springframework.security.web.authentication.preauth.PreAuthenticatedGrantedAuthoritiesWebAuthenticationDetails;
 
 public class GrantedAuthoritiesAuthenticationDetailsSource implements
@@ -19,16 +13,8 @@ public class GrantedAuthoritiesAuthenticationDetailsSource implements
   public PreAuthenticatedGrantedAuthoritiesWebAuthenticationDetails buildDetails(
       HttpServletRequest request) {
     var userRoles = request.getHeader(USER_ROLES_PRE_AUTH_HEADER);
-    var authorities = StringUtils.isBlank(userRoles) ? List.<GrantedAuthority>of() :
-        getAuthorities(userRoles);
+    var authorities = GrantedAuthorityFactory.getAuthoritiesFrom(userRoles);
     return new PreAuthenticatedGrantedAuthoritiesWebAuthenticationDetails(request, authorities);
   }
 
-  private List<GrantedAuthority> getAuthorities(String userRoles) {
-    return Set.of(userRoles.split(","))
-        .stream()
-        .map(SimpleGrantedAuthority::new)
-        .collect(Collectors.toList());
-  }
-
 }

src/main/java/com/learn/graphql/config/security/GrantedAuthorityFactory.java

@@ -0,0 +1,26 @@
+package com.learn.graphql.config.security;
+
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import lombok.AccessLevel;
+import lombok.NoArgsConstructor;
+import org.apache.commons.lang3.StringUtils;
+import org.springframework.security.core.GrantedAuthority;
+import org.springframework.security.core.authority.SimpleGrantedAuthority;
+
+@NoArgsConstructor(access = AccessLevel.PRIVATE)
+public class GrantedAuthorityFactory {
+
+  public static List<GrantedAuthority> getAuthoritiesFrom(String userRoles) {
+    if (StringUtils.isBlank(userRoles)) {
+      return List.of();
+    }
+
+    return Set.of(userRoles.split(","))
+        .stream()
+        .map(SimpleGrantedAuthority::new)
+        .collect(Collectors.toList());
+  }
+
+}

src/main/java/com/learn/graphql/config/security/GraphQLSecurityConfig.java

@@ -77,7 +77,7 @@ public void configure(WebSecurity web) {
         .antMatchers("/actuator/health")
         // Permit playground for development
         .antMatchers("/playground", "/vendor/playground/**")
-        // Disable security for subscription example
+        // Subscription are secured via AuthenticationConnectionListener
         .antMatchers("/subscriptions");
   }
 

src/main/java/com/learn/graphql/resolver/bank/mutation/BankAccountMutation.java

@@ -29,9 +29,20 @@ public class BankAccountMutation implements GraphQLMutationResolver {
    */
   public BankAccount createBankAccount(@Valid CreateBankAccountInput input) {
     log.info("Creating bank account for {}", input);
+    return getBankAccount(UUID.randomUUID());
+  }
+
+  /**
+   * Schema Directive Validation (Chapter 32)
+   */
+  public BankAccount updateBankAccount(UUID id, String name, int age) {
+    log.info("Updating bank account for {}. Name: {}, age: {}", id, name, age);
+    return getBankAccount(id);
+  }
 
+  private BankAccount getBankAccount(UUID id) {
     var bankAccount = BankAccount.builder()
-        .id(UUID.randomUUID())
+        .id(id)
         .currency(Currency.USD)
         .createdAt(ZonedDateTime.now(clock))
         .createdOn(LocalDate.now(clock))
@@ -45,17 +56,4 @@ public BankAccount createBankAccount(@Valid CreateBankAccountInput input) {
     return bankAccount;
   }
 
-  /**
-   * Schema Directive Validation (Chapter 32)
-   */
-  public BankAccount updateBankAccount(UUID id, String name, int age) {
-    log.info("Updating bank account for {}. Name: {}, age: {}", id, name, age);
-    return BankAccount.builder()
-        .id(UUID.randomUUID())
-        .currency(Currency.USD)
-        .createdAt(ZonedDateTime.now(clock))
-        .createdOn(LocalDate.now(clock))
-        .build();
-  }
-
 }

src/main/java/com/learn/graphql/resolver/bank/subscription/BankAccountSubscription.java

@@ -1,28 +1,52 @@
 package com.learn.graphql.resolver.bank.subscription;
 
+import com.learn.graphql.config.security.AuthenticationConnectionListener;
 import com.learn.graphql.domain.bank.BankAccount;
 import com.learn.graphql.publisher.BankAccountPublisher;
+import graphql.kickstart.servlet.context.GraphQLWebSocketContext;
 import graphql.kickstart.tools.GraphQLSubscriptionResolver;
+import graphql.schema.DataFetchingEnvironment;
 import java.util.UUID;
 import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
 import org.reactivestreams.Publisher;
+import org.springframework.security.access.prepost.PreAuthorize;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.SecurityContextHolder;
 import org.springframework.stereotype.Component;
 
 /**
  * Subscription (Chapter 33)
  */
+@Slf4j
 @Component
 @RequiredArgsConstructor
 public class BankAccountSubscription implements GraphQLSubscriptionResolver {
 
   private final BankAccountPublisher bankAccountPublisher;
 
+  @PreAuthorize("hasAuthority('get:bank_account')")
   public Publisher<BankAccount> bankAccounts() {
     return bankAccountPublisher.getBankAccountPublisher();
   }
 
-  public Publisher<BankAccount> bankAccount(UUID id) {
+  @PreAuthorize("hasAuthority('get:bank_account')")
+  public Publisher<BankAccount> bankAccount(UUID id, DataFetchingEnvironment e) {
+    log.info("Creating bank account publisher for user Id: {}",
+        SecurityContextHolder.getContext().getAuthentication().getPrincipal());
+
+    // As an alternative to spring-security, you can access the authentication via the DataFetchingEnvironment
+    GraphQLWebSocketContext context = getContext(e);
+    var authentication = (Authentication) context.getSession().getUserProperties()
+        .get(AuthenticationConnectionListener.AUTHENTICATION);
+    log.info("Creating bank account publisher for user Id: {}",
+        authentication.getPrincipal());
+
     return bankAccountPublisher.getBankAccountPublisherFor(id);
   }
 
+  private <T> T getContext(DataFetchingEnvironment e) {
+    return e.getContext();
+  }
+
 }