Skip to content

Commit 4d32ac1

Browse files
phrockerphrocker
committed
fix bug
1 parent 6941aa7 commit 4d32ac1

9 files changed

Lines changed: 76 additions & 47 deletions

File tree

‎.local.env‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
SENTRIUS_VERSION=1.1.92
1+
SENTRIUS_VERSION=1.1.95
22
SENTRIUS_SSH_VERSION=1.1.18
33
SENTRIUS_KEYCLOAK_VERSION=1.1.25
44
SENTRIUS_AGENT_VERSION=1.1.18
55
SENTRIUS_AI_AGENT_VERSION=1.1.33
66
LLMPROXY_VERSION=1.0.18
7-
LAUNCHER_VERSION=1.0.27
7+
LAUNCHER_VERSION=1.0.29

‎.local.env.bak‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
SENTRIUS_VERSION=1.1.92
1+
SENTRIUS_VERSION=1.1.95
22
SENTRIUS_SSH_VERSION=1.1.18
33
SENTRIUS_KEYCLOAK_VERSION=1.1.25
44
SENTRIUS_AGENT_VERSION=1.1.18
55
SENTRIUS_AI_AGENT_VERSION=1.1.33
66
LLMPROXY_VERSION=1.0.18
7-
LAUNCHER_VERSION=1.0.26
7+
LAUNCHER_VERSION=1.0.28

