JSON Web Token (JWT) with Spring based SockJS / STOMP Web Socket

Question

Background

I am in the process of setting up a RESTful web application using Spring Boot (1.3.0.BUILD-SNAPSHOT) that includes a STOMP/SockJS WebSocket, which I intend to consume from an iOS app as well as web browsers. I want to use JSON Web Tokens (JWT) to secure the REST requests and the WebSocket interface but I’m having difficulty with the latter.

The app is secured with Spring Security:-

@Configuration
@EnableWebSecurity
public class WebSecurityConfiguration extends WebSecurityConfigurerAdapter {

    public WebSecurityConfiguration() {
        super(true);
    }

    @Autowired
    public void configureGlobal(AuthenticationManagerBuilder auth) throws Exception {
        auth.inMemoryAuthentication()
                .withUser("steve").password("steve").roles("USER");
    }

    @Override
    protected void configure(HttpSecurity http) throws Exception {
        http
            .exceptionHandling().and()
            .anonymous().and()
            .servletApi().and()
            .headers().cacheControl().and().and()

            // Relax CSRF on the WebSocket due to needing direct access from apps
            .csrf().ignoringAntMatchers("/ws/**").and()

            .authorizeRequests()

            //allow anonymous resource requests
            .antMatchers("/", "/index.html").permitAll()
            .antMatchers("/resources/**").permitAll()

            //allow anonymous POSTs to JWT
            .antMatchers(HttpMethod.POST, "/rest/jwt/token").permitAll()

            // Allow anonymous access to websocket 
            .antMatchers("/ws/**").permitAll()

            //all other request need to be authenticated
            .anyRequest().hasRole("USER").and()

            // Custom authentication on requests to /rest/jwt/token
            .addFilterBefore(new JWTLoginFilter("/rest/jwt/token", authenticationManagerBean()), UsernamePasswordAuthenticationFilter.class)

            // Custom JWT based authentication
            .addFilterBefore(new JWTTokenFilter(), UsernamePasswordAuthenticationFilter.class);
    }

}

The WebSocket configuration is standard:-

@Configuration
@EnableScheduling
@EnableWebSocketMessageBroker
public class WebSocketConfiguration extends AbstractWebSocketMessageBrokerConfigurer {

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        config.enableSimpleBroker("/topic");
        config.setApplicationDestinationPrefixes("/app");
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/ws").withSockJS();
    }

}

I also have a subclass of AbstractSecurityWebSocketMessageBrokerConfigurer to secure the WebSocket:-

@Configuration
public class WebSocketSecurityConfiguration extends AbstractSecurityWebSocketMessageBrokerConfigurer {

    @Override
    protected void configureInbound(MessageSecurityMetadataSourceRegistry messages) {
        messages.anyMessage().hasRole("USER");
    }

    @Override
    protected boolean sameOriginDisabled() {
        // We need to access this directly from apps, so can't do cross-site checks
        return true;
    }

}

There is also a couple of @RestController annotated classes to handle various bits of functionality and these are secured successfully via the JWTTokenFilter registered in my WebSecurityConfiguration class.

Problem

However I can't seem to get the WebSocket to be secured with JWT. I am using SockJS 1.1.0 and STOMP 1.7.1 in the browser and can't figure out how to pass the token. It would appear that SockJS does not allow parameters to be sent with the initial /info and/or handshake requests.

The Spring Security for WebSockets documentation states that the AbstractSecurityWebSocketMessageBrokerConfigurer ensures that:

Any inbound CONNECT message requires a valid CSRF token to enforce Same Origin Policy

Which seems to imply that the initial handshake should be unsecured and authentication invoked at the point of receiving a STOMP CONNECT message. Unfortunately I can't seem to find any information with regards to implementing this. Additionally this approach would require additional logic to disconnect a rogue client that opens a WebSocket connection and never sends a STOMP CONNECT.

Being (very) new to Spring I'm also not sure if or how Spring Sessions fits into this. While the documentation is very detailed there doesn't appear to a nice and simple (aka idiots) guide to how the various components fit together / interact with each other.

Question

How do I go about securing the SockJS WebSocket by providing a JSON Web Token, preferably at the point of handshake (is it even possible)?

Answer 1

Current Situation

UPDATE 2016-12-13 : the issue referenced below is now marked fixed, so the hack below is no longer necessary which Spring 4.3.5 or above. See https://github.com/spring-projects/spring-framework/blob/master/src/docs/asciidoc/web/websocket.adoc#token-authentication.

Previous Situation

Currently (Sep 2016), this is not supported by Spring except via query parameter as answered by @rossen-stoyanchev, who wrote a lot (all?) of the Spring WebSocket support. I don't like the query parameter approach because of potential HTTP referrer leakage and storage of the token in server logs. In addition, if the security ramifications don't bother you, note that I have found this approach works for true WebSocket connections, but if you are using SockJS with fallbacks to other mechanisms, the determineUser method is never called for the fallback. See Spring 4.x token-based WebSocket SockJS fallback authentication.

