Standalone Java Websocket Server

WebSocketStandalone.java

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.Set;


/**
 * A standalone Java implementation of a WebSocket server,
 * based on ServerSocketChannel and SocketChannel classes from the java.nio package,
 * for non-blocking communication with clients. It listens on port 8000 and waits
 * for incoming connections using a Selector. When a client connects, the server reads incoming
 * data using a ByteBuffer and performs a WebSocket protocol handshake.
 * If the incoming data is a GET request, the server generates and returns a Sec-WebSocket-Accept header in response.
 * If the incoming data is a WebSocket message, the server extracts the message data and performs
 * an XOR operation on the message data with a mask key if the message is masked.
 *
 * Client is located here https://gist.github.com/patrikbego/acd1eb069195c78d9db7edb8a0eb6a3c
 */
class WebSocketServer {
    private static final String GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"; // This is a global unique identifier used in the WebSocket protocol, as per https://www.rfc-editor.org/rfc/rfc6455

    public static void main(String[] args) throws IOException, NoSuchAlgorithmException {
        try (ServerSocketChannel serverSocketChannel = ServerSocketChannel.open()) {
            serverSocketChannel.configureBlocking(false); // set the channel to non-blocking mode, meaning that it will not block the main thread when waiting for a connection.
            serverSocketChannel.bind(new InetSocketAddress(8000));
            System.out.println("Server has started on 127.0.0.1:8000).\r\nWaiting for a connection…");

            Selector selector = Selector.open();
            serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);


            while (true) {
                selector.select(1); // blocks until at least one channel is ready for I/O operations, or until the timeout has been reached. In this case, the timeout is set to 1 millisecond.
                Set<SelectionKey> selectionKeys = selector.selectedKeys();
                for (SelectionKey key : selectionKeys) {
                    if (key.isAcceptable()) { // If the key is acceptable, a new connection is accepted by the channel and a new SocketChannel is created.
                        SocketChannel channel = serverSocketChannel.accept();
                        if (channel != null) {
                            channel.configureBlocking(false);
                            channel.register(selector, SelectionKey.OP_READ);
                        }
                    }
                    if (key.isReadable()) { // I=ff the key is readable, data can be read from the SocketChannel
                        SocketChannel socketChannel = (SocketChannel) key.channel();
                        ByteBuffer buffer = ByteBuffer.allocate(2048);
                        int count = socketChannel.read(buffer);
                        buffer.flip();

                        // We need Sec-WebSocket-Key header as specified in RFC 6455.
                        String string = new String(buffer.array()); // creates a string from the byte buffer, which should contain the data received from the SocketChannel.
                        if (string.startsWith("GET / HTTP/1.1")) {
                            int swkIndex = string.indexOf("Sec-WebSocket-Key:");
                            int endIndex = string.indexOf("\r\n", swkIndex);
                            String swk = string.substring(swkIndex + 19, endIndex);

                            String s = encodeSha1AndBase64(swk);

                            String response = "HTTP/1.1 101 Switching Protocols\r\n" +
                                    "Upgrade: websocket\r\n" +
                                    "Connection: Upgrade\r\n" +
                                    "Sec-WebSocket-Accept: " + s + "\r\n\r\n";
                            socketChannel.write(ByteBuffer.wrap(response.getBytes()));
                        } else {
                            // If the buffer is not empty, the code then reads the first two bytes from the buffer.
                            if (buffer.limit() <= 0)
                                continue;

                            byte firstByte = buffer.get(); // The first byte contains information about the message, such as whether it is the last message in a frame, and the type of message.
                            byte secondByte = buffer.get(); // The second byte contains the length of the message.

                            byte fin = (byte) ((firstByte & 128) >> 7); // The fin byte indicates whether this is the last message in a frame
                            byte rsv1 = (byte) ((firstByte & 64) >> 6); // The rsv bytes are reserved for future use
                            byte rsv2 = (byte) ((firstByte & 32) >> 5);
                            byte rsv3 = (byte) ((firstByte & 16) >> 4);

                            byte opCode = (byte) ((firstByte & 8) | (firstByte & 4) | (firstByte & 2) | (firstByte & 1)); // The opCode byte indicates the type of message.

                            byte mask = (byte) ((secondByte & 128) >> 7); // The mask byte indicates whether the message is masked
                            long length = (secondByte & 64) | (secondByte & 32) | (secondByte & 16) | (secondByte & 8) | (secondByte & 4) | (secondByte & 2) | (secondByte & 1);

                            if (length == 126) {
                                length = buffer.getShort();
                            } else if (length == 127) {
                                length = buffer.getLong();
                            }

                            byte[] dataBytes = new byte[(int) length];

                            byte[] maskValue = null;
                            if (mask == 1) {
                                maskValue = new byte[4]; //reads four bytes from the buffer to get the mask value
                                buffer.get(maskValue);

                                for (int i = 0; i < length; i++) {
                                    byte data = buffer.get();
                                    dataBytes[i] = (byte) (data ^ maskValue[i % 4]); // The data is unmasked using a bitwise XOR operation with maskValue[i % 4]. The result is stored in dataBytes[i].
                                }
                            }
                            System.out.println(new String(dataBytes));
                            if(new String(dataBytes).equalsIgnoreCase("hello")) {
                                byte[] textBytes = "world".getBytes(StandardCharsets.UTF_8);
                                ByteArrayOutputStream bao = new ByteArrayOutputStream();

                                // Add start of for a single-frame unmasked text message
                                bao.write(0x81); //  the byte 0x81 is written to the bao stream. This is the start of the WebSocket message and 0x1 means a text frame (0x2 would be binary data).
                                bao.write((byte) textBytes.length);
                                bao.write(textBytes); // the length of the response message is written to the bao stream.
                                bao.flush();
                                bao.close();
                                socketChannel.write(ByteBuffer.wrap(bao.toByteArray(), 0, bao.size()));
                            }
                        }
                    }
                }
            }
        }
    }

    private static String encodeSha1AndBase64(String secWebSocketKey) throws NoSuchAlgorithmException {
        String string = secWebSocketKey + GUID;
        MessageDigest digest = MessageDigest.getInstance("SHA-1");
        byte[] sha1Result = digest.digest(string.getBytes());
        return Base64.getEncoder().encodeToString(sha1Result);
    }
}
// ref client https://developer.mozilla.org/en-US/docs/Web/API/WebSockets_API/Writing_WebSocket_client_applications
// CLIENT: in browser open the console and paste in the code
// const exampleSocket = new WebSocket("ws://localhost:8000");
// exampleSocket.send("Hello ");
// Better client is in ./WebSocketClient.js