Spring Websockets Authentication with Spring Security and Keycloak

I was able to do websocket authentication/authorization without using Spring Security and SockJS:

@Configuration
@EnableWebSocketMessageBroker
@RequiredArgsConstructor
public class StompConfiguration implements WebSocketMessageBrokerConfigurer {

    private final KeycloakSpringBootProperties configuration;

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        config.setApplicationDestinationPrefixes("/stompy");  // prefix for incoming messages in @MessageMapping
        config.enableSimpleBroker("/broker");                 // enabling broker @SendTo("/broker/blabla")
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/stomp")
                .addInterceptors(new StompHandshakeInterceptor(configuration))
                .setAllowedOrigins("*");
    }
}

Handshake interceptor:

@Slf4j
@RequiredArgsConstructor
public class StompHandshakeInterceptor implements HandshakeInterceptor {

    private final KeycloakSpringBootProperties configuration;

    @Override
    public boolean beforeHandshake(ServerHttpRequest req, ServerHttpResponse resp, WebSocketHandler h, Map<String, Object> atts) {
        List<String> protocols = req.getHeaders().get("Sec-WebSocket-Protocol");
        try {
            String token = protocols.get(0).split(", ")[2];
            log.debug("Token: " + token);
            AdapterTokenVerifier.verifyToken(token, KeycloakDeploymentBuilder.build(configuration));
            resp.setStatusCode(HttpStatus.SWITCHING_PROTOCOLS);
            log.debug("token valid");
        } catch (IndexOutOfBoundsException e) {
            resp.setStatusCode(HttpStatus.UNAUTHORIZED);
            return false;
        }
        catch (VerificationException e) {
            resp.setStatusCode(HttpStatus.FORBIDDEN);
            log.error(e.getMessage());
            return false;
        }
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest rq, ServerHttpResponse rp, WebSocketHandler h, @Nullable Exception e) {}
}

Websocket controller:

@Controller
public class StompController {
    @MessageMapping("/test")
    @SendTo("/broker/lol")
    public String lol(String message) {
        System.out.println("Incoming message: " + message);
        return message;
    }
}

Client side (javascript):

function connect() {
    let protocols = ['v10.stomp', 'v11.stomp'];
    protocols.push("KEYCLOAK TOKEN");
    const url = "ws://localhost:8080/stomp";

    client = Stomp.client(url, protocols);
    client.connect(
        {},
        () => {
            console.log("Connection established");
            client.subscribe("/broker/lol", function (mes) {
                console.log("New message for /broker/lol: " + mes.body);
            });
        },
        error => { console.log("ERROR: " + error); }
    );
}

function sendMessage() {
    let message = "test message";
    if (client) client.send("/stompy/test", {}, message);
}

build.gradle:

dependencies {
    implementation 'org.springframework.boot:spring-boot-starter-web'
    implementation 'org.springframework.boot:spring-boot-starter-websocket'
    compileOnly 'org.projectlombok:lombok'
    testImplementation 'org.springframework.boot:spring-boot-starter-test'

    // keycloak
    implementation 'org.keycloak:keycloak-spring-boot-starter'

    // stomp.js
    implementation("org.webjars:webjars-locator-core")
    implementation("org.webjars:stomp-websocket:2.3.3")
}

dependencyManagement {
    imports {
        mavenBom "org.keycloak.bom:keycloak-adapter-bom:$keycloakVersion"
    }
}

As you can see the client is authenticated during the handshake. The HandshakeInterceptor class extracts the token from the Sec-WebSocket-Protocol header. No SockJS or Spring Security is needed. Hope this helps :)

I like the answer of adrianmoya except for the part of the KeycloakTokenVerifier. I use the following instead:

public class KeycloakWebSocketAuthManager implements AuthenticationManager {

  private final KeycloakSpringBootConfigResolver keycloakSpringBootConfigResolver;

  @Override
  public Authentication authenticate(final Authentication authentication) throws AuthenticationException {
     final JWSAuthenticationToken token = (JWSAuthenticationToken) authentication;
     final String tokenString = (String) token.getCredentials();
     try {
        final KeycloakDeployment resolve = keycloakSpringBootConfigResolver.resolve(null);
        final AccessToken accessToken = AdapterRSATokenVerifier.verifyToken(tokenString, resolve);
       ...
      }
}

I was able to enable token based authentication, following the recomendations by Raman on this question. Here's the final code to make it work:

1) First, create a class that represent the JWS auth token:

public class JWSAuthenticationToken extends AbstractAuthenticationToken implements Authentication {

  private static final long serialVersionUID = 1L;

  private String token;
  private User principal;

  public JWSAuthenticationToken(String token) {
    this(token, null, null);
  }

  public JWSAuthenticationToken(String token, User principal, Collection<GrantedAuthority> authorities) {
    super(authorities);
    this.token = token;
    this.principal = principal;
  }

  @Override
  public Object getCredentials() {
    return token;
  }

  @Override
  public Object getPrincipal() {
    return principal;
  }

}

2) Then, create an authenticator that handles the JWSToken, validating against keycloak. User is my own app class that represents a user:

@Slf4j
@Component
@Qualifier("websocket")
@AllArgsConstructor
public class KeycloakWebSocketAuthManager implements AuthenticationManager {

  private final KeycloakTokenVerifier tokenVerifier;