I've created a Spring issue to improve support for token-based WebSocket authentication: https://jira.spring.io/browse/SPR-14690

Hacking It

In the meantime, I've found a hack that works well in testing. Bypass the built-in Spring connection-level Spring auth machinery. Instead, set the authentication token at the message-level by sending it in the Stomp headers on the client side (this nicely mirrors what you are already doing with regular HTTP XHR calls) e.g.:

stompClient.connect({'X-Authorization': 'token'}, ...);
stompClient.subscribe(..., {'X-Authorization': 'token'});
stompClient.send("/wherever", {'X-Authorization': 'token'}, ...);

On the server-side, obtain the token from the Stomp message using a ChannelInterceptor

@Override
public void configureClientInboundChannel(ChannelRegistration registration) {
  registration.setInterceptors(new ChannelInterceptorAdapter() {
     Message<*> preSend(Message<*> message,  MessageChannel channel) {
      StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
      List tokenList = accessor.getNativeHeader("X-Authorization");
      String token = null;
      if(tokenList == null || tokenList.size < 1) {
        return message;
      } else {
        token = tokenList.get(0);
        if(token == null) {
          return message;
        }
      }

      // validate and convert to a Principal based on your own requirements e.g.
      // authenticationManager.authenticate(JwtAuthentication(token))
      Principal yourAuth = [...];

      accessor.setUser(yourAuth);

      // not documented anywhere but necessary otherwise NPE in StompSubProtocolHandler!
      accessor.setLeaveMutable(true);
      return MessageBuilder.createMessage(message.payload, accessor.messageHeaders)
    }
  })

This is simple and gets us 85% of the way there, however, this approach does not support sending messages to specific users. This is because Spring's machinery to associate users to sessions is not affected by the result of the ChannelInterceptor. Spring WebSocket assumes authentication is done at the transport layer, not the message layer, and thus ignores the message-level authentication.

The hack to make this work anyway, is to create our instances of DefaultSimpUserRegistry and DefaultUserDestinationResolver, expose those to the environment, and then use the interceptor to update those as if Spring itself was doing it. In other words, something like:

@Configuration
@EnableWebSocketMessageBroker
@Order(HIGHEST_PRECEDENCE + 50)
class WebSocketConfig extends AbstractWebSocketMessageBrokerConfigurer() {
  private DefaultSimpUserRegistry userRegistry = new DefaultSimpUserRegistry();
  private DefaultUserDestinationResolver resolver = new DefaultUserDestinationResolver(userRegistry);

  @Bean
  @Primary
  public SimpUserRegistry userRegistry() {
    return userRegistry;
  }

  @Bean
  @Primary
  public UserDestinationResolver userDestinationResolver() {
    return resolver;
  }


  @Override
  public configureMessageBroker(MessageBrokerRegistry registry) {
    registry.enableSimpleBroker("/queue", "/topic");
  }

  @Override
  public registerStompEndpoints(StompEndpointRegistry registry) {
    registry
      .addEndpoint("/stomp")
      .withSockJS()
      .setWebSocketEnabled(false)
      .setSessionCookieNeeded(false);
  }

  @Override public configureClientInboundChannel(ChannelRegistration registration) {
    registration.setInterceptors(new ChannelInterceptorAdapter() {
       Message<*> preSend(Message<*> message,  MessageChannel channel) {
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);

        List tokenList = accessor.getNativeHeader("X-Authorization");
        accessor.removeNativeHeader("X-Authorization");

        String token = null;
        if(tokenList != null && tokenList.size > 0) {
          token = tokenList.get(0);
        }

        // validate and convert to a Principal based on your own requirements e.g.
        // authenticationManager.authenticate(JwtAuthentication(token))
        Principal yourAuth = token == null ? null : [...];

        if (accessor.messageType == SimpMessageType.CONNECT) {
          userRegistry.onApplicationEvent(SessionConnectedEvent(this, message, yourAuth));
        } else if (accessor.messageType == SimpMessageType.SUBSCRIBE) {
          userRegistry.onApplicationEvent(SessionSubscribeEvent(this, message, yourAuth));
        } else if (accessor.messageType == SimpMessageType.UNSUBSCRIBE) {
          userRegistry.onApplicationEvent(SessionUnsubscribeEvent(this, message, yourAuth));
        } else if (accessor.messageType == SimpMessageType.DISCONNECT) {
          userRegistry.onApplicationEvent(SessionDisconnectEvent(this, message, accessor.sessionId, CloseStatus.NORMAL));
        }

        accessor.setUser(yourAuth);

        // not documented anywhere but necessary otherwise NPE in StompSubProtocolHandler!
        accessor.setLeaveMutable(true);
        return MessageBuilder.createMessage(message.payload, accessor.messageHeaders);
      }
    })
  }
}

Now Spring is fully aware of the the authentication i.e. it injects the Principal into any controller methods that require it, exposes it to the context for Spring Security 4.x, and associates the user to the WebSocket session for sending messages to specific users/sessions.

