Fixed Undertow and added websocket support for it

Author: JanHolger

src/main/java/org/javawebstack/httpserver/adapter/simple/SimpleHTTPSocketServer.java

@@ -2,6 +2,7 @@
 
 import org.javawebstack.httpserver.adapter.IHTTPSocketHandler;
 import org.javawebstack.httpserver.adapter.IHTTPSocketServer;
+import org.javawebstack.httpserver.util.websocket.WebSocketUtil;
 
 import java.io.IOException;
 import java.net.ServerSocket;
@@ -23,7 +24,17 @@ public SimpleHTTPSocketServer() {
                 try {
                     Socket socket = serverSocket.accept();
                     SimpleHTTPSocket httpSocket = new SimpleHTTPSocket(socket);
-                    executorService.execute(() -> handler.handle(httpSocket));
+                    executorService.execute(() -> {
+                        if(httpSocket.getRequestHeaderNames().contains("sec-websocket-key")) {
+                            try {
+                                if(!WebSocketUtil.accept(httpSocket, null))
+                                    return;
+                            } catch (IOException e) {
+                                return;
+                            }
+                        }
+                        handler.handle(httpSocket);
+                    });
                 } catch (IOException exception) {}
             }
         });

src/main/java/org/javawebstack/httpserver/adapter/undertow/StreamSinkOutputStream.java

@@ -0,0 +1,29 @@
+package org.javawebstack.httpserver.adapter.undertow;
+
+import org.xnio.channels.StreamSinkChannel;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.ByteBuffer;
+
+public class StreamSinkOutputStream extends OutputStream {
+
+    private final StreamSinkChannel sink;
+    private final ByteBuffer byteBuffer = ByteBuffer.allocate(1);
+
+    public StreamSinkOutputStream(StreamSinkChannel sink) {
+        this.sink = sink;
+    }
+
+    public void write(int i) throws IOException {
+        byteBuffer.position(0);
+        byteBuffer.put((byte) i);
+        byteBuffer.position(0);
+        sink.write(byteBuffer);
+    }
+
+    public void close() throws IOException {
+        sink.close();
+    }
+
+}

src/main/java/org/javawebstack/httpserver/adapter/undertow/StreamSourceInputStream.java

@@ -0,0 +1,29 @@
+package org.javawebstack.httpserver.adapter.undertow;
+
+import org.xnio.channels.StreamSourceChannel;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+
+public class StreamSourceInputStream extends InputStream {
+
+    private final StreamSourceChannel source;
+    private final ByteBuffer byteBuffer = ByteBuffer.allocate(1);
+
+    public StreamSourceInputStream(StreamSourceChannel source) {
+        this.source = source;
+    }
+
+    public synchronized int read() throws IOException {
+        byteBuffer.position(0);
+        int r;
+        while ((r = source.read(byteBuffer)) == 0)
+            Thread.yield();
+        if(r == -1)
+            return -1;
+        int b = byteBuffer.position(0).get();
+        return b;
+    }
+
+}

…adapter/untertow/UndertowHTTPSocket.java …adapter/undertow/UndertowHTTPSocket.javasrc/main/java/org/javawebstack/httpserver/adapter/untertow/UndertowHTTPSocket.java renamed to src/main/java/org/javawebstack/httpserver/adapter/undertow/UndertowHTTPSocket.java src/main/java/org/javawebstack/httpserver/adapter/untertow/UndertowHTTPSocket.java renamed to src/main/java/org/javawebstack/httpserver/adapter/undertow/UndertowHTTPSocket.java

@@ -1,6 +1,5 @@
-package org.javawebstack.httpserver.adapter.untertow;
+package org.javawebstack.httpserver.adapter.undertow;
 
-import io.undertow.server.BlockingHttpExchange;
 import io.undertow.server.HttpServerExchange;
 import io.undertow.util.HeaderValues;
 import io.undertow.util.HttpString;