‎api/src/main/java/io/sentrius/sso/locator/KubernetesAgentLocator.java‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -9,16 +9,12 @@
99
@Component
1010
public class KubernetesAgentLocator {
1111

12-
@Value("${sentrius.agent.namespace}")
13-
private String agentNamespace;
1412

15-
@Value("${sentrius.agent.port:8080}")
16-
private int agentPort;
17-
18-
public URI resolveWebSocketUri(String agentId) {
13+
public URI resolveWebSocketUri(String host, String sessionId, String chatGroupId, String ztat) {
1914
// DNS: sentrius-agent-[ID].[namespace].svc.cluster.local
20-
String fqdn = String.format("ws://sentrius-agent-%s.%s.svc.cluster.local:%d/ws",
21-
agentId, agentNamespace, agentPort);
15+
///api/v1/chat/attach/subscribe?sessionId=${encodeURIComponent(this.sessionId)}&chatGroupId=${this.chatGroupId}&ztat=${encodeURIComponent(jwt)
16+
String fqdn = String.format("%s/api/v1/chat/attach/subscribe?sessionId=%s&chatGroupId=%s&ztat=%s",
17+
host, sessionId, chatGroupId, ztat);
2218
return URI.create(fqdn);
2319
}
2420
}

‎api/src/main/java/io/sentrius/sso/startup/ConfigurationApplicationTask.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -683,16 +683,16 @@ protected List<User> createUsers(
683683
}
684684
}
685685
if (action){
686-
user = userService.getUser(user.getId()).orElseThrow();
686+
var newUser = userService.getUser(user.getId());
687687
var definition = userDTO.getAtlpDefinition();
688688
if (null != definition && !definition.isEmpty()) {
689689
Optional<ATPLPolicyEntity> policy = policyList.stream()
690690
.filter(p -> p.getPolicyId().equals(definition))
691691
.findFirst();
692-
if (policy.isPresent()) {
693-
atplPolicyService.assignPolicyToUser(user, policy.get());
692+
if (policy.isPresent() & newUser.isPresent()) {
693+
atplPolicyService.assignPolicyToUser(newUser.get(), policy.get());
694694
} else {
695-
log.warn("No ATPL policy found for user {} with policy id {}", user.getUsername(),
695+
log.warn("No ATPL policy found for user {} with policy id {}", newUser.get().getUsername(),
696696
definition);
697697
}
698698
}
@@ -811,16 +811,16 @@ protected List<User> createNPEs(
811811
}
812812
}
813813
if (action){
814-
user = userService.getUser(user.getId()).orElseThrow();
814+
var newUser = userService.getUser(user.getId());
815815
var definition = userDTO.getAtlpDefinition();
816816
if (null != definition && !definition.isEmpty()) {
817817
Optional<ATPLPolicyEntity> policy = policyList.stream()
818818
.filter(p -> p.getPolicyId().equals(definition))
819819
.findFirst();
820820
if (policy.isPresent()) {
821-
atplPolicyService.assignPolicyToUser(user, policy.get());
821+
atplPolicyService.assignPolicyToUser(newUser.get(), policy.get());
822822
} else {
823-
log.warn("No ATPL policy found for user {} with policy id {}", user.getUsername(),
823+
log.warn("No ATPL policy found for user {} with policy id {}", newUser.get().getUsername(),
824824
definition);
825825
}
826826
}

‎api/src/main/java/io/sentrius/sso/websocket/AgentHandshakeInterceptor.java‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77

88
import java.util.Map;
99

10+
import org.springframework.web.util.UriComponentsBuilder;
11+
1012
public class AgentHandshakeInterceptor implements HandshakeInterceptor {
1113

1214
@Override
@@ -15,12 +17,16 @@ public boolean beforeHandshake(ServerHttpRequest request,
1517
WebSocketHandler wsHandler,
1618
Map<String, Object> attributes) {
1719

18-
String path = request.getURI().getPath(); // e.g. /api/v1/agents/ws/agent-123
19-
String[] segments = path.split("/");
20-
String agentId = segments[segments.length - 1]; // assumes agentId is at the end
20+
String query = request.getURI().getQuery();
21+
Map<String, String> queryParams = UriComponentsBuilder.fromUri(request.getURI()).build().getQueryParams().toSingleValueMap();
22+
23+
attributes.put("host", queryParams.get("phost"));
24+
attributes.put("sessionId", queryParams.get("sessionId"));
25+
attributes.put("chatGroupId", queryParams.get("chatGroupId"));
26+
attributes.put("ztat", queryParams.get("ztat"));
2127

22-
attributes.put("agentId", agentId);
2328
return true;
29+
2430
}
2531

2632
@Override

‎api/src/main/java/io/sentrius/sso/websocket/AgentWebSocketProxyHandler.java‎

Lines changed: 39 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,43 +1,61 @@
11
package io.sentrius.sso.websocket;
22

33
import java.net.URI;
4+
import java.security.GeneralSecurityException;
5+
46
import io.sentrius.sso.locator.KubernetesAgentLocator;
57
import lombok.RequiredArgsConstructor;
8+
69
import org.springframework.stereotype.Component;
710
import org.springframework.web.reactive.socket.WebSocketHandler;
811
import org.springframework.web.reactive.socket.WebSocketMessage;
912
import org.springframework.web.reactive.socket.WebSocketSession;
1013
import org.springframework.web.reactive.socket.client.ReactorNettyWebSocketClient;
14+
15+
import io.sentrius.sso.core.services.security.CryptoService;
16+
import lombok.extern.slf4j.Slf4j;
1117
import reactor.core.publisher.Mono;
1218

1319
@Component
20+
@Slf4j
1421
@RequiredArgsConstructor
1522
public class AgentWebSocketProxyHandler implements WebSocketHandler {
1623

1724
private final KubernetesAgentLocator agentLocator;
25+
private final CryptoService cryptoService;
1826

1927
@Override
2028
public Mono<Void> handle(WebSocketSession clientSession) {
21-
String agentId = (String) clientSession.getAttributes().get("agentId");
22-
URI agentUri = agentLocator.resolveWebSocketUri(agentId);
23-
24-
ReactorNettyWebSocketClient proxyClient = new ReactorNettyWebSocketClient();
25-
26-
return proxyClient.execute(agentUri, agentSession -> {
27-
// Forward messages from client to agent
28-
Mono<Void> clientToAgent = clientSession.receive()
29-
.map(WebSocketMessage::getPayload)
30-
.map(dataBuffer -> agentSession.binaryMessage(factory -> dataBuffer))
31-
.as(agentSession::send);
32-
33-
// Forward messages from agent to client
34-
Mono<Void> agentToClient = agentSession.receive()
35-
.map(WebSocketMessage::getPayload)
36-
.map(dataBuffer -> clientSession.binaryMessage(factory -> dataBuffer))
37-
.as(clientSession::send);
38-
39-
// Run both directions in parallel, complete when both are done
40-
return Mono.zip(clientToAgent, agentToClient).then();
41-
});
29+
try {
30+
String host = (String) clientSession.getAttributes().get("host");
31+
var decryptedHost = cryptoService.decrypt(host); // Ensure host is decrypted if necessary
32+
String sessionId = (String) clientSession.getAttributes().get("sessionId");
33+
String chatGroupId = (String) clientSession.getAttributes().get("chatGroupId");
34+
String ztat = (String) clientSession.getAttributes().get("ztat");
35+
log.info("Handling WebSocket connection for host: {}, sessionId: {}, chatGroupId: {}, ztat: {}",
36+
decryptedHost, sessionId, chatGroupId, ztat);
37+
URI agentUri = agentLocator.resolveWebSocketUri(decryptedHost, sessionId, chatGroupId, ztat);
38+
39+
ReactorNettyWebSocketClient proxyClient = new ReactorNettyWebSocketClient();
40+
41+
return proxyClient.execute(agentUri, agentSession -> {
42+
// Forward messages from client to agent
43+
Mono<Void> clientToAgent = clientSession.receive()
44+
.map(WebSocketMessage::getPayload)
45+
.map(dataBuffer -> agentSession.binaryMessage(factory -> dataBuffer))
46+
.as(agentSession::send);
47+
48+
// Forward messages from agent to client
49+
Mono<Void> agentToClient = agentSession.receive()
50+
.map(WebSocketMessage::getPayload)
51+
.map(dataBuffer -> clientSession.binaryMessage(factory -> dataBuffer))
52+
.as(clientSession::send);
53+
54+
// Run both directions in parallel, complete when both are done
55+
return Mono.zip(clientToAgent, agentToClient).then();
56+
});
57+
} catch (GeneralSecurityException ex) {
58+
throw new RuntimeException("Failed to decrypt host", ex);
59+
}
4260
}
4361
}

‎api/src/main/java/io/sentrius/sso/websocket/WebSocketRouteConfig.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ public class WebSocketRouteConfig {
2121
@Bean
2222
public WebSocketHandlerMapping webSocketMapping() {
2323
Map<String, WebSocketHandler> map = new HashMap<>();
24-
map.put("/api/v1/agents/ws/{agentId}", agentWebSocketProxyHandler);
24+
map.put("/api/v1/agents/ws/{host}/{sessionId}/{chatGroupId}/{ztat}", agentWebSocketProxyHandler);
2525

2626
WebSocketHandlerMapping mapping = new WebSocketHandlerMapping();
2727
mapping.setUrlMap(map);

‎api/src/main/resources/static/js/chat.js‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,10 @@ class ChatSession {
7373

7474

7575
// Step 3: Open WebSocket with ZTAT token
76-
const uri = `${phost}/api/v1/chat/attach/subscribe?sessionId=${encodeURIComponent(this.sessionId)}&chatGroupId=${this.chatGroupId}&ztat=${encodeURIComponent(jwt)}`;
76+
//const uri = `${phost}/api/v1/chat/attach/subscribe?sessionId=${encodeURIComponent(this.sessionId)}&chatGroupId=${this.chatGroupId}&ztat=${encodeURIComponent(jwt)}`;
77+
//const uri = `/api/v1/agents/ws/${encodeURIComponent(phost)}/${encodeURIComponent(this.sessionId)}/${encodeURIComponent(this.chatGroupId)}/${encodeURIComponent(jwt)}`;
78+
const uri = `/api/v1/agents/ws?phost=${encodeURIComponent(phost)}&sessionId=${encodeURIComponent(this.sessionId)}&chatGroupId=${encodeURIComponent(this.chatGroupId)}&jwt=${encodeURIComponent(jwt)}`;
79+
7780
console.log("Connecting to chat server with ZTAT at:", uri);
7881
this.connection = new WebSocket(uri);
7982

‎dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,13 @@ public List<AgentDTO> getAllAgents(boolean encryptId, List<String> filteredIds,
129129
dtoBuilder.agentName(heartbeat.getAgentName());
130130
var callback = callbackUrls.get(heartbeat.getAgentId());
131131
if (callback != null) {
132-
dtoBuilder.agentCallback(callback);
132+
try {
133+
134+
var encryptedCallback = cryptoService.encrypt(callback); // Ensure callback is decrypted
135+
dtoBuilder.agentCallback(encryptedCallback);
136+
} catch (GeneralSecurityException e) {
137+
throw new RuntimeException("Error encrypting callback URL", e);
138+
}
133139
}
134140
}
135141
if (encryptId){

0 commit comments

Comments
 (0)