diff --git a/llm-service/llm-service-server/src/main/java/com/knecon/fforesight/llm/service/websocket/WebSocketConfiguration.java b/llm-service/llm-service-server/src/main/java/com/knecon/fforesight/llm/service/websocket/WebSocketConfiguration.java index 8e0695e..028a63f 100644 --- a/llm-service/llm-service-server/src/main/java/com/knecon/fforesight/llm/service/websocket/WebSocketConfiguration.java +++ b/llm-service/llm-service-server/src/main/java/com/knecon/fforesight/llm/service/websocket/WebSocketConfiguration.java @@ -50,7 +50,6 @@ public class WebSocketConfiguration implements WebSocketMessageBrokerConfigurer @Override public void registerStompEndpoints(StompEndpointRegistry registry) { - registry.addEndpoint("/api/llm/llm-websocket").setAllowedOrigins("*"); } @@ -65,14 +64,13 @@ public class WebSocketConfiguration implements WebSocketMessageBrokerConfigurer 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); - AuthenticationManager authenticationManager = tenantAuthenticationManagerResolver.resolve(bearerToken); - JwtAuthenticationToken token = (JwtAuthenticationToken) authenticationManager.authenticate(new BearerTokenAuthenticationToken(bearerToken)); - accessor.setUser(token); - }); + Optional.ofNullable(accessor.getNativeHeader("Authorization")).ifPresent(ah -> { + String bearerToken = ah.get(0).replace("Bearer ", ""); + log.debug("Received bearer token {}", bearerToken); + AuthenticationManager authenticationManager = tenantAuthenticationManagerResolver.resolve(bearerToken); + JwtAuthenticationToken token = (JwtAuthenticationToken) authenticationManager.authenticate(new BearerTokenAuthenticationToken(bearerToken)); + accessor.setUser(token); + }); } return message; } diff --git a/llm-service/llm-service-server/src/main/java/com/knecon/fforesight/llm/service/websocket/WebSocketSecurityConfig.java b/llm-service/llm-service-server/src/main/java/com/knecon/fforesight/llm/service/websocket/WebSocketSecurityConfig.java index b192d3e..dfe6a55 100644 --- a/llm-service/llm-service-server/src/main/java/com/knecon/fforesight/llm/service/websocket/WebSocketSecurityConfig.java +++ b/llm-service/llm-service-server/src/main/java/com/knecon/fforesight/llm/service/websocket/WebSocketSecurityConfig.java @@ -63,9 +63,10 @@ public class WebSocketSecurityConfig extends AbstractSecurityWebSocketMessageBro } - public boolean checkAuthAndSetTenant(JwtAuthenticationToken authentication, Message message){ + public boolean checkAuthAndSetTenant(JwtAuthenticationToken authentication, Message message) { + Optional tenantId = getCurrentTenant(authentication); - if (tenantId.isPresent()){ + if (tenantId.isPresent()) { TenantContext.setTenantId(tenantId.get()); return true; } @@ -84,6 +85,7 @@ public class WebSocketSecurityConfig extends AbstractSecurityWebSocketMessageBro } + private Optional extractTenantId(Message message) { StompHeaderAccessor sha = StompHeaderAccessor.wrap(message);