io.undertow.websockets.core.WebSocketChannel Java Examples

@Override
protected void onFullTextMessage(WebSocketChannel channel, BufferedTextMessage textMessage) {
    var wrapper = (ChannelImpl) channel.getAttribute(WebSocketHandler.CHANNEL_KEY);
    ActionLog actionLog = logManager.begin("=== ws message handling begin ===");
    try {
        actionLog.action(wrapper.action);
        linkContext(channel, wrapper, actionLog);

        String data = textMessage.getData();
        logger.debug("[channel] message={}", data);     // not mask, assume ws message not containing sensitive info, the data can be json or plain text
        actionLog.track("ws", 0, 1, 0);

        validateRate(wrapper);

        Object message = wrapper.handler.fromClientMessage(data);
        wrapper.handler.listener.onMessage(wrapper, message);
    } catch (Throwable e) {
        logManager.logError(e);
        WebSockets.sendClose(closeCode(e), e.getMessage(), channel, ChannelCallback.INSTANCE);
    } finally {
        logManager.end("=== ws message handling end ===");
    }
}
 

@Override
public void onConnect(WebSocketHttpExchange exchange, WebSocketChannel channel) {
    if (this.hasAuthentication) {
        String header = null;
        if (exchange.getRequestHeader(Header.COOKIE.toString()) != null) {
            header = exchange.getRequestHeader(Header.COOKIE.toString());
        }
        
        if (RequestUtils.hasValidAuthentication(header)) {
            channel.getReceiveSetter().set((ChannelListener<? super WebSocketChannel>) Application.getInstance(this.controllerClass));
            channel.resumeReceives();
            channel.addCloseTask(Application.getInstance(WebSocketCloseListener.class));
            Application.getInstance(WebSocketService.class).addChannel(channel);
        } else {
            MangooUtils.closeQuietly(channel);
        }
    } else {
        channel.getReceiveSetter().set((ChannelListener<? super WebSocketChannel>) Application.getInstance(this.controllerClass));
        channel.resumeReceives();
        channel.addCloseTask(Application.getInstance(WebSocketCloseListener.class));
        Application.getInstance(WebSocketService.class).addChannel(channel);
    }
}
 

private void registerWebSocketChannelListener(WebSocketChannel webSocketChannel) {
  ChannelListener<WebSocketChannel> listener = new AbstractReceiveListener() {

    @Override
    protected void onFullBinaryMessage(WebSocketChannel channel, BufferedBinaryMessage message) throws IOException {
      log.trace("Server received full binary message");
      Pooled<ByteBuffer[]> pulledData = message.getData();
      try {
        ByteBuffer[] resource = pulledData.getResource();
        ByteBuffer byteBuffer = WebSockets.mergeBuffers(resource);
        String msg = new String(byteBuffer.array());
        log.trace("Sending message to decoder: {}", msg);
        writeToDecoder(msg);
      } finally {
        pulledData.discard();
      }
    }
  };
  webSocketChannel.getReceiveSetter().set(listener);
}
 

private void registerWebSocketChannelListener(WebSocketChannel webSocketChannel) {
  ChannelListener<WebSocketChannel> listener = new AbstractReceiveListener() {

    @Override
    protected void onFullBinaryMessage(WebSocketChannel channel, BufferedBinaryMessage message) throws IOException {
      log.trace("Server received full binary message");
      Pooled<ByteBuffer[]> pulledData = message.getData();
      try {
        ByteBuffer[] resource = pulledData.getResource();
        ByteBuffer byteBuffer = WebSockets.mergeBuffers(resource);
        String msg = new String(byteBuffer.array());
        log.trace("Sending message to decoder: {}", msg);
        writeToDecoder(msg);
      } finally {
        pulledData.discard();
      }
    }
  };
  webSocketChannel.getReceiveSetter().set(listener);
}
 

@Override
public void onConnect(WebSocketHttpExchange exchange, WebSocketChannel channel) {
    try {
        Cookie sid = getRequestCookies(exchange).get("PYX-Session");

        User user;
        if (sid == null || (user = Sessions.get().getUser(sid.getValue())) == null) {
            sendConnectionError(exchange, channel, new JsonWrapper(Consts.ErrorCode.NOT_REGISTERED));
        } else if (!user.isValid()) {
            sendConnectionError(exchange, channel, new JsonWrapper(Consts.ErrorCode.SESSION_EXPIRED));
        } else {
            if (user.getEventsSender() == null) user.establishedEventsConnection(new EventsSender(user, channel));
            else user.getEventsSender().addChannel(channel);

            channel.getCloseSetter().set((ChannelListener<AbstractFramedChannel>) newChannel -> {
                if (user.getEventsSender() != null)
                    user.getEventsSender().removeChannel((WebSocketChannel) newChannel);
            });
        }
    } catch (Throwable ex) {
        logger.error("Failed handling incoming connection.", ex);
        throw ex;
    }
}
 

@Override
protected void onCloseMessage(CloseMessage message, WebSocketChannel channel) {
    var wrapper = (ChannelImpl) channel.getAttribute(WebSocketHandler.CHANNEL_KEY);
    ActionLog actionLog = logManager.begin("=== ws close message handling begin ===");
    try {
        actionLog.action(wrapper.action + ":close");
        linkContext(channel, wrapper, actionLog);

        int code = message.getCode();
        String reason = message.getReason();
        actionLog.context("code", code);
        logger.debug("[channel] reason={}", reason);
        actionLog.track("ws", 0, 1, 0);

        wrapper.handler.listener.onClose(wrapper, code, reason);
    } catch (Throwable e) {
        logManager.logError(e);
    } finally {
        logManager.end("=== ws close message handling end ===");
    }
}
 

@Override
public void onConnect(WebSocketHttpExchange exchange, WebSocketChannel channel) {

    String uri = exchange.getRequestURI();
    final String subscriptionId = uri.substring(path.length() + 1);
    EventReservationsHandler.Reservation reservation = eventReservationsHandler.claimReservation(subscriptionId);
    if (reservation == null) {
        send(channel, "error", "Invalid subscription: not reserved");
        safeClose(channel);
        return;
    }
    LOG.debug("Principal is: {}", reservation.getPrincipal());
    send(channel, "message", "connected");
    bus.subscribe(subscriptionId, (type, data) -> {
        if (channel.isOpen()) {
            send(channel, type, data);
        } else {
            bus.unsubscribe(subscriptionId);
        }
    });
}
 

private PooledByteBuffer decompress(WebSocketChannel channel, PooledByteBuffer pooled) throws IOException {
    ByteBuffer buffer = pooled.getBuffer();

    while (!decompress.needsInput() && !decompress.finished()) {
        if (!buffer.hasRemaining()) {
            pooled = largerBuffer(pooled, channel, buffer.capacity() * 2);
            buffer = pooled.getBuffer();
        }

        int n;

        try {
            n = decompress.inflate(buffer.array(),
                    buffer.arrayOffset() + buffer.position(),
                    buffer.remaining());
        } catch (DataFormatException e) {
            WebSocketLogger.EXTENSION_LOGGER.debug(e.getMessage(), e);
            throw WebSocketMessages.MESSAGES.badCompressedPayload(e);
        }
        buffer.position(buffer.position() + n);
    }

    return pooled;
}
 

private void processWebSocketRequest(Request request, Response response) {
        HttpServletRequest servletRequest = request.getHttpServletRequest();
        HttpServletResponse servletResponse = response.getHttpServletResponse();

        WebSocketHttpExchange exchange = new ServletWebSocketHttpExchange(servletRequest, servletResponse, peerConnections);
//        exchange.putAttachment(HandshakeUtil.PATH_PARAMS, Collections.<String, String>emptyMap());

        Handshake handshake = getHandshake(exchange);

        exchange.upgradeChannel((connection, serverExchange) -> {
            WebSocketChannel channel = handshake.createChannel(exchange, connection, exchange.getBufferPool());
            peerConnections.add(channel);
            createWebSocketAdapter(request).onConnect(exchange, channel);
        });

        handshake.handshake(exchange);
    }
 

/**
 * Adds a new channel to the manager
 *
 * @param channel channel The channel to put
 */
public void addChannel(WebSocketChannel channel) {
    Objects.requireNonNull(channel, Required.CHANNEL.toString());

    final String url = RequestUtils.getWebSocketURL(channel);
    Set<WebSocketChannel> channels = getChannels(url);
    if (channels == null) {
        channels = new HashSet<>();
        channels.add(channel);
    } else {
        channels.add(channel);
    }
    setChannels(url, channels);
}
 

@Test
public void testRemoveChannel() {
    //given
    final WebSocketService webSocketService = Application.getInstance(WebSocketService.class);
    final WebSocketChannel channel = Mockito.mock(WebSocketChannel.class);
    when(channel.getUrl()).thenReturn("/websocket");

    //when
    webSocketService.addChannel(channel);
    webSocketService.removeChannels("/websocket");

    //then
    assertThat(webSocketService.getChannels("/websocket"), not(nullValue()));
    assertThat(webSocketService.getChannels("/websocket").size(), equalTo(0));
}
 

@Override
public void onConnect(WebSocketHttpExchange exchange, WebSocketChannel channel) {
	UndertowWebSocketSession session = createSession(channel);
	UndertowWebSocketHandlerAdapter adapter = new UndertowWebSocketHandlerAdapter(session);

	channel.getReceiveSetter().set(adapter);
	channel.resumeReceives();

	this.handler.handle(session)
			.checkpoint(exchange.getRequestURI() + " [UndertowRequestUpgradeStrategy]")
			.subscribe(session);
}
 

public WebSocketTtyConnection(WebSocketChannel webSocketChannel, ScheduledExecutorService executor) {
    this.webSocketChannel = webSocketChannel;
    this.executor = executor;

    registerWebSocketChannelListener(webSocketChannel);
    webSocketChannel.resumeReceives();
}
 

public void handle(HttpServerExchange exchange, RequestImpl request, ActionLog actionLog) {
    String path = exchange.getRequestPath();
    String action = "ws:" + path;
    actionLog.action(action + ":open");

    ChannelHandler handler = handlers.get(path);
    if (handler == null) throw new NotFoundException("not found, path=" + path, "PATH_NOT_FOUND");

    request.session = loadSession(request, actionLog);  // load session as late as possible, so for sniffer/scan request with sessionId, it won't call redis every time even for 404/405

    var webSocketExchange = new AsyncWebSocketHttpServerExchange(exchange, channels);
    exchange.upgradeChannel((connection, httpServerExchange) -> {
        WebSocketChannel channel = handshake.createChannel(webSocketExchange, connection, webSocketExchange.getBufferPool());
        try {
            var wrapper = new ChannelImpl(channel, context, handler);
            wrapper.action = action;
            wrapper.clientIP = request.clientIP();
            wrapper.refId = actionLog.id;   // with ws, correlationId and refId are same as parent http action id
            actionLog.context("channel", wrapper.id);
            channel.setAttribute(CHANNEL_KEY, wrapper);
            channel.addCloseTask(channelCloseListener);
            context.add(wrapper);

            handler.listener.onConnect(request, wrapper);
            actionLog.context("room", wrapper.rooms.toArray()); // may join room onConnect
            channel.getReceiveSetter().set(messageListener);
            channel.resumeReceives();

            channels.add(channel);
        } catch (Throwable e) {
            // upgrade is handled by io.undertow.server.protocol.http.HttpReadListener.exchangeComplete, and it catches all exceptions during onConnect
            logManager.logError(e);
            IoUtils.safeClose(connection);
        }
    });
    handshake.handshake(webSocketExchange);
}
 

public UndertowRequestUpgradeStrategy() {
	if (exchangeConstructorWithPeerConnections) {
		this.peerConnections = Collections.newSetFromMap(new ConcurrentHashMap<WebSocketChannel, Boolean>());
	}
	else {
		this.peerConnections = null;
	}
}
 

@Override
public void init(FilterConfig filterConfig) throws ServletException {
    super.init(filterConfig);

    peerConnections = Collections.newSetFromMap(new ConcurrentHashMap<WebSocketChannel, Boolean>());

    handshakes = new HashSet<>();
    handshakes.add(new Hybi13Handshake());
    handshakes.add(new Hybi08Handshake());
    handshakes.add(new Hybi07Handshake());
}
 

@Override
public void onConnect(WebSocketHttpExchange httpExchange, WebSocketChannel channel) {
	UndertowWebSocketSession session = createSession(channel);
	UndertowWebSocketHandlerAdapter adapter = new UndertowWebSocketHandlerAdapter(session);

	channel.getReceiveSetter().set(adapter);
	channel.resumeReceives();

	this.handler.handle(session).subscribe(session);
}
 

@Override
public void onConnect(WebSocketHttpExchange exchange, WebSocketChannel channel) {
    connection = new UndertowWebSocketConnection(exchange, channel);
    connections.add(connection);

    context = new WebSocketContext(connectionsReadOnly, connection, pathParameters);

    handler.onOpen(context);

    channel.addCloseTask(ch -> connections.remove(connection));
    channel.getReceiveSetter().set(this);
    channel.resumeReceives();
}
 

@Override
protected void write(byte[] buffer) {
  if (isOpen()) {
    sendBinary(buffer, webSocketChannel);
  }
  for (WebSocketChannel channel : readonlyChannels) {
    sendBinary(buffer, channel);
  }
}
 

/**
 * Sets the URI resources for a given URL
 *
 * @param uri The URI resource for the connection
 * @param channels The channels for the URI resource
 */
public void setChannels(String uri, Set<WebSocketChannel> channels) {
    Objects.requireNonNull(uri, Required.URI.toString());
    Objects.requireNonNull(channels, Required.URI_CONNECTIONS.toString());

    this.cache.put(Default.WSS_CACHE_PREFIX.toString() + uri, channels);
}
 

@Test
public void testAddChannel() {
    //given
    final WebSocketService webSocketService = Application.getInstance(WebSocketService.class);
    final WebSocketChannel channel = Mockito.mock(WebSocketChannel.class);
    when(channel.getUrl()).thenReturn("/websocket");

    //when
    webSocketService.removeChannels("/websocket");
    webSocketService.addChannel(channel);

    //then
    assertThat(webSocketService.getChannels("/websocket"), not(nullValue()));
    assertThat(webSocketService.getChannels("/websocket").size(), equalTo(1));
}
 

public WebSocketTtyConnection(WebSocketChannel webSocketChannel, ScheduledExecutorService executor) {
  this.webSocketChannel = webSocketChannel;
  this.executor = executor;

  registerWebSocketChannelListener(webSocketChannel);
  webSocketChannel.resumeReceives();
}
 

private void linkContext(WebSocketChannel channel, ChannelImpl wrapper, ActionLog actionLog) {
    actionLog.context("channel", wrapper.id);
    logger.debug("refId={}", wrapper.refId);
    List<String> refIds = List.of(wrapper.refId);
    actionLog.refIds = refIds;
    actionLog.correlationIds = refIds;
    logger.debug("[channel] url={}", channel.getUrl());
    logger.debug("[channel] remoteAddress={}", channel.getSourceAddress().getAddress().getHostAddress());
    actionLog.context("client_ip", wrapper.clientIP);
    actionLog.context("listener", wrapper.handler.listener.getClass().getCanonicalName());
    actionLog.context("room", wrapper.rooms.toArray());
}
 

private void handleReceivedMessage(WebSocketChannel channel, Object message, HttpServerExchange exchange) {
    executor.execute(new Runnable() {

        @Override
        public void run() {
            try {
                HttpServletRequest request = new WebSocketUndertowServletRequest(channel, message, exchange);
                HttpServletResponse response = new WebSocketUndertowServletResponse(channel);
                if (request.getHeader(WebSocketConstants.DEFAULT_REQUEST_ID_KEY) != null) {
                    String headerValue = request.getHeader(WebSocketConstants.DEFAULT_REQUEST_ID_KEY);
                    if (WebSocketUtils.isContainingCRLF(headerValue)) {
                        LOG.warning("Invalid characters (CR/LF) in header "
                            + WebSocketConstants.DEFAULT_REQUEST_ID_KEY);
                    } else {
                        response.setHeader(WebSocketConstants.DEFAULT_RESPONSE_ID_KEY, headerValue);
                    }
                }
                handleNormalRequest(request, response);
            } catch (Exception ex) {
                LOG.log(Level.WARNING, "Failed to invoke service", ex);
            }

        }

    });

}
 

private PooledByteBuffer largerBuffer(PooledByteBuffer smaller, WebSocketChannel channel, int newSize) {
    ByteBuffer smallerBuffer = smaller.getBuffer();

    smallerBuffer.flip();

    PooledByteBuffer larger = allocateBufferWithArray(channel, newSize);
    larger.getBuffer().put(smallerBuffer);

    smaller.close();
    return larger;
}
 

@Override
public void handleEvent(WebSocketChannel channel) {
    final String url = RequestUtils.getWebSocketURL(channel);
    final Set<WebSocketChannel> channels = this.cache.get(Default.WSS_CACHE_PREFIX.toString() + url);
    if (channels != null) {
        channels.remove(channel);
        this.cache.put(Default.WSS_CACHE_PREFIX.toString() + url, channels);
    }
}
 

private static AbstractReceiveListener getListener() {
    return new AbstractReceiveListener() {
        @Override
        protected void onFullTextMessage(WebSocketChannel channel, BufferedTextMessage message) {
            final String messageData = message.getData();
            for (WebSocketChannel session : channel.getPeerConnections()) {
                WebSockets.sendText(messageData, session, null);
            }
        }
    };
}
 

@Override
protected void upgradeInternal(ServerHttpRequest request, ServerHttpResponse response,
		String selectedProtocol, List<Extension> selectedExtensions, final Endpoint endpoint)
		throws HandshakeFailureException {

	HttpServletRequest servletRequest = getHttpServletRequest(request);
	HttpServletResponse servletResponse = getHttpServletResponse(response);

	final ServletWebSocketHttpExchange exchange = createHttpExchange(servletRequest, servletResponse);
	exchange.putAttachment(HandshakeUtil.PATH_PARAMS, Collections.<String, String>emptyMap());

	ServerWebSocketContainer wsContainer = (ServerWebSocketContainer) getContainer(servletRequest);
	final EndpointSessionHandler endpointSessionHandler = new EndpointSessionHandler(wsContainer);

	final ConfiguredServerEndpoint configuredServerEndpoint = createConfiguredServerEndpoint(
			selectedProtocol, selectedExtensions, endpoint, servletRequest);

	final Handshake handshake = getHandshakeToUse(exchange, configuredServerEndpoint);

	exchange.upgradeChannel(new HttpUpgradeListener() {
		@Override
		public void handleUpgrade(StreamConnection connection, HttpServerExchange serverExchange) {
			Object bufferPool = ReflectionUtils.invokeMethod(getBufferPoolMethod, exchange);
			WebSocketChannel channel = (WebSocketChannel) ReflectionUtils.invokeMethod(
					createChannelMethod, handshake, exchange, connection, bufferPool);
			if (peerConnections != null) {
				peerConnections.add(channel);
			}
			endpointSessionHandler.onConnect(exchange, channel);
		}
	});

	handshake.handshake(exchange);
}
 

private void removeChannel(WebSocketChannel channel) {
    synchronized (channels) {
        channels.remove(channel);
    }
}
 

private void handleIoException(IOException ex, WebSocketChannel channel) {
    if (ex instanceof ClosedChannelException) {
        removeChannel(channel);
    }
}