UNDERTOW-224 Provide a means to get "peer" connections to a WebSocketChannel

stuartwdouglas · stuartwdouglas · commit 17bf53ec804b · 2014-05-16T14:28:25.000-05:00
diff --git a/core/src/main/java/io/undertow/websockets/WebSocketProtocolHandshakeHandler.java b/core/src/main/java/io/undertow/websockets/WebSocketProtocolHandshakeHandler.java
@@ -33,8 +33,10 @@
 import org.xnio.StreamConnection;
 
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * {@link HttpHandler} which will process the {@link HttpServerExchange} and do the actual handshake/upgrade
@@ -52,6 +54,8 @@ public class WebSocketProtocolHandshakeHandler implements HttpHandler {
 
     private final WebSocketConnectionCallback callback;
 
+    private final Set<WebSocketChannel> peerConnections = Collections.newSetFromMap(new ConcurrentHashMap<WebSocketChannel, Boolean>());
+
     /**
      * The handler that is invoked if there are no web socket headers
      */
@@ -169,7 +173,7 @@ public void handleRequest(final HttpServerExchange exchange) throws Exception {
             next.handleRequest(exchange);
             return;
         }
-        final AsyncWebSocketHttpServerExchange facade = new AsyncWebSocketHttpServerExchange(exchange);
+        final AsyncWebSocketHttpServerExchange facade = new AsyncWebSocketHttpServerExchange(exchange, peerConnections);
         Handshake handshaker = null;
         for (Handshake method : handshakes) {
             if (method.matches(facade)) {
@@ -187,6 +191,7 @@ public void handleRequest(final HttpServerExchange exchange) throws Exception {
                     @Override
                     public void handleUpgrade(StreamConnection streamConnection, HttpServerExchange exchange) {
                         WebSocketChannel channel = selected.createChannel(facade, streamConnection, facade.getBufferPool());
+                        peerConnections.add(channel);
                         callback.onConnect(facade, channel);
                     }
                 });
@@ -196,4 +201,8 @@ public void handleUpgrade(StreamConnection streamConnection, HttpServerExchange
             handshaker.handshake(facade);
         }
     }
+
+    public Set<WebSocketChannel> getPeerConnections() {
+        return peerConnections;
+    }
 }
diff --git a/core/src/main/java/io/undertow/websockets/client/WebSocket13ClientHandshake.java b/core/src/main/java/io/undertow/websockets/client/WebSocket13ClientHandshake.java
@@ -38,6 +38,7 @@
 import java.security.SecureRandom;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