Spring Security Messaging

Lastly, if you use Spring Security 4.x Messaging support, make sure to set the @Order of your AbstractWebSocketMessageBrokerConfigurer to a higher value than Spring Security's AbstractSecurityWebSocketMessageBrokerConfigurer (Ordered.HIGHEST_PRECEDENCE + 50 would work, as shown above). That way, your interceptor sets the Principal before Spring Security executes its check and sets the security context.

Creating a Principal (Update June 2018)

Lots of people seem to be confused by this line in the code above:

  // validate and convert to a Principal based on your own requirements e.g.
  // authenticationManager.authenticate(JwtAuthentication(token))
  Principal yourAuth = [...];

This is pretty much out of scope for the question as it is not Stomp-specific, but I'll expand on it a little bit anyway, because its related to using auth tokens with Spring. When using token-based authentication, the Principal you need will generally be a custom JwtAuthentication class that extends Spring Security's AbstractAuthenticationToken class. AbstractAuthenticationToken implements the Authentication interface which extends the Principal interface, and contains most of the machinery to integrate your token with Spring Security.

So, in Kotlin code (sorry I don't have the time or inclination to translate this back to Java), your JwtAuthentication might look something like this, which is a simple wrapper around AbstractAuthenticationToken:

import my.model.UserEntity
import org.springframework.security.authentication.AbstractAuthenticationToken
import org.springframework.security.core.GrantedAuthority

class JwtAuthentication(
  val token: String,
  // UserEntity is your application's model for your user
  val user: UserEntity? = null,
  authorities: Collection<GrantedAuthority>? = null) : AbstractAuthenticationToken(authorities) {

  override fun getCredentials(): Any? = token

  override fun getName(): String? = user?.id

  override fun getPrincipal(): Any? = user
}

Now you need an AuthenticationManager that knows how to deal with it. This might look something like the following, again in Kotlin:

@Component
class CustomTokenAuthenticationManager @Inject constructor(
  val tokenHandler: TokenHandler,
  val authService: AuthService) : AuthenticationManager {

  val log = logger()

  override fun authenticate(authentication: Authentication?): Authentication? {
    return when(authentication) {
      // for login via username/password e.g. crash shell
      is UsernamePasswordAuthenticationToken -> {
        findUser(authentication).let {
          //checkUser(it)
          authentication.withGrantedAuthorities(it).also { setAuthenticated(true) }
        }
      }
      // for token-based auth
      is JwtAuthentication -> {
        findUser(authentication).let {
          val tokenTypeClaim = tokenHandler.parseToken(authentication.token)[CLAIM_TOKEN_TYPE]
          when(tokenTypeClaim) {
            TOKEN_TYPE_ACCESS -> {
              //checkUser(it)
              authentication.withGrantedAuthorities(it).also { setAuthenticated(true) }
            }
            TOKEN_TYPE_REFRESH -> {
              //checkUser(it)
              JwtAuthentication(authentication.token, it, listOf(SimpleGrantedAuthority(Authorities.REFRESH_TOKEN)))
            }
            else -> throw IllegalArgumentException("Unexpected token type claim $tokenTypeClaim.")
          }
        }
      }
      else -> null
    }
  }

  private fun findUser(authentication: JwtAuthentication): UserEntity =
    authService.login(authentication.token) ?:
      throw BadCredentialsException("No user associated with token or token revoked.")

  private fun findUser(authentication: UsernamePasswordAuthenticationToken): UserEntity =
    authService.login(authentication.principal.toString(), authentication.credentials.toString()) ?:
      throw BadCredentialsException("Invalid login.")

  @Suppress("unused", "UNUSED_PARAMETER")
  private fun checkUser(user: UserEntity) {
    // TODO add these and lock account on x attempts
    //if(!user.enabled) throw DisabledException("User is disabled.")
    //if(user.accountLocked) throw LockedException("User account is locked.")
  }

  fun JwtAuthentication.withGrantedAuthorities(user: UserEntity): JwtAuthentication {
    return JwtAuthentication(token, user, authoritiesOf(user))
  }

  fun UsernamePasswordAuthenticationToken.withGrantedAuthorities(user: UserEntity): UsernamePasswordAuthenticationToken {
    return UsernamePasswordAuthenticationToken(principal, credentials, authoritiesOf(user))
  }

  private fun authoritiesOf(user: UserEntity) = user.authorities.map(::SimpleGrantedAuthority)
}

The injected TokenHandler abstracts away the JWT token parsing, but should use a common JWT token library like jjwt. The injected AuthService is your abstraction that actually creates your UserEntity based on the claims in the token, and may talk to your user database or other backend system(s).

Now, coming back to the line we started with, it might look something like this, where authenticationManager is an AuthenticationManager injected into our adapter by Spring, and is an instance of CustomTokenAuthenticationManager we defined above:

Principal yourAuth = token == null ? null : authenticationManager.authenticate(new JwtAuthentication(token));

This principal is then attached to the message as described above. HTH!