From 49483e0420ca8c7ff1afb31dacc7a427f93c7a3a Mon Sep 17 00:00:00 2001
From: Anton Tananaev <anton@traccar.org>
Date: Sun, 5 May 2024 09:50:19 -0700
Subject: Support WebSocket auth token

---
 .../java/org/traccar/api/AsyncSocketServlet.java   | 30 +++++++++++++++++-----
 1 file changed, 24 insertions(+), 6 deletions(-)

(limited to 'src/main/java/org/traccar/api/AsyncSocketServlet.java')

diff --git a/src/main/java/org/traccar/api/AsyncSocketServlet.java b/src/main/java/org/traccar/api/AsyncSocketServlet.java
index cd2c1639e..e1e7ee977 100644
--- a/src/main/java/org/traccar/api/AsyncSocketServlet.java
+++ b/src/main/java/org/traccar/api/AsyncSocketServlet.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2015 - 2022 Anton Tananaev (anton@traccar.org)
+ * Copyright 2015 - 2024 Anton Tananaev (anton@traccar.org)
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,6 +19,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import org.eclipse.jetty.websocket.server.JettyWebSocketServlet;
 import org.eclipse.jetty.websocket.server.JettyWebSocketServletFactory;
 import org.traccar.api.resource.SessionResource;
+import org.traccar.api.security.LoginService;
 import org.traccar.config.Config;
 import org.traccar.config.Keys;
 import org.traccar.session.ConnectionManager;
@@ -27,7 +28,12 @@ import org.traccar.storage.Storage;
 import jakarta.inject.Inject;
 import jakarta.inject.Singleton;
 import jakarta.servlet.http.HttpSession;
+import org.traccar.storage.StorageException;
+
+import java.io.IOException;
+import java.security.GeneralSecurityException;
 import java.time.Duration;
+import java.util.List;
 
 @Singleton
 public class AsyncSocketServlet extends JettyWebSocketServlet {
@@ -36,25 +42,37 @@ public class AsyncSocketServlet extends JettyWebSocketServlet {
     private final ObjectMapper objectMapper;
     private final ConnectionManager connectionManager;
     private final Storage storage;
+    private final LoginService loginService;
 
     @Inject
     public AsyncSocketServlet(
-            Config config, ObjectMapper objectMapper, ConnectionManager connectionManager, Storage storage) {
+            Config config, ObjectMapper objectMapper, ConnectionManager connectionManager, Storage storage,
+            LoginService loginService) {
         this.config = config;
         this.objectMapper = objectMapper;
         this.connectionManager = connectionManager;
         this.storage = storage;
+        this.loginService = loginService;
     }
 
     @Override
     public void configure(JettyWebSocketServletFactory factory) {
         factory.setIdleTimeout(Duration.ofMillis(config.getLong(Keys.WEB_TIMEOUT)));
         factory.setCreator((req, resp) -> {
-            if (req.getSession() != null) {
-                Long userId = (Long) ((HttpSession) req.getSession()).getAttribute(SessionResource.USER_ID_KEY);
-                if (userId != null) {
-                    return new AsyncSocket(objectMapper, connectionManager, storage, userId);
+            Long userId = null;
+            List<String> tokens = req.getParameterMap().get("token");
+            if (tokens != null && !tokens.isEmpty()) {
+                String token = tokens.iterator().next();
+                try {
+                    userId = loginService.login(token).getUser().getId();
+                } catch (StorageException | GeneralSecurityException | IOException e) {
+                    throw new RuntimeException(e);
                 }
+            } else if (req.getSession() != null) {
+                userId = (Long) ((HttpSession) req.getSession()).getAttribute(SessionResource.USER_ID_KEY);
+            }
+            if (userId != null) {
+                return new AsyncSocket(objectMapper, connectionManager, storage, userId);
             }
             return null;
         });
-- 
cgit v1.2.3