@@ -63,7 +64,7 @@ public WebSocket13ClientHandshake(final URI url) {
 
     @Override
     public WebSocketChannel createChannel(final StreamConnection channel, final String wsUri, final Pool<ByteBuffer> bufferPool) {
-        return new WebSocket13Channel(channel, bufferPool, wsUri, negotiation != null ? negotiation.getSelectedSubProtocol() : "", true, false);
+        return new WebSocket13Channel(channel, bufferPool, wsUri, negotiation != null ? negotiation.getSelectedSubProtocol() : "", true, false, new HashSet<WebSocketChannel>());
     }
 
 
diff --git a/core/src/main/java/io/undertow/websockets/core/WebSocketChannel.java b/core/src/main/java/io/undertow/websockets/core/WebSocketChannel.java
@@ -20,6 +20,7 @@
 import io.undertow.server.protocol.framed.AbstractFramedChannel;
 import io.undertow.server.protocol.framed.FrameHeaderData;
 import org.xnio.ChannelExceptionHandler;
+import org.xnio.ChannelListener;
 import org.xnio.ChannelListeners;
 import org.xnio.IoUtils;
 import org.xnio.Pool;
@@ -62,24 +63,37 @@ public abstract class WebSocketChannel extends AbstractFramedChannel<WebSocketCh
 
     protected StreamSourceFrameChannel fragmentedChannel;
 
+    /**
+     * Represents all web socket channels that are attached to the same endpoint.
+     */
+    private final Set<WebSocketChannel> peerConnections;
+
     /**
      * Create a new {@link WebSocketChannel}
      * 8
      *
      * @param connectedStreamChannel The {@link org.xnio.channels.ConnectedStreamChannel} over which the WebSocket Frames should get send and received.
      *                               Be aware that it already must be "upgraded".
      * @param bufferPool             The {@link org.xnio.Pool} which will be used to acquire {@link java.nio.ByteBuffer}'s from.
-     * @param version                The {@link io.undertow.websockets.core.WebSocketVersion} of the {@link io.undertow.websockets.core.WebSocketChannel}
+     * @param version                The {@link WebSocketVersion} of the {@link WebSocketChannel}
      * @param wsUrl                  The url for which the channel was created.
      * @param client
+     * @param peerConnections        The concurrent set that is used to track open connections associtated with an endpoint
      */
-    protected WebSocketChannel(final StreamConnection connectedStreamChannel, Pool<ByteBuffer> bufferPool, WebSocketVersion version, String wsUrl, String subProtocol, final boolean client, boolean extensionsSupported) {
+    protected WebSocketChannel(final StreamConnection connectedStreamChannel, Pool<ByteBuffer> bufferPool, WebSocketVersion version, String wsUrl, String subProtocol, final boolean client, boolean extensionsSupported, Set<WebSocketChannel> peerConnections) {
         super(connectedStreamChannel, bufferPool, new WebSocketFramePriority(), null);
         this.client = client;
         this.version = version;
         this.wsUrl = wsUrl;
         this.extensionsSupported = extensionsSupported;
         this.subProtocol = subProtocol;
+        this.peerConnections = peerConnections;
+        addCloseTask(new ChannelListener<WebSocketChannel>() {
+            @Override
+            public void handleEvent(WebSocketChannel channel) {
+                WebSocketChannel.this.peerConnections.remove(WebSocketChannel.this);
+            }
+        });
     }
 
     @Override
@@ -343,6 +357,15 @@ protected WebSocketFramePriority getFramePriority() {
         return (WebSocketFramePriority) super.getFramePriority();
     }
 
+    /**
+     * Returns all 'peer' web socket connections that were created from the same endpoint.
+     *
+     *
+     * @return all 'peer' web socket connections
+     */
+    public Set<WebSocketChannel> getPeerConnections() {
+        return Collections.unmodifiableSet(peerConnections);
+    }
 
     /**
      * Interface that represents a frame channel that is in the process of being created
diff --git a/core/src/main/java/io/undertow/websockets/core/protocol/version07/Hybi07Handshake.java b/core/src/main/java/io/undertow/websockets/core/protocol/version07/Hybi07Handshake.java
@@ -100,6 +100,6 @@ protected final String solve(final String nonceBase64) throws NoSuchAlgorithmExc
 
     @Override
     public WebSocketChannel createChannel(WebSocketHttpExchange exchange, final StreamConnection channel, final Pool<ByteBuffer> pool) {
-        return new WebSocket07Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions);
+        return new WebSocket07Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions, exchange.getPeerConnections());
     }
 }
diff --git a/core/src/main/java/io/undertow/websockets/core/protocol/version07/WebSocket07Channel.java b/core/src/main/java/io/undertow/websockets/core/protocol/version07/WebSocket07Channel.java
@@ -33,6 +33,7 @@
 import org.xnio.StreamConnection;
 
 import java.nio.ByteBuffer;
+import java.util.Set;
 
 
 /**
@@ -83,8 +84,8 @@ private enum State {
      * @param wsUrl      The url for which the {@link WebSocket07Channel} was created.
      */
     public WebSocket07Channel(StreamConnection channel, Pool<ByteBuffer> bufferPool,
-                              String wsUrl, String subProtocol, final boolean client, boolean allowExtensions) {
-        super(channel, bufferPool, WebSocketVersion.V08, wsUrl, subProtocol, client, allowExtensions);
+                              String wsUrl, String subProtocol, final boolean client, boolean allowExtensions, Set<WebSocketChannel> openConnections) {
+        super(channel, bufferPool, WebSocketVersion.V08, wsUrl, subProtocol, client, allowExtensions, openConnections);
     }
 
     @Override
diff --git a/core/src/main/java/io/undertow/websockets/core/protocol/version08/Hybi08Handshake.java b/core/src/main/java/io/undertow/websockets/core/protocol/version08/Hybi08Handshake.java
@@ -47,7 +47,7 @@ public Hybi08Handshake(Set<String> subprotocols, boolean allowExtensions) {
 
     @Override
     public WebSocketChannel createChannel(final WebSocketHttpExchange exchange, final StreamConnection channel, final Pool<ByteBuffer> pool) {
-        return new WebSocket08Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions);
+        return new WebSocket08Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions, exchange.getPeerConnections());
 
     }
 }