@@ -17,24 +16,29 @@
 public class UndertowHTTPSocket implements IHTTPSocket {
 
     private final HttpServerExchange exchange;
-    private final BlockingHttpExchange blockingExchange;
+    private final InputStream inputStream;
+    private final OutputStream outputStream;
+    private boolean closed;
 
-    public UndertowHTTPSocket(HttpServerExchange exchange) {
+    public UndertowHTTPSocket(HttpServerExchange exchange, InputStream inputStream, OutputStream outputStream) {
         this.exchange = exchange;
-        this.blockingExchange = exchange.startBlocking();
+        this.inputStream = inputStream == null ? exchange.getInputStream() : inputStream;
+        this.outputStream = outputStream == null ? exchange.getOutputStream() : outputStream;
     }
 
     public InputStream getInputStream() throws IOException {
-        return exchange.getInputStream();
+        return inputStream;
     }
 
     public OutputStream getOutputStream() throws IOException {
-        return exchange.getOutputStream();
+        return outputStream;
     }
 
     public void close() throws IOException {
-        blockingExchange.close();
-        exchange.endExchange();
+        if(closed)
+            return;
+        closed = true;
+        exchange.getOutputStream().close();
     }
 
     public boolean isClosed() {
@@ -102,7 +106,7 @@ public String getResponseStatusMessage() {
     }
 
     public void writeHeaders() throws IOException {
-        exchange.getResponseSender().send("");
+        exchange.getOutputStream().write(new byte[0]);
     }
 
     public String getRemoteAddress() {

src/main/java/org/javawebstack/httpserver/adapter/undertow/UndertowHTTPSocketServer.java

@@ -0,0 +1,75 @@
+package org.javawebstack.httpserver.adapter.undertow;
+
+import io.undertow.Undertow;
+import io.undertow.server.handlers.BlockingHandler;
+import io.undertow.websockets.core.WebSocketVersion;
+import org.javawebstack.httpserver.adapter.IHTTPSocketHandler;
+import org.javawebstack.httpserver.adapter.IHTTPSocketServer;
+import org.javawebstack.httpserver.util.websocket.WebSocketUtil;
+import org.xnio.Options;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+public class UndertowHTTPSocketServer implements IHTTPSocketServer {
+
+    private int port = 80;
+    private Undertow server;
+    private IHTTPSocketHandler handler;
+    private ExecutorService executorService;
+
+    public void setPort(int port) {
+        this.port = port;
+    }
+
+    public int getPort() {
+        return port;
+    }
+
+    public void start() throws IOException {
+        executorService = Executors.newCachedThreadPool();
+        server = Undertow.builder()
+                .addHttpListener(port, "0.0.0.0")
+                .setServerOption(Options.KEEP_ALIVE, true)
+                .setHandler(new BlockingHandler(httpServerExchange -> {
+                    if(httpServerExchange.getRequestHeaders().contains("sec-websocket-key")) {
+                        httpServerExchange.upgradeChannel((streamConnection, httpServerExchange1) -> {
+                            InputStream inputStream = new StreamSourceInputStream(streamConnection.getSourceChannel());
+                            OutputStream outputStream = new StreamSinkOutputStream(streamConnection.getSinkChannel());
+                            handler.handle(new UndertowHTTPSocket(httpServerExchange1, inputStream, outputStream));
+                        });
+                        httpServerExchange.putAttachment(WebSocketVersion.ATTACHMENT_KEY, WebSocketVersion.V13);
+                        if(!WebSocketUtil.accept(new UndertowHTTPSocket(httpServerExchange, null, null), null))
+                            return;
+                        httpServerExchange.endExchange();
+                    } else {
+                        handler.handle(new UndertowHTTPSocket(httpServerExchange, null, null));
+                    }
+                }))
+                .build();
+        server.start();
+    }
+
+    public void stop() {
+        executorService.shutdown();
+        server.stop();
+    }
+
+    public void join() {
+        try {
+            server.getWorker().awaitTermination();
+        } catch (InterruptedException e) {}
+    }
+
+    public void setHandler(IHTTPSocketHandler handler) {
+        this.handler = handler;
+    }
+
+    public boolean isWebSocketSupported() {
+        return true;
+    }
+
+}

src/main/java/org/javawebstack/httpserver/adapter/untertow/UndertowHTTPSocketServer.java

@@ -98,11 +98,11 @@ public void write(OutputStream stream) throws IOException {
 
     public static WebSocketFrame read(InputStream stream) throws IOException {
         WebSocketFrame frame = new WebSocketFrame();
-        int b = safeRead(stream);
+        byte b = safeRead(stream);
         frame.flags = (byte) (b & 0xF0);
         frame.opcode = (byte) (b & 0x0F);
         b = safeRead(stream);
-        frame.maskKey = (b >> 7) == 1 ? new byte[4] : null;
+        frame.maskKey = ((b & 0xFF) >> 7) == 1 ? new byte[4] : null;
         int len = b & 0b0111_1111;
         if(len == 126) {
             len = safeRead(stream) << 8;
@@ -114,27 +114,27 @@ public static WebSocketFrame read(InputStream stream) throws IOException {
             len |= safeRead(stream);
         }
         if(frame.maskKey != null) {
-            frame.maskKey[0] = (byte) safeRead(stream);
-            frame.maskKey[1] = (byte) safeRead(stream);
-            frame.maskKey[2] = (byte) safeRead(stream);
-            frame.maskKey[3] = (byte) safeRead(stream);
+            frame.maskKey[0] = safeRead(stream);
+            frame.maskKey[1] = safeRead(stream);
+            frame.maskKey[2] = safeRead(stream);
+            frame.maskKey[3] = safeRead(stream);
         }
         frame.payload = new byte[len];
         if(frame.maskKey != null) {
             for(int i=0; i<len; i++)
                 frame.payload[i] = (byte) (safeRead(stream) ^ frame.maskKey[i % 4]);
         } else {
             for(int i=0; i<len; i++)
-                frame.payload[i] = (byte) safeRead(stream);
+                frame.payload[i] = safeRead(stream);
         }
         return frame;
     }
 
-    private static int safeRead(InputStream stream) throws IOException {
+    private static byte safeRead(InputStream stream) throws IOException {
         int b = stream.read();
         if(b == -1)
             throw new IOException("Unexpected end of stream");
-        return b;
+        return (byte) b;
     }
 
 }

src/main/java/org/javawebstack/httpserver/util/websocket/WebSocketFrame.java

@@ -21,8 +21,6 @@ public InternalWebSocketRequestHandler(WebSocketHandler handler) {
     public Object handle(Exchange exchange) {
         IHTTPSocket socket = exchange.socket();
         try {
-            if(!WebSocketUtil.accept(socket, null))
-                return null;
             WebSocket webSocket = new WebSocket(exchange);
             handler.onConnect(webSocket);
             WebSocketFrame frame;