Makes SocketContext available to a WsSession (helidon-io#8944)

* Makes SocketContext available to a WsSession to access peer and socket information. See issue 8856. Signed-off-by: Santiago Pericas-Geertsen <[email protected]> * Fixes javadoc --------- Signed-off-by: Santiago Pericas-Geertsen <[email protected]>
arjav-desai · Jul 12, 2024 · 5a54e24 · 5a54e24
1 parent 89caae9
commit 5a54e24
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 2 deletions.
diff --git a/microprofile/websocket/src/main/java/io/helidon/microprofile/tyrus/TyrusConnection.java b/microprofile/websocket/src/main/java/io/helidon/microprofile/tyrus/TyrusConnection.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, 2023 Oracle and/or its affiliates.
+ * Copyright (c) 2022, 2024 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -25,6 +25,7 @@
 
 import io.helidon.common.buffers.BufferData;
 import io.helidon.common.buffers.DataReader;
+import io.helidon.common.socket.SocketContext;
 import io.helidon.http.DateTime;
 import io.helidon.webserver.ConnectionContext;
 import io.helidon.webserver.spi.ServerConnection;
@@ -130,6 +131,11 @@ public Optional<String> subProtocol() {
         return Optional.empty();
     }
 
+    @Override
+    public SocketContext socketContext() {
+        return ctx;
+    }
+
     @Override
     public Duration idleTime() {
         return Duration.between(lastRequestTimestamp, DateTime.timestamp());

diff --git a/webclient/websocket/src/main/java/io/helidon/webclient/websocket/ClientWsConnection.java b/webclient/websocket/src/main/java/io/helidon/webclient/websocket/ClientWsConnection.java
@@ -22,6 +22,7 @@
 import io.helidon.common.buffers.BufferData;
 import io.helidon.common.buffers.DataReader;
 import io.helidon.common.socket.HelidonSocket;
+import io.helidon.common.socket.SocketContext;
 import io.helidon.webclient.api.ClientConnection;
 import io.helidon.websocket.ClientWsFrame;
 import io.helidon.websocket.ServerWsFrame;
@@ -169,6 +170,11 @@ public Optional<String> subProtocol() {
         return Optional.ofNullable(subProtocol);
     }
 
+    @Override
+    public SocketContext socketContext() {
+        return helidonSocket;
+    }
+
     private ClientWsConnection send(ClientWsFrame frame) {
         WsOpCode opCode = frame.opCode();
         if (opCode == WsOpCode.TEXT || opCode == WsOpCode.BINARY) {

diff --git a/...erver/tests/websocket/src/main/java/io/helidon/webserver/tests/websocket/EchoService.java b/...erver/tests/websocket/src/main/java/io/helidon/webserver/tests/websocket/EchoService.java
@@ -40,6 +40,9 @@ public void onOpen(WsSession session) {
         if (subProtocol != null && !subProtocol.equals(p)) {
             throw new InternalError("Invalid sub-protocol in session");
         }
+        if (session.socketContext().remotePeer() == null) {
+            throw new InternalError("Unable to access remote peer info");
+        }
     }
 
     @Override

diff --git a/...sts/websocket/src/test/java/io/helidon/webserver/tests/websocket/WebSocketClientTest.java b/...sts/websocket/src/test/java/io/helidon/webserver/tests/websocket/WebSocketClientTest.java
@@ -88,6 +88,9 @@ public void onMessage(WsSession session, String text, boolean last) {
 
             @Override
             public void onOpen(WsSession session) {
+                if (session.socketContext().remotePeer() == null) {
+                    throw new InternalError("Unable to access remote peer info");
+                }
                 for (String s : text) {
                     session.send(s, false);
                 }

diff --git a/webserver/websocket/src/main/java/io/helidon/webserver/websocket/WsConnection.java b/webserver/websocket/src/main/java/io/helidon/webserver/websocket/WsConnection.java
@@ -25,6 +25,7 @@
 
 import io.helidon.common.buffers.BufferData;
 import io.helidon.common.buffers.DataReader;
+import io.helidon.common.socket.SocketContext;
 import io.helidon.http.DateTime;
 import io.helidon.http.Headers;
 import io.helidon.http.HttpPrologue;
@@ -200,6 +201,11 @@ public Optional<String> subProtocol() {
         return upgradeHeaders.first(WsUpgrader.PROTOCOL);
     }
 
+    @Override
+    public SocketContext socketContext() {
+        return ctx;
+    }
+
     @Override
     public Duration idleTime() {
         return Duration.between(lastRequestTimestamp, DateTime.timestamp());

diff --git a/websocket/src/main/java/io/helidon/websocket/WsSession.java b/websocket/src/main/java/io/helidon/websocket/WsSession.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, 2023 Oracle and/or its affiliates.
+ * Copyright (c) 2022, 2024 Oracle and/or its affiliates.
  *
  * 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 java.util.Optional;
 
 import io.helidon.common.buffers.BufferData;
+import io.helidon.common.socket.SocketContext;
 
 /**
  * WebSocket session.
@@ -82,4 +83,11 @@ public interface WsSession {
     default Optional<String> subProtocol() {
         return Optional.empty();
     }
+
+    /**
+     * The underlying socket context.
+     *
+     * @return socket context
+     */
+    SocketContext socketContext();
 }