add advanced error-handling

Author: guFalcon

pom.xml

@@ -10,7 +10,7 @@
 
 	<modelVersion>4.0.0</modelVersion>
 	<artifactId>websocket-server</artifactId>
-	<version>1.0.7</version>
+	<version>1.0.8</version>
 	<name>WebsocketServer</name>
 	<packaging>jar</packaging>
 

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

@@ -5,6 +5,7 @@
 
 import info.unterrainer.oauthtokenmanager.OauthTokenManager;
 import io.javalin.Javalin;
+import io.javalin.websocket.WsExceptionHandler;
 import io.javalin.websocket.WsHandler;
 import lombok.extern.slf4j.Slf4j;
 
@@ -19,27 +20,47 @@ public class WebsocketServer {
 	private boolean isOauthEnabled = false;
 
 	public WebsocketServer() {
-		wss = Javalin.create();
+		this(null);
 	}
 
-	public WebsocketServer(String keycloakHost, String keycloakRealm) {
-		this.keycloakHost = keycloakHost;
-		this.realm = keycloakRealm;
-
+	public WebsocketServer(WsExceptionHandler<Exception> exceptionHandler) {
 		try {
-			tokenManager = new OauthTokenManager(this.keycloakHost, this.realm);
-			tokenManager.initPublicKey();
 			wss = Javalin.create();
 			wss.exception(Exception.class, (e, ctx) -> {
 				log.error("Uncaught exception in Websocket-Server: {}", e);
 			});
+			if (exceptionHandler != null)
+				wss.wsException(Exception.class, exceptionHandler);
 			wss.wsException(Exception.class, (e, ctx) -> {
 				log.error("Uncaught websocket-exception in Websocket-Server: {}", e);
 			});
+		} catch (Exception e) {
+			log.error("Error initializing Websocket-Server.", e);
+		}
+	}
+
+	public WebsocketServer(String keycloakHost, String keycloakRealm) {
+		this(keycloakHost, keycloakRealm, null);
+	}
+
+	public WebsocketServer(String keycloakHost, String keycloakRealm, WsExceptionHandler<Exception> exceptionHandler) {
+		this(exceptionHandler);
+		if (keycloakHost == null || keycloakHost.isEmpty()) {
+			throw new IllegalArgumentException("Keycloak host must not be null or empty.");
+		}
+		if (keycloakRealm == null || keycloakRealm.isEmpty()) {
+			throw new IllegalArgumentException("Keycloak realm must not be null or empty.");
+		}
+		this.keycloakHost = keycloakHost;
+		this.realm = keycloakRealm;
+
+		try {
+			tokenManager = new OauthTokenManager(this.keycloakHost, this.realm);
+			tokenManager.initPublicKey();
 			isOauthEnabled = true;
 		} catch (Exception e) {
-			// Exceptions will terminate a request later on, but should not terminate the
-			// main-thread here.
+			log.error("Error initializing OauthTokenManager.", e);
+			return;
 		}
 	}
 

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

@@ -68,7 +68,7 @@ public void onConnect(WsConnectContext ctx) throws Exception {
 			clientsConnected.add(ctx);
 		} catch (Exception e) {
 			log.debug("Token validation failed for client [{}]. Disconnecting.", ctx.session.getRemoteAddress(), e);
-			ctx.session.close();
+			ctx.session.close(1000, "Unauthorized access with invalid token");
 			return;
 		}
 	}
@@ -83,7 +83,7 @@ public void onMessage(WsMessageContext ctx) throws Exception {
 				log.warn("Invalid message from quarantined client [{}]. Disconnecting.",
 						ctx.session.getRemoteAddress());
 				removeClient(ctx.session);
-				ctx.session.close();
+				ctx.session.close(1000, "Unauthorized access from quarantined client");
 				return;
 			}
 			try {
@@ -94,8 +94,8 @@ public void onMessage(WsMessageContext ctx) throws Exception {
 				clientsQuarantined.removeIf(c -> c.session.equals(ctx.session));
 				clientsConnected.add(client);
 			} catch (Exception e) {
-				ctx.session.close();
 				log.debug("Token validation failed for client [{}]. Disconnecting.", ctx.session.getRemoteAddress(), e);
+				ctx.session.close(1000, "Unauthorized access with invalid token");
 				return;
 			}
 		}
@@ -107,7 +107,7 @@ public void onBinaryMessage(WsBinaryMessageContext ctx) throws Exception {
 		if (isQuarantined(ctx.session)) {
 			log.warn("Invalid Message from quarantined client [{}]. Disconnecting.", ctx.session.getRemoteAddress());
 			removeClient(ctx.session);
-			ctx.session.close();
+			ctx.session.close(1000, "Unauthorized access from quarantined client");
 			return;
 		}
 	}