From bda04a188b72ea2b8515540770db9fdf79636943 Mon Sep 17 00:00:00 2001
From: jeffcheasey88 <66554203+jeffcheasey88@users.noreply.github.com>
Date: Sat, 23 Sep 2023 17:40:51 +0200
Subject: [PATCH] Add WebSocket Authentification

---
 src/dev/peerat/framework/Context.java     |  4 ++++
 src/dev/peerat/framework/RouteMapper.java | 21 +++++++++++++++------
 src/dev/peerat/framework/Router.java      |  6 +++---
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/src/dev/peerat/framework/Context.java b/src/dev/peerat/framework/Context.java
index 9a0142b..8ea54c5 100644
--- a/src/dev/peerat/framework/Context.java
+++ b/src/dev/peerat/framework/Context.java
@@ -29,6 +29,10 @@ public class Context{
 		return user != null;
 	}
 	
+	void setUser(User user){
+		this.user = user;
+	}
+	
 	public <U extends User> U getUser(){
 		return (U) user;
 	}
diff --git a/src/dev/peerat/framework/RouteMapper.java b/src/dev/peerat/framework/RouteMapper.java
index fe84a13..cf323ef 100644
--- a/src/dev/peerat/framework/RouteMapper.java
+++ b/src/dev/peerat/framework/RouteMapper.java
@@ -4,7 +4,11 @@ import java.security.MessageDigest;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
-public class RouteMapper{
+import dev.peerat.framework.utils.json.JsonMap;
+
+public class RouteMapper<U extends User>{
+	
+	private Router<U> router;
 	
 	private Response[] responses;
 	private Route[] routes;
@@ -12,7 +16,8 @@ public class RouteMapper{
 	private int[] calls;
 	private int dif;
 	
-	public RouteMapper(){
+	public RouteMapper(Router<U> router){
+		this.router = router;
 		this.responses = new Response[0];
 		this.routes = new Route[0];
 		this.patterns = new Pattern[0];
@@ -44,14 +49,18 @@ public class RouteMapper{
 				Pattern pattern = patterns[i];
 				matcher = pattern.matcher(path);
 				if(matcher.matches()){
-					if((!context.isLogged()) && route.needLogin()){
-						writer.response(401, "Access-Control-Allow-Origin: *");
-						return true;
-					}
 					if(route.websocket()){
 						switchToWebSocket(reader, writer);
 						reader = new WebSocketReader(reader);
 						writer = new WebSocketWriter(writer);
+						if(route.needLogin()){
+							String token = reader.<JsonMap>readJson().get("token");
+							context.setUser(this.router.getUser(token));
+						}
+					}
+					if((!context.isLogged()) && route.needLogin()){
+						writer.response(401, router.getDefaultHeaders(context.getType()));
+						return true;
 					}
 					order(i);
 					result = response;
diff --git a/src/dev/peerat/framework/Router.java b/src/dev/peerat/framework/Router.java
index b631155..a222eca 100644
--- a/src/dev/peerat/framework/Router.java
+++ b/src/dev/peerat/framework/Router.java
@@ -28,7 +28,7 @@ public class Router<U extends User>{
 	}
 	
 	private Locker<Context> logger;
-	private RouteMapper[] mappers;
+	private RouteMapper<U>[] mappers;
 	private Response noFileFound;
 	private RsaJsonWebKey rsaJsonWebKey;
 	private JwtConsumer jwtConsumer;
@@ -41,7 +41,7 @@ public class Router<U extends User>{
 		this.logger = new Locker<>();
 		int types = RequestType.values().length;
 		this.mappers = new RouteMapper[types];
-		for(RequestType type : RequestType.values()) this.mappers[type.ordinal()] = new RouteMapper();
+		for(RequestType type : RequestType.values()) this.mappers[type.ordinal()] = new RouteMapper<>(this);
 		this.rsaJsonWebKey = RsaJwkGenerator.generateJwk(2048);
 		this.headers = new String[types][0];
 	}
@@ -89,7 +89,7 @@ public class Router<U extends User>{
 	}
 	
 	public Router<U> activeReOrdering(){
-		for(RouteMapper mapper : this.mappers) mapper.activeReOrdering();
+		for(RouteMapper<?> mapper : this.mappers) mapper.activeReOrdering();
 		return this;
 	}