Skip to content

Commit 0c0cfb3

Browse files
guFalconguFalcon
committed
update
1 parent b19fbc6 commit 0c0cfb3

File tree

2 files changed

+84
-39
lines changed

2 files changed

+84
-39
lines changed

‎src/main/java/info/unterrainer/websocketserver/WebsocketServer.java‎

Lines changed: 51 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
@Slf4j
1313
public class WebsocketServer {
1414

15+
private String name;
1516
private String keycloakHost;
1617
private String realm;
1718
private OauthTokenManager tokenManager;
@@ -20,18 +21,35 @@ public class WebsocketServer {
2021
private boolean isOauthEnabled = false;
2122

2223
public WebsocketServer() {
23-
this((Javalin) null);
24+
this("", (Javalin) null);
25+
}
26+
27+
public WebsocketServer(String name) {
28+
this(name, (Javalin) null);
2429
}
2530

2631
public WebsocketServer(Javalin server) {
27-
this(server, null);
32+
this("", server);
33+
}
34+
35+
public WebsocketServer(String name, Javalin server) {
36+
this(name, server, null);
2837
}
2938

3039
public WebsocketServer(WsExceptionHandler<Exception> exceptionHandler) {
31-
this(null, exceptionHandler);
40+
this("", (Javalin) null, exceptionHandler);
41+
}
42+
43+
public WebsocketServer(String name, WsExceptionHandler<Exception> exceptionHandler) {
44+
this(name, (Javalin) null, exceptionHandler);
3245
}
3346

3447
public WebsocketServer(Javalin server, WsExceptionHandler<Exception> exceptionHandler) {
48+
this("", server, exceptionHandler);
49+
}
50+
51+
public WebsocketServer(String name, Javalin server, WsExceptionHandler<Exception> exceptionHandler) {
52+
this.name = name;
3553
try {
3654
wss = server;
3755
if (wss == null)
@@ -41,33 +59,51 @@ public WebsocketServer(Javalin server, WsExceptionHandler<Exception> exceptionHa
4159
wss.wsException(Exception.class, exceptionHandler);
4260

4361
wss.wsException(Exception.class, (e, ctx) -> {
44-
log.error("Uncaught websocket-exception in Websocket-Server: {}", e);
62+
log.error("(" + name + ") Uncaught websocket-exception in Websocket-Server: {}", e);
4563
});
4664
} catch (Exception e) {
47-
log.error("Error initializing Websocket-Server.", e);
65+
log.error("(" + name + ") Error initializing Websocket-Server.", e);
4866
}
4967
}
5068

5169
public WebsocketServer(String keycloakHost, String keycloakRealm) {
52-
this(null, keycloakHost, keycloakRealm);
70+
this("", null, keycloakHost, keycloakRealm);
71+
}
72+
73+
public WebsocketServer(String name, String keycloakHost, String keycloakRealm) {
74+
this(name, null, keycloakHost, keycloakRealm);
5375
}
5476

5577
public WebsocketServer(Javalin server, String keycloakHost, String keycloakRealm) {
56-
this(server, keycloakHost, keycloakRealm, null);
78+
this("", server, keycloakHost, keycloakRealm, null);
79+
}
80+
81+
public WebsocketServer(String name, Javalin server, String keycloakHost, String keycloakRealm) {
82+
this(name, server, keycloakHost, keycloakRealm, null);
5783
}
5884

5985
public WebsocketServer(String keycloakHost, String keycloakRealm, WsExceptionHandler<Exception> exceptionHandler) {
60-
this(null, keycloakHost, keycloakRealm, exceptionHandler);
86+
this("", null, keycloakHost, keycloakRealm, exceptionHandler);
87+
}
88+
89+
public WebsocketServer(String name, String keycloakHost, String keycloakRealm,
90+
WsExceptionHandler<Exception> exceptionHandler) {
91+
this(name, null, keycloakHost, keycloakRealm, exceptionHandler);
6192
}
6293

6394
public WebsocketServer(Javalin server, String keycloakHost, String keycloakRealm,
6495
WsExceptionHandler<Exception> exceptionHandler) {
65-
this(server, exceptionHandler);
96+
this("", server, keycloakHost, keycloakRealm, exceptionHandler);
97+
}
98+
99+
public WebsocketServer(String name, Javalin server, String keycloakHost, String keycloakRealm,
100+
WsExceptionHandler<Exception> exceptionHandler) {
101+
this(name, server, exceptionHandler);
66102
if (keycloakHost == null || keycloakHost.isEmpty()) {
67-
throw new IllegalArgumentException("Keycloak host must not be null or empty.");
103+
throw new IllegalArgumentException("(" + name + ") Keycloak host must not be null or empty.");
68104
}
69105
if (keycloakRealm == null || keycloakRealm.isEmpty()) {
70-
throw new IllegalArgumentException("Keycloak realm must not be null or empty.");
106+
throw new IllegalArgumentException("(" + name + ") Keycloak realm must not be null or empty.");
71107
}
72108
this.keycloakHost = keycloakHost;
73109
this.realm = keycloakRealm;
@@ -77,7 +113,7 @@ public WebsocketServer(Javalin server, String keycloakHost, String keycloakRealm
77113
tokenManager.initPublicKey();
78114
isOauthEnabled = true;
79115
} catch (Exception e) {
80-
log.error("Error initializing OauthTokenManager.", e);
116+
log.error("(" + name + ") Error initializing OauthTokenManager.", e);
81117
return;
82118
}
83119
}
@@ -92,7 +128,7 @@ public WebsocketServer(Javalin server, String keycloakHost, String keycloakRealm
92128
*/
93129
public WebsocketServer start(int port) {
94130
wss.start("0.0.0.0", port);
95-
log.debug("Websocket server started on port: {}", port);
131+
log.debug("(" + name + ") Websocket server started on port: {}", port);
96132
return this;
97133
}
98134

@@ -103,7 +139,8 @@ public WebsocketServer ws(String path, Consumer<WsHandler> ws) {
103139

104140
public WebsocketServer wsOauth(String path, WsOauthHandlerBase handler) {
105141
if (!isOauthEnabled) {
106-
throw new IllegalStateException("Websocket server is not configured for OAuth2/JWT support.");
142+
throw new IllegalStateException(
143+
"(" + name + ") Websocket server is not configured for OAuth2/JWT support.");
107144
}
108145

109146
handler.setTokenHandler(tokenManager);

‎src/main/java/info/unterrainer/websocketserver/WsOauthHandlerBase.java‎

Lines changed: 33 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
@Slf4j
2525
public class WsOauthHandlerBase extends WsHandlerBase {
2626

27+
protected String name;
2728
protected OauthTokenManager tokenHandler;
2829
protected Set<WsConnectContext> clientsConnected = ConcurrentHashMap.newKeySet();
2930
protected Set<WsConnectContext> clientsQuarantined = ConcurrentHashMap.newKeySet();
@@ -35,8 +36,9 @@ public class WsOauthHandlerBase extends WsHandlerBase {
3536
return t;
3637
});
3738

38-
public WsOauthHandlerBase() {
39+
public WsOauthHandlerBase(String name) {
3940
super();
41+
this.name = name;
4042
ShutdownHook.register(() -> {
4143
hb.close();
4244
hb = null;
@@ -50,7 +52,7 @@ public WsOauthHandlerBase() {
5052
s.getRemote().sendPing(ByteBuffer.allocate(1));
5153
} catch (Exception e) {
5254
try {
53-
s.close(1000, "heartbeat failed");
55+
s.close(1000, "(" + name + ") heartbeat failed");
5456
} catch (Exception ignore) {
5557
}
5658
}
@@ -64,49 +66,50 @@ void setTokenHandler(OauthTokenManager tokenHandler) {
6466
}
6567

6668
public void removeClient(Session session) {
67-
log.debug("Removing client: [{}]", session.getRemoteAddress());
69+
log.debug("(" + name + ") Removing client: [{}]", session.getRemoteAddress());
6870
clientsConnected.removeIf(client -> client.session.equals(session));
6971
clientsQuarantined.removeIf(client -> client.session.equals(session));
7072
}
7173

7274
public WsConnectContext getClient(Session session) {
73-
log.debug("Getting client: [{}]", session.getRemoteAddress());
75+
log.debug("(" + name + ") Getting client: [{}]", session.getRemoteAddress());
7476
return clientsConnected.stream().filter(client -> client.session.equals(session)).findFirst().orElse(null);
7577
}
7678

7779
public WsConnectContext getQuarantinedClient(Session session) {
78-
log.debug("Getting quarantined client: [{}]", session.getRemoteAddress());
80+
log.debug("(" + name + ") Getting quarantined client: [{}]", session.getRemoteAddress());
7981
return clientsQuarantined.stream().filter(client -> client.session.equals(session)).findFirst().orElse(null);
8082
}
8183

8284
public boolean isQuarantined(Session session) {
83-
log.debug("Checking if client is quarantined: [{}]", session.getRemoteAddress());
85+
log.debug("(" + name + ") Checking if client is quarantined: [{}]", session.getRemoteAddress());
8486
return clientsQuarantined.stream().anyMatch(client -> client.session.equals(session));
8587
}
8688

8789
public boolean isConnected(Session session) {
88-
log.debug("Checking if client is connected: [{}]", session.getRemoteAddress());
90+
log.debug("(" + name + ") Checking if client is connected: [{}]", session.getRemoteAddress());
8991
return clientsConnected.stream().anyMatch(client -> client.session.equals(session));
9092
}
9193

9294
@Override
9395
public void onConnect(WsConnectContext ctx) throws Exception {
94-
log.debug("New client tries to connect: [{}]", ctx.session.getRemoteAddress());
96+
log.debug("(" + name + ") New client tries to connect: [{}]", ctx.session.getRemoteAddress());
9597
String token = ctx.header("Authorization");
9698
if (token == null || token.isEmpty()) {
97-
log.warn("No token provided for client: [{}]\nSending connection into quarantine.",
99+
log.warn("(" + name + ") No token provided for client: [{}]\nSending connection into quarantine.",
98100
ctx.session.getRemoteAddress());
99101
clientsQuarantined.add(ctx);
100102
return;
101103
}
102-
log.debug("New client token: [{}]", token);
104+
log.debug("(" + name + ") New client token: [{}]", token);
103105
try {
104106
String tenantId = tokenHandler.checkAccess(token);
105107
tenantIdsBySession.put(ctx.session, tenantId);
106108
clientsConnected.add(ctx);
107109
} catch (Exception e) {
108-
log.debug("Token validation failed for client [{}]. Disconnecting.", ctx.session.getRemoteAddress(), e);
109-
ctx.session.close(1000, "Unauthorized access with invalid token");
110+
log.debug("(" + name + ") Token validation failed for client [{}]. Disconnecting.",
111+
ctx.session.getRemoteAddress(), e);
112+
ctx.session.close(1000, "(" + name + ") Unauthorized access with invalid token");
110113
return;
111114
}
112115
}
@@ -116,29 +119,32 @@ public void onMsg(WsMessageContext ctx) throws Exception {
116119
}
117120

118121
public final void onMessage(WsMessageContext ctx) throws Exception {
119-
log.debug("Received from [{}]: [{}]", ctx.session.getRemoteAddress(), ctx.message());
122+
log.debug("(" + name + ") Received from [{}]: [{}]", ctx.session.getRemoteAddress(), ctx.message());
120123
if (isQuarantined(ctx.session)) {
121-
log.warn("Client [{}] is quarantined, checking message for standard authorization-bearer-token.",
124+
log.warn(
125+
"(" + name
126+
+ ") Client [{}] is quarantined, checking message for standard authorization-bearer-token.",
122127
ctx.session.getRemoteAddress());
123128
if (ctx.message() == null || !ctx.message().startsWith("Bearer ")) {
124-
log.warn("Invalid message from quarantined client [{}]. Disconnecting.",
129+
log.warn("(" + name + ") Invalid message from quarantined client [{}]. Disconnecting.",
125130
ctx.session.getRemoteAddress());
126131
removeClient(ctx.session);
127-
ctx.session.close(1000, "Unauthorized access from quarantined client");
132+
ctx.session.close(1000, "(" + name + ") Unauthorized access from quarantined client");
128133
return;
129134
}
130135
try {
131136
String tenantId = tokenHandler.checkAccess(ctx.message());
132137
tenantIdsBySession.put(ctx.session, tenantId);
133138
WsConnectContext client = getQuarantinedClient(ctx.session);
134-
log.debug("Client [{}] passed token validation. Moving from quarantine to connected.",
139+
log.debug("(" + name + ") Client [{}] passed token validation. Moving from quarantine to connected.",
135140
ctx.session.getRemoteAddress());
136141
clientsQuarantined.removeIf(c -> c.session.equals(ctx.session));
137142
clientsConnected.add(client);
138143
return;
139144
} catch (Exception e) {
140-
log.debug("Token validation failed for client [{}]. Disconnecting.", ctx.session.getRemoteAddress(), e);
141-
ctx.session.close(1000, "Unauthorized access with invalid token");
145+
log.debug("(" + name + ") Token validation failed for client [{}]. Disconnecting.",
146+
ctx.session.getRemoteAddress(), e);
147+
ctx.session.close(1000, "(" + name + ") Unauthorized access with invalid token");
142148
return;
143149
}
144150
}
@@ -147,28 +153,30 @@ public final void onMessage(WsMessageContext ctx) throws Exception {
147153

148154
@Override
149155
public void onBinaryMessage(WsBinaryMessageContext ctx) throws Exception {
150-
log.debug("Received binary message from [{}]: [{}] bytes", ctx.session.getRemoteAddress(), ctx.data().length);
156+
log.debug("(" + name + ") Received binary message from [{}]: [{}] bytes", ctx.session.getRemoteAddress(),
157+
ctx.data().length);
151158
if (isQuarantined(ctx.session)) {
152-
log.warn("Invalid Message from quarantined client [{}]. Disconnecting.", ctx.session.getRemoteAddress());
159+
log.warn("(" + name + ") Invalid Message from quarantined client [{}]. Disconnecting.",
160+
ctx.session.getRemoteAddress());
153161
removeClient(ctx.session);
154-
ctx.session.close(1000, "Unauthorized access from quarantined client");
162+
ctx.session.close(1000, "(" + name + ") Unauthorized access from quarantined client");
155163
return;
156164
}
157165
}
158166

159167
@Override
160168
public void onClose(WsCloseContext ctx) throws Exception {
161-
log.debug("Disconnected client: [{}]", ctx.session.getRemoteAddress());
169+
log.debug("(" + name + ") Disconnected client: [{}]", ctx.session.getRemoteAddress());
162170
removeClient(ctx.session);
163171
}
164172

165173
@Override
166174
public void onError(WsErrorContext ctx) throws Exception {
167175
Throwable t = ctx.error();
168176
if (t instanceof EOFException || t instanceof IOException) {
169-
log.debug("Client disconnected [{}].", ctx.session.getRemoteAddress());
177+
log.debug("(" + name + ") Client disconnected [{}].", ctx.session.getRemoteAddress());
170178
} else {
171-
log.error("Unexpected error on [{}].", ctx.session.getRemoteAddress(), t);
179+
log.error("(" + name + ") Unexpected error on [{}].", ctx.session.getRemoteAddress(), t);
172180
}
173181
removeClient(ctx.session);
174182
}

0 commit comments

Comments
 (0)