  @Override
  public Authentication authenticate(Authentication authentication) throws AuthenticationException {
    JWSAuthenticationToken token = (JWSAuthenticationToken) authentication;
    String tokenString = (String) token.getCredentials();
    try {
      AccessToken accessToken = tokenVerifier.verifyToken(tokenString);
      List<GrantedAuthority> authorities = accessToken.getRealmAccess().getRoles().stream()
          .map(SimpleGrantedAuthority::new).collect(Collectors.toList());
      User user = new User(accessToken.getName(), accessToken.getEmail(), accessToken.getPreferredUsername(),
          accessToken.getRealmAccess().getRoles());
      token = new JWSAuthenticationToken(tokenString, user, authorities);
      token.setAuthenticated(true);
    } catch (VerificationException e) {
      log.debug("Exception authenticating the token {}:", tokenString, e);
      throw new BadCredentialsException("Invalid token");
    }
    return token;
  }

}

3) The class that actually validates the token against keycloak by calling the certs endpoint to validate the token signature, based on this gists. It returns a keycloak AccessToken:

@Component
@AllArgsConstructor
public class KeycloakTokenVerifier {

  private final KeycloakProperties config;

  /**
   * Verifies a token against a keycloak instance
   * @param tokenString the string representation of the jws token
   * @return a validated keycloak AccessToken
   * @throws VerificationException when the token is not valid
   */
  public AccessToken verifyToken(String tokenString) throws VerificationException {
    RSATokenVerifier verifier = RSATokenVerifier.create(tokenString);
    PublicKey publicKey = retrievePublicKeyFromCertsEndpoint(verifier.getHeader());
    return verifier.realmUrl(getRealmUrl()).publicKey(publicKey).verify().getToken();
  }

  @SuppressWarnings("unchecked")
  private PublicKey retrievePublicKeyFromCertsEndpoint(JWSHeader jwsHeader) {
    try {
      ObjectMapper om = new ObjectMapper();
      Map<String, Object> certInfos = om.readValue(new URL(getRealmCertsUrl()).openStream(), Map.class);
      List<Map<String, Object>> keys = (List<Map<String, Object>>) certInfos.get("keys");

      Map<String, Object> keyInfo = null;
      for (Map<String, Object> key : keys) {
        String kid = (String) key.get("kid");
        if (jwsHeader.getKeyId().equals(kid)) {
          keyInfo = key;
          break;
        }
      }

      if (keyInfo == null) {
        return null;
      }

      KeyFactory keyFactory = KeyFactory.getInstance("RSA");
      String modulusBase64 = (String) keyInfo.get("n");
      String exponentBase64 = (String) keyInfo.get("e");
      Decoder urlDecoder = Base64.getUrlDecoder();
      BigInteger modulus = new BigInteger(1, urlDecoder.decode(modulusBase64));
      BigInteger publicExponent = new BigInteger(1, urlDecoder.decode(exponentBase64));

      return keyFactory.generatePublic(new RSAPublicKeySpec(modulus, publicExponent));

    } catch (Exception e) {
      e.printStackTrace();
    }
    return null;
  }

  public String getRealmUrl() {
    return String.format("%s/realms/%s", config.getAuthServerUrl(), config.getRealm());
  }

  public String getRealmCertsUrl() {
    return getRealmUrl() + "/protocol/openid-connect/certs";
  }

}

4) Finally, inject the authenticator in the Websocket configuration and complete the piece of code as recommended by spring docs:

@Slf4j
@Configuration
@EnableWebSocketMessageBroker
@AllArgsConstructor
public class WebSocketConfiguration extends AbstractWebSocketMessageBrokerConfigurer {

  @Qualifier("websocket")
  private AuthenticationManager authenticationManager;

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

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

  @Override
  public void configureClientInboundChannel(ChannelRegistration registration) {
    registration.interceptors(new ChannelInterceptorAdapter() {
      @Override
      public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        if (StompCommand.CONNECT.equals(accessor.getCommand())) {
          Optional.ofNullable(accessor.getNativeHeader("Authorization")).ifPresent(ah -> {
            String bearerToken = ah.get(0).replace("Bearer ", "");
            log.debug("Received bearer token {}", bearerToken);
            JWSAuthenticationToken token = (JWSAuthenticationToken) authenticationManager
                .authenticate(new JWSAuthenticationToken(bearerToken));
            accessor.setUser(token);
          });
        }
        return message;
      }
    });
  }

}

I also changed my security configuration a bit. First, I excluded the WS endpoint from spring web securty, and also let the connection methods open to anyone in the websocket security:

In WebSecurityConfiguration:

  @Override
  public void configure(WebSecurity web) throws Exception {
    web.ignoring()
        .antMatchers("/ws-endpoint/**");
  }

And in the class WebSocketSecurityConfig:

@Configuration
public class WebSocketSecurityConfig extends AbstractSecurityWebSocketMessageBrokerConfigurer {

  @Override
  protected void configureInbound(MessageSecurityMetadataSourceRegistry messages) {
    messages.simpTypeMatchers(CONNECT, UNSUBSCRIBE, DISCONNECT, HEARTBEAT).permitAll()
    .simpDestMatchers("/app/**", "/topic/**").authenticated().simpSubscribeDestMatchers("/topic/**").authenticated()
        .anyMessage().denyAll();
  }

  @Override
  protected boolean sameOriginDisabled() {
    return true;
  }
}

So the final result is: anybody in the local network can connect to the socket, but to actually subscribe to any channel, you have to be authenticated, so you need to send the Bearer token with the original CONNECT message or you'll get UnauthorizedException. Hope it helps others with this requeriment!