diff --git a/core/src/main/java/io/undertow/websockets/core/protocol/version08/WebSocket08Channel.java b/core/src/main/java/io/undertow/websockets/core/protocol/version08/WebSocket08Channel.java
@@ -17,12 +17,14 @@
  */
 package io.undertow.websockets.core.protocol.version08;
 
+import io.undertow.websockets.core.WebSocketChannel;
 import io.undertow.websockets.core.WebSocketVersion;
 import io.undertow.websockets.core.protocol.version07.WebSocket07Channel;
 import org.xnio.Pool;
 import org.xnio.StreamConnection;
 
 import java.nio.ByteBuffer;
+import java.util.Set;
 
 
 /**
@@ -31,8 +33,8 @@
  * @author <a href="mailto:nmaurer@redhat.com">Norman Maurer</a>
  */
 public class WebSocket08Channel extends WebSocket07Channel {
-    public WebSocket08Channel(StreamConnection channel, Pool<ByteBuffer> bufferPool, String wsUrl, String subProtocols, final boolean client, boolean allowExtensions) {
-        super(channel, bufferPool, wsUrl, subProtocols, client, allowExtensions);
+    public WebSocket08Channel(StreamConnection channel, Pool<ByteBuffer> bufferPool, String wsUrl, String subProtocols, final boolean client, boolean allowExtensions, Set<WebSocketChannel> openConnections) {
+        super(channel, bufferPool, wsUrl, subProtocols, client, allowExtensions, openConnections);
     }
 
     @Override
diff --git a/core/src/main/java/io/undertow/websockets/core/protocol/version13/Hybi13Handshake.java b/core/src/main/java/io/undertow/websockets/core/protocol/version13/Hybi13Handshake.java
@@ -70,6 +70,6 @@ protected void handshakeInternal(final WebSocketHttpExchange exchange) {
 
     @Override
     public WebSocketChannel createChannel(WebSocketHttpExchange exchange, final StreamConnection channel, final Pool<ByteBuffer> pool) {
-        return new WebSocket13Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions);
+        return new WebSocket13Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions, exchange.getPeerConnections());
     }
 }
diff --git a/core/src/main/java/io/undertow/websockets/core/protocol/version13/WebSocket13Channel.java b/core/src/main/java/io/undertow/websockets/core/protocol/version13/WebSocket13Channel.java
@@ -17,12 +17,14 @@
  */
 package io.undertow.websockets.core.protocol.version13;
 
+import io.undertow.websockets.core.WebSocketChannel;
 import io.undertow.websockets.core.WebSocketVersion;
 import io.undertow.websockets.core.protocol.version07.WebSocket07Channel;
 import org.xnio.Pool;
 import org.xnio.StreamConnection;
 
 import java.nio.ByteBuffer;
+import java.util.Set;
 
 /**
  *
@@ -31,8 +33,8 @@
  * @author <a href="mailto:nmaurer@redhat.com">Norman Maurer</a>
  */
 public class WebSocket13Channel extends WebSocket07Channel {
-    public WebSocket13Channel(StreamConnection channel, Pool<ByteBuffer> bufferPool, String wsUrl, String subProtocols, final boolean client, boolean allowExtensions) {
-        super(channel, bufferPool, wsUrl, subProtocols, client, allowExtensions);
+    public WebSocket13Channel(StreamConnection channel, Pool<ByteBuffer> bufferPool, String wsUrl, String subProtocols, final boolean client, boolean allowExtensions, Set<WebSocketChannel> openConnections) {
+        super(channel, bufferPool, wsUrl, subProtocols, client, allowExtensions, openConnections);
     }
 
     @Override
diff --git a/core/src/main/java/io/undertow/websockets/spi/AsyncWebSocketHttpServerExchange.java b/core/src/main/java/io/undertow/websockets/spi/AsyncWebSocketHttpServerExchange.java
@@ -28,6 +28,7 @@
 import io.undertow.util.AttachmentKey;
 import io.undertow.util.HeaderMap;
 import io.undertow.util.HttpString;
+import io.undertow.websockets.core.WebSocketChannel;
 import org.xnio.ChannelListener;
 import org.xnio.FinishedIoFuture;
 import org.xnio.FutureResult;
@@ -47,6 +48,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * @author Stuart Douglas
@@ -55,9 +57,11 @@ public class AsyncWebSocketHttpServerExchange implements WebSocketHttpExchange {
 
     private final HttpServerExchange exchange;
     private Sender sender;
+    private final Set<WebSocketChannel> peerConnections;
 
-    public AsyncWebSocketHttpServerExchange(final HttpServerExchange exchange) {
+    public AsyncWebSocketHttpServerExchange(final HttpServerExchange exchange, Set<WebSocketChannel> peerConnections) {
         this.exchange = exchange;
+        this.peerConnections = peerConnections;
     }
 
 
@@ -273,4 +277,9 @@ public boolean isUserInRole(String role) {
         }
         return authenticatedAccount.getRoles().contains(role);
     }
+
+    @Override
+    public Set<WebSocketChannel> getPeerConnections() {
+        return peerConnections;
+    }
 }
diff --git a/core/src/main/java/io/undertow/websockets/spi/BlockingWebSocketHttpServerExchange.java b/core/src/main/java/io/undertow/websockets/spi/BlockingWebSocketHttpServerExchange.java
@@ -19,6 +19,7 @@
 package io.undertow.websockets.spi;
 
 import io.undertow.server.HttpServerExchange;
+import io.undertow.websockets.core.WebSocketChannel;
 import org.xnio.FinishedIoFuture;
 import org.xnio.FutureResult;
 import org.xnio.IoFuture;
@@ -28,6 +29,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.util.Set;
 
 /**
  * @author Stuart Douglas
@@ -37,8 +39,8 @@ public class BlockingWebSocketHttpServerExchange extends AsyncWebSocketHttpServe
     private final OutputStream out;
     private final InputStream in;
 
-    public BlockingWebSocketHttpServerExchange(final HttpServerExchange exchange) {
-        super(exchange);
+    public BlockingWebSocketHttpServerExchange(final HttpServerExchange exchange, Set<WebSocketChannel> peerConnections) {
+        super(exchange, peerConnections);
         out = exchange.getOutputStream();
         in = exchange.getInputStream();
     }
diff --git a/core/src/main/java/io/undertow/websockets/spi/WebSocketHttpExchange.java b/core/src/main/java/io/undertow/websockets/spi/WebSocketHttpExchange.java
@@ -20,6 +20,7 @@
 
 import io.undertow.server.HttpUpgradeListener;
 import io.undertow.util.AttachmentKey;
+import io.undertow.websockets.core.WebSocketChannel;
 import org.xnio.IoFuture;
 import org.xnio.Pool;
 
@@ -28,6 +29,7 @@
 import java.security.Principal;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 
 /**
@@ -156,4 +158,6 @@ public interface WebSocketHttpExchange extends Closeable {
     Principal getUserPrincipal();
 
     boolean isUserInRole(String role);
+
+    Set<WebSocketChannel> getPeerConnections();
 }
diff --git a/examples/src/main/java/io/undertow/examples/chat/ChatServer.java b/examples/src/main/java/io/undertow/examples/chat/ChatServer.java
@@ -21,17 +21,12 @@
 import io.undertow.Undertow;
 import io.undertow.examples.UndertowExample;
 import io.undertow.server.handlers.resource.ClassPathResourceManager;
+import io.undertow.websockets.WebSocketConnectionCallback;
 import io.undertow.websockets.core.AbstractReceiveListener;
 import io.undertow.websockets.core.BufferedTextMessage;
 import io.undertow.websockets.core.WebSocketChannel;
 import io.undertow.websockets.core.WebSockets;
-import io.undertow.websockets.WebSocketConnectionCallback;
 import io.undertow.websockets.spi.WebSocketHttpExchange;
-import org.xnio.ChannelListener;
-
-import java.nio.channels.Channel;
-import java.util.ArrayList;
-import java.util.List;
 
 import static io.undertow.Handlers.path;
 import static io.undertow.Handlers.resource;
@@ -43,44 +38,30 @@
 @UndertowExample("Chat")
 public class ChatServer {
 
-    private static final List<WebSocketChannel> sessions = new ArrayList<WebSocketChannel>();
-
     public static void main(final String[] args) {
 
         System.out.println("To see chat in action is to open two different browsers and point them at http://localhost:8080");
 
         Undertow server = Undertow.builder()
-                .addListener(8080, "localhost")
+                .addHttpListener(8080, "localhost")
                 .setHandler(path()
                         .addPrefixPath("/myapp", websocket(new WebSocketConnectionCallback() {
 
                             @Override
                             public void onConnect(WebSocketHttpExchange exchange, WebSocketChannel channel) {
-                                synchronized (sessions) {
-                                    sessions.add(channel);
-                                    channel.getCloseSetter().set(new ChannelListener<Channel>() {
-                                        @Override
-                                        public void handleEvent(Channel channel) {
-                                            synchronized (sessions) {
-                                                sessions.remove(channel);
-                                            }
-                                        }
-                                    });
-                                    channel.getReceiveSetter().set(new AbstractReceiveListener() {
+                                channel.getReceiveSetter().set(new AbstractReceiveListener() {
 
-                                        @Override
-                                        protected void onFullTextMessage(WebSocketChannel channel, BufferedTextMessage message) {
-                                            final String messageData = message.getData();
-                                            synchronized (sessions) {
-                                                for (WebSocketChannel session : sessions) {
-                                                    WebSockets.sendText(messageData, session, null);
-                                                }
-                                            }
+                                    @Override
+                                    protected void onFullTextMessage(WebSocketChannel channel, BufferedTextMessage message) {
+                                        final String messageData = message.getData();
+                                        for (WebSocketChannel session : channel.getPeerConnections()) {
+                                            WebSockets.sendText(messageData, session, null);
                                         }
-                                    });
-                                    channel.resumeReceives();
-                                }
+                                    }
+                                });
+                                channel.resumeReceives();
                             }
+
                         }))
                         .addPrefixPath("/", resource(new ClassPathResourceManager(ChatServer.class.getClassLoader(), ChatServer.class.getPackage()))
                                 .addWelcomeFiles("index.html")))
diff --git a/servlet/src/main/java/io/undertow/servlet/websockets/ServletWebSocketHttpExchange.java b/servlet/src/main/java/io/undertow/servlet/websockets/ServletWebSocketHttpExchange.java
diff --git a/servlet/src/main/java/io/undertow/servlet/websockets/WebSocketServlet.java b/servlet/src/main/java/io/undertow/servlet/websockets/WebSocketServlet.java
diff --git a/websockets-jsr/src/main/java/io/undertow/websockets/jsr/JsrWebSocketFilter.java b/websockets-jsr/src/main/java/io/undertow/websockets/jsr/JsrWebSocketFilter.java

Original file line number	Diff line number	Diff line change
`@@ -100,6 +100,6 @@ protected final String solve(final String nonceBase64) throws NoSuchAlgorithmExc`
`100`	`100`
`101`	`101`	`@Override`
`102`	`102`	`public WebSocketChannel createChannel(WebSocketHttpExchange exchange, final StreamConnection channel, final Pool<ByteBuffer> pool) {`
`103`		`- return new WebSocket07Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions);`
	`103`	`+ return new WebSocket07Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions, exchange.getPeerConnections());`
`104`	`104`	`}`
`105`	`105`	`}`
Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ public Hybi08Handshake(Set<String> subprotocols, boolean allowExtensions) {`
`47`	`47`
`48`	`48`	`@Override`
`49`	`49`	`public WebSocketChannel createChannel(final WebSocketHttpExchange exchange, final StreamConnection channel, final Pool<ByteBuffer> pool) {`
`50`		`- return new WebSocket08Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions);`
	`50`	`+ return new WebSocket08Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions, exchange.getPeerConnections());`
`51`	`51`
`52`	`52`	`}`
`53`	`53`	`}`
Original file line number	Diff line number	Diff line change
`@@ -70,6 +70,6 @@ protected void handshakeInternal(final WebSocketHttpExchange exchange) {`
`70`	`70`
`71`	`71`	`@Override`
`72`	`72`	`public WebSocketChannel createChannel(WebSocketHttpExchange exchange, final StreamConnection channel, final Pool<ByteBuffer> pool) {`
`73`		`- return new WebSocket13Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions);`
	`73`	`+ return new WebSocket13Channel(channel, pool, getWebSocketLocation(exchange), exchange.getResponseHeader(Headers.SEC_WEB_SOCKET_PROTOCOL_STRING), false, allowExtensions, exchange.getPeerConnections());`
`74`	`74`	`}`
`75`	`75`	`}`