org.eclipse.jetty.websocket.api.Session Java Examples

/**
 *
 * Removes the closing client's session from 'sessionMap' (using the handshake identifier).
 *
 * @param session
 */
public void removeSession(final Session session) {
    logger.debug("removeSession called");

    final String sessionHandshake = session.getUpgradeResponse().getHeader(WS_HAND_SHAKE_KEY);

    sessionMap.entrySet().forEach(entry -> {
        Set<SessionIdWrapper> sessionSet = entry.getValue();
        sessionSet.forEach( s -> {
            if (s.getSession().getUpgradeResponse().getHeader(WS_HAND_SHAKE_KEY).equals(sessionHandshake)) {
                sessionSet.remove(s);

                liveLoggingHandler.broadcast(LiveLoggingUtils.buildLiveLogOutboundDTO(s.getTraceId(), null, null, "Websocket closed", false, false));

                return;
            }
        });
    });

}
 

@Override
public String onConnect(Session session) {
    for (HttpCookie cookie : session.getUpgradeRequest().getCookies()) {
        if ("auth-token".equals(cookie.getName())) {
            String authToken = cookie.getValue();
            TokenAuthenticator authenticator = getAuthenticator();
            org.hibernate.Session hSession = sessionFactory.openSession();
            ManagedSessionContext.bind(hSession);
            Optional<BasicToken> token;
            try {
                token = authenticator.authenticate(authToken);
            } catch (AuthenticationException e) {
                e.printStackTrace();
                return null;
            }
            if (!token.isPresent()) {
                return null;
            }
            hSession.close();
            return token.get().getUserId();
        }
    }
    return null;
}
 

public WebSocketVirtualServletRequest(WebSocketServletHolder websocket, InputStream in, Session session)
    throws IOException {
    this.webSocketHolder = websocket;
    this.in = in;

    Map<String, List<String>> ugHeaders = session.getUpgradeRequest().getHeaders();
    this.requestHeaders = WebSocketUtils.readHeaders(in);
    for (Map.Entry<String, List<String>> ent : ugHeaders.entrySet()) {
        if (!requestHeaders.containsKey(ent.getKey())) {
            requestHeaders.put(ent.getKey(), ent.getValue().get(0));
        }
    }
    String path = requestHeaders.get(WebSocketUtils.URI_KEY);
    String origin = websocket.getRequestURI();
    if (!path.startsWith(origin)) {
        LOG.log(Level.WARNING, "invalid path: {0} not within {1}", new Object[]{path, origin});
        throw new InvalidPathException();
    }
    this.attributes = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
    Object v = websocket.getAttribute("org.apache.cxf.transport.endpoint.address");
    if (v != null) {
        attributes.put("org.apache.cxf.transport.endpoint.address", v);
    }
}
 

private Session openNoteSession(String noteId, String principal, String ticket) {
  ClientUpgradeRequest request = new ClientUpgradeRequest();
  request.setHeader(ORIGIN, "*");
  ZeppelinWebsocket socket = new ZeppelinWebsocket(noteId);
  Future<Session> future = null;
  Session session = null;
  try {
    future = wsClient.connect(socket, zeppelinWebsocketUrl, request);
    session = future.get();
  } catch (IOException | InterruptedException | ExecutionException e) {
    LOG.error("Couldn't establish websocket connection to Zeppelin ", e);
    return session;
  }

  if (notesConnection.containsKey(noteId)) {
    session.close();
    session = notesConnection.get(noteId);
  } else {
    String getNote = serialize(zeppelinGetNoteMsg(noteId, principal, ticket));
    session.getRemote().sendStringByFuture(getNote);
    notesConnection.put(noteId, session);
  }
  return session;
}
 

@OnWebSocketConnect
public void onConnect(Session session) {
    LOG.info(webSocketClient.getSubProtocol() + " - Got connect.");
    if (mode == SDKConstants.INTERACTION_CLIENT_INPUT_MODE_INPUTSTREAM) {
        if (inputBuf == null) {
            this.inputBuf = ByteBuffer.allocate(SDKConstants.WEBSOCKET_MAX_BINARY_MESSAGE_SIZE);
            this.inputBuf.limit(0);
            this.inputStream = new ByteBufferBackedInputStream(inputBuf);
        }
    }
    if (outputBuf == null) {
        this.outputBuf = ByteBuffer.allocate(SDKConstants.WEBSOCKET_MAX_BINARY_MESSAGE_SIZE);
        this.outputStream = new ByteBufferBackedOutputStream(outputBuf, session);
    } else {
        synchronized (outputBuf) {
            ((ByteBufferBackedOutputStream) this.outputStream).setSession(session);
            outputBuf.notify();
        }
    }
    this.closed = false;
}
 

@Test
public void testSessionsManagement() {

	Assert.assertEquals( 0, WebSocketHandler.getSessions().size());

	RoboconfWebSocket ws1 = new RoboconfWebSocket();
	Session session1 = Mockito.mock( Session.class );
	ws1.onWebSocketConnect( session1 );
	Assert.assertEquals( 1, WebSocketHandler.getSessions().size());
	Assert.assertEquals( session1, WebSocketHandler.getSessions().iterator().next());

	RoboconfWebSocket ws2 = new RoboconfWebSocket();
	Session session2 = Mockito.mock( Session.class );
	ws2.onWebSocketConnect( session2 );
	Assert.assertEquals( 2, WebSocketHandler.getSessions().size());

	ws1.onWebSocketClose( 0, "whatever" );
	Assert.assertEquals( 1, WebSocketHandler.getSessions().size());
	Assert.assertEquals( session2, WebSocketHandler.getSessions().iterator().next());

	ws2.onWebSocketClose( 1, "whatever" );
	Assert.assertEquals( 0, WebSocketHandler.getSessions().size());
}
 

private void extractLoadBalanceKey(Session session) {
    String setCookieValue =
            session.getUpgradeResponse().getHeader(SDKConstants.CUPID_INTERACTION_HEADER_SET_COOKIE);
    if (setCookieValue != null) {
        setCookieValue = setCookieValue.trim();
        String[] kv = setCookieValue.split(";");
        for (String kvStr : kv) {
            if (kvStr.contains("=")) {
                String[] kAndV = kvStr.split("=");
                if (kAndV.length == 2 && kAndV[0] != null && kAndV[1] != null) {
                    if (SDKConstants.CUPID_INTERACTION_COOKIE_HASH_KEY.equals(kAndV[0].trim())) {
                        loadBalanceHashKey = kAndV[1].trim();
                        LOG.info(subProtocol + " - loadbalance key:" + loadBalanceHashKey);
                        return;
                    }
                }
            }
        }
    }
}
 

@OnWebSocketConnect @SuppressWarnings("unused")
public void onWebSocketConnect(final Session session)
{
    SocketAddress localAddress = session.getLocalAddress();
    SocketAddress remoteAddress = session.getRemoteAddress();
    _protocolEngine = _factory.newProtocolEngine(remoteAddress);

    // Let AMQP do timeout handling
    session.setIdleTimeout(0);

    _connectionWrapper = new ConnectionWrapper(session, localAddress, remoteAddress, _protocolEngine, _server.getThreadPool());
    if (session.getUpgradeRequest() instanceof ServletUpgradeRequest)
    {
        ServletUpgradeRequest upgradeRequest = (ServletUpgradeRequest) session.getUpgradeRequest();
        if (upgradeRequest.getCertificates() != null && upgradeRequest.getCertificates().length > 0)
        {
            _connectionWrapper.setPeerCertificate(upgradeRequest.getCertificates()[0]);
        }
    }
    _protocolEngine.setNetworkConnection(_connectionWrapper);
    _protocolEngine.setWorkListener(object -> _server.getThreadPool().execute(() -> _connectionWrapper.doWork()));
    _activeConnections.add(_connectionWrapper);
    _idleTimeoutChecker.wakeup();

}
 

/**
 * Fetches and submits all the data (e.g. {@link MediaObjectDescriptor}, {@link MediaSegmentDescriptor}) associated with the
 * raw results produced by a similarity search in a specific category.
 *
 * @param session The {@link Session} object used to transmit the results.
 * @param queryId ID of the running query.
 * @param category Name of the query category.
 * @param raw List of raw per-category results (segmentId -> score).
 */
private void finalizeAndSubmitResults(Session session, String queryId, String category, List<StringDoublePair> raw) {
    final int stride = 1000;
    for (int i=0; i<Math.floorDiv(raw.size(), stride)+1; i++) {
        final List<StringDoublePair> sub = raw.subList(i*stride, Math.min((i+1)*stride, raw.size()));
        final List<String> segmentIds = sub.stream().map(s -> s.key).collect(Collectors.toList());

        /* Load segment & object information. */
        final List<MediaSegmentDescriptor> segments = this.loadSegments(segmentIds);
        final List<String> objectIds = segments.stream().map(MediaSegmentDescriptor::getObjectId).collect(Collectors.toList());
        final List<MediaObjectDescriptor> objects = this.loadObjects(objectIds);
        if (segments.isEmpty() || objects.isEmpty()) {
            continue;
        }

        /* Write segments, objects and similarity search data to stream. */
        this.write(session, new MediaObjectQueryResult(queryId, objects));
        this.write(session, new MediaSegmentQueryResult(queryId, segments));
        this.write(session, new SimilarityQueryResult(queryId, category, sub));

        /* Load and transmit segment & object metadata. */
        this.loadAndWriteSegmentMetadata(session, queryId, segmentIds);
        this.loadAndWriteObjectMetadata(session, queryId, objectIds);
    }
}
 

@Test
public void testSocket_write_string() throws IOException {
    final String queryString = ParseQS.encode(new HashMap<String, String>() {{
        put("transport", WebSocket.NAME);
    }});
    final UpgradeRequest upgradeRequest = Mockito.mock(UpgradeRequest.class);
    final RemoteEndpoint remoteEndpoint = Mockito.mock(RemoteEndpoint.class);
    final Session session = Mockito.mock(Session.class);
    final EngineIoServer server = Mockito.spy(new EngineIoServer());

    Mockito.doAnswer(invocationOnMock -> queryString).when(upgradeRequest).getQueryString();
    Mockito.doAnswer(invocationOnMock -> upgradeRequest).when(session).getUpgradeRequest();
    Mockito.doAnswer(invocationOnMock -> remoteEndpoint).when(session).getRemote();
    Mockito.doAnswer(invocationOnMock -> null).when(server).handleWebSocket(Mockito.any(EngineIoWebSocket.class));

    final JettyWebSocketHandler handler = new JettyWebSocketHandler(server);
    handler.onWebSocketConnect(session);

    final String message = "FooBar";
    handler.write(message);
    Mockito.verify(remoteEndpoint, Mockito.times(1))
            .sendString(Mockito.eq(message));
}
 

@Override
public ListenableFuture<WebSocketSession> doHandshakeInternal(WebSocketHandler wsHandler,
		HttpHeaders headers, final URI uri, List<String> protocols,
		List<WebSocketExtension> extensions,  Map<String, Object> attributes) {

	final ClientUpgradeRequest request = new ClientUpgradeRequest();
	request.setSubProtocols(protocols);

	for (WebSocketExtension e : extensions) {
		request.addExtensions(new WebSocketToJettyExtensionConfigAdapter(e));
	}

	headers.forEach(request::setHeader);

	Principal user = getUser();
	final JettyWebSocketSession wsSession = new JettyWebSocketSession(attributes, user);
	final JettyWebSocketHandlerAdapter listener = new JettyWebSocketHandlerAdapter(wsHandler, wsSession);

	Callable<WebSocketSession> connectTask = () -> {
		Future<Session> future = this.client.connect(listener, uri, request);
		future.get();
		return wsSession;
	};

	if (this.taskExecutor != null) {
		return this.taskExecutor.submitListenable(connectTask);
	}
	else {
		ListenableFutureTask<WebSocketSession> task = new ListenableFutureTask<>(connectTask);
		task.run();
		return task;
	}
}
 

private synchronized double getSampleRate(Session session) {
  if (!this.sampleRate.containsKey(session)) {
    return 1.0D;
  } else {
    return Double.longBitsToDouble(this.sampleRate.get(session));
  }
}
 

@Test
public void testDepth5StreamWatcher() throws Exception, BinanceApiException {
    Session session = binanceApi.websocketDepth5(symbol, new BinanceWebSocketAdapterDepthLevel() {
        @Override
        public void onMessage(BinanceEventDepthLevelUpdate message) {
            log.info(message.toString());
        }
    });
    Thread.sleep(3000);
    session.close();
}
 

private void sendMessage(Session session, Object o) {
	MessageWrapper wrapper = new MessageWrapper(o.getClass().getSimpleName(), GSON.toJson(o));
	try {
		session.getRemote().sendString(GSON.toJson(wrapper));
	} catch (IOException e) {
		Logger.error(e, "Error: {}", e);
		// TODO Cleanup this session?
	}
}
 

@OnWebSocketConnect
public void onConnect(Session session) {
    // актуально только для LocalWSGameRunner
    if (WebSocketRunner.this.session == null) {
        WebSocketRunner.this.session = session;
    }

    print("Opened connection " + session.toString());
}
 

public void send(Message msg, String noteId) {
  Session noteSession = getZeppelinConnection(noteId, msg.principal, msg.ticket);
  if (!isSessionOpen(noteSession)) {
    LOG.error("Cannot open websocket connection to Zeppelin note {}", noteId);
    return;
  }
  noteSession.getRemote().sendStringByFuture(serialize(msg));
}
 

@Override
public Object createWebSocket(ServletUpgradeRequest servletUpgradeRequest, ServletUpgradeResponse servletUpgradeResponse) {
    final URI requestURI = servletUpgradeRequest.getRequestURI();
    final int port = servletUpgradeRequest.getLocalPort();
    final JettyWebSocketServer service = portToControllerService.get(port);

    if (service == null) {
        throw new RuntimeException("No controller service is bound with port: " + port);
    }

    final String path = requestURI.getPath();
    final WebSocketMessageRouter router;
    try {
        router = service.routers.getRouterOrFail(path);
    } catch (WebSocketConfigurationException e) {
        throw new IllegalStateException("Failed to get router due to: "  + e, e);
    }

    final RoutingWebSocketListener listener = new RoutingWebSocketListener(router) {
        @Override
        public void onWebSocketConnect(Session session) {
            final WebSocketPolicy currentPolicy = session.getPolicy();
            currentPolicy.setInputBufferSize(service.configuredPolicy.getInputBufferSize());
            currentPolicy.setMaxTextMessageSize(service.configuredPolicy.getMaxTextMessageSize());
            currentPolicy.setMaxBinaryMessageSize(service.configuredPolicy.getMaxBinaryMessageSize());
            super.onWebSocketConnect(session);
        }
    };

    return listener;
}
 

@OnWebSocketMessage
public void onMessage(Session session, String message) {
    OnlineUser onlineUser = onlineUserManager.getUser(session);

    //Shouldn't happen
    if (onlineUser == null) {
        LogManager.LOGGER.severe("(WS) FIXME: SocketServer:onMessage");
        return;
    }

    //Authenticated user
    if (onlineUser.isAuthenticated()) {
        messageDispatcher.dispatch(onlineUser, message);
        return;
    }

    //Handle auth request
    if (message.length() == AUTH_TOKEN_LEN) {
        LogManager.LOGGER.info("(WS) Received message from unauthenticated user " + session.getRemoteAddress().getAddress());

        User user = GameServer.INSTANCE.getUserManager().validateAuthToken(message);

        if (user != null) {
            doPostAuthUser(session, onlineUser, user);
        } else if (this.guestPolicy != GuestPolicy.BLOCK) {
            doPostAuthGuest(session, onlineUser);
        } else {
            LogManager.LOGGER.info("(WS) Blocked guest user " + session.getRemoteAddress().getAddress());
            kickOnlineUser(session);
        }
    }

    //Ignore other cases
}
 

@Before
public void setup() {
	this.session = mock(Session.class);
	given(this.session.getUpgradeRequest()).willReturn(Mockito.mock(UpgradeRequest.class));
	given(this.session.getUpgradeResponse()).willReturn(Mockito.mock(UpgradeResponse.class));

	this.webSocketHandler = mock(WebSocketHandler.class);
	this.webSocketSession = new JettyWebSocketSession(null, null);
	this.adapter = new JettyWebSocketHandlerAdapter(this.webSocketHandler, this.webSocketSession);
}
 

@OnWebSocketConnect
public void onWebSocketConnect(Session session) {
	try {
		this.wsSession.initializeNativeSession(session);
		this.webSocketHandler.afterConnectionEstablished(this.wsSession);
	}
	catch (Throwable ex) {
		ExceptionWebSocketHandlerDecorator.tryCloseWithError(this.wsSession, ex, logger);
	}
}
 

@Override
public boolean process(Session session, String message) throws IOException {
  String requestId = System.currentTimeMillis() + "." + counter.getAndIncrement();
  try (DataParser parser = parserFactory.getParser(requestId, message)) {
    return process(session, parser);
  } catch (DataParserException ex) {
    throw new IOException(ex);
  }
}
 

@Before
public void setup() {
	this.session = mock(Session.class);
	given(this.session.getUpgradeRequest()).willReturn(Mockito.mock(UpgradeRequest.class));
	given(this.session.getUpgradeResponse()).willReturn(Mockito.mock(UpgradeResponse.class));

	this.webSocketHandler = mock(WebSocketHandler.class);
	this.webSocketSession = new JettyWebSocketSession(null, null);
	this.adapter = new JettyWebSocketHandlerAdapter(this.webSocketHandler, this.webSocketSession);
}
 

/** AMQP frames MUST be sent as binary data payloads of WebSocket messages.*/
@OnWebSocketMessage @SuppressWarnings("unused")
public void onWebSocketText(Session sess, String text)
{
    LOGGER.info("Unexpected websocket text message received, closing connection");
    sess.close();
}
 

@OnWebSocketConnect
public void onConnect(Session session) {
  System.out.printf("Got connect: %s%n",session);
  try {
    Future<Void> fut = session.getRemote().sendStringByFuture("{\"field1\" : \"value\"}");
    fut.get(2, TimeUnit.SECONDS); // wait for send to complete.

    session.close(StatusCode.NORMAL,"I'm done");
  } catch (Throwable t) {
    t.printStackTrace();
  }
}
 

@Test(timeOut = 10000)
public void consumeMessagesInPartitionedTopicTest() throws Exception {
    final String namespace = "my-property/my-ns";
    final String topic = namespace + "/" + "my-topic7";
    admin.topics().createPartitionedTopic("persistent://" + topic, 3);

    final String subscription = "my-sub";
    final String consumerUri = "ws://localhost:" + proxyServer.getListenPortHTTP().get() + "/ws/v2/consumer/persistent/" + topic + "/" + subscription;
    final String producerUri = "ws://localhost:" + proxyServer.getListenPortHTTP().get() + "/ws/v2/producer/persistent/" + topic;

    URI consumeUri = URI.create(consumerUri);
    URI produceUri = URI.create(producerUri);

    WebSocketClient consumeClient = new WebSocketClient();
    WebSocketClient produceClient = new WebSocketClient();

    SimpleConsumerSocket consumeSocket = new SimpleConsumerSocket();
    SimpleProducerSocket produceSocket = new SimpleProducerSocket();

    try {
        produceClient.start();
        ClientUpgradeRequest produceRequest = new ClientUpgradeRequest();
        Future<Session> producerFuture = produceClient.connect(produceSocket, produceUri, produceRequest);
        producerFuture.get();
        produceSocket.sendMessage(100);
    } finally {
        stopWebSocketClient(produceClient);
    }

    Thread.sleep(500);

    try {
        consumeClient.start();
        ClientUpgradeRequest consumeRequest = new ClientUpgradeRequest();
        Future<Session> consumerFuture = consumeClient.connect(consumeSocket, consumeUri, consumeRequest);
        consumerFuture.get();
    } finally {
        stopWebSocketClient(consumeClient);
    }
}
 

@Test
public void testProcess() throws Exception {

	HttpRoutingContext routingContext = new HttpRoutingContext();
	HttpDmClient httpDmClient = new HttpDmClient( routingContext );

	// Not connected => no processing
	Message message = Mockito.mock( Message.class );
	Session session = Mockito.mock( Session.class );
	Mockito.when( session.isOpen()).thenReturn( false );

	httpDmClient.process( session, message );
	Mockito.verifyZeroInteractions( message );
	Mockito.verify( session, Mockito.only()).isOpen();

	// Connected => processing...
	Mockito.reset( session );
	Future<Void> future = Mockito.mock( Future.class );
	RemoteEndpoint remote = Mockito.mock( RemoteEndpoint.class );
	Mockito.when( remote.sendBytesByFuture( Mockito.any( ByteBuffer.class ) )).thenReturn( future );

	Mockito.when( session.getRemote()).thenReturn( remote );
	Mockito.when( session.isOpen()).thenReturn( true );

	httpDmClient.process( session, message );

	Mockito.verifyZeroInteractions( message );
	Mockito.verify( session, Mockito.times( 1 )).isOpen();
	Mockito.verify( session, Mockito.times( 1 )).getRemote();
	Mockito.verifyNoMoreInteractions( session );
}
 

private synchronized void unsubscribe(Session session, List<Metadata> metadatas) {   
  
  if (metadatas.isEmpty()) {
    return;
  }
  
  // 128BITS
  byte[] bytes = new byte[16];
  ByteBuffer bb = ByteBuffer.wrap(bytes).order(ByteOrder.BIG_ENDIAN);
  
  if (!this.subscriptions.containsKey(session)) {
    return;
  }
  
  for (Metadata metadata: metadatas) {
    bb.rewind();
    bb.putLong(metadata.getClassId());
    bb.putLong(metadata.getLabelsId());
    
    BigInteger id = new BigInteger(bytes);

    if (subscriptions.get(session).contains(id)) {
      subscriptions.get(session).remove(id);
      if (0 == this.refcounts.get(id).addAndGet(-1)) {
        this.metadatas.remove(id);
        this.refcounts.remove(id);
      }
    }
  }
  
  if (null != this.subscriptionListener) {
    this.subscriptionListener.onChange();
  }
}
 

private OUTPUT_FORMAT getOutputFormat(Session session) {
  if (this.outputFormat.containsKey(session)) {
    return this.outputFormat.get(session);
  } else {
    return OUTPUT_FORMAT.TEXT;
  }
}
 

/**
 * Send text frame to each WebSocket session.
 *
 * @param text text to send.
 */
public void broadcast(final String text) {
  for (final Session session : sessions) {
    try {
      session.getRemote().sendString(text);
    } catch (final IOException e) {
      LOG.warn("Failed to send string to remote session {}.", session.getRemoteAddress().toString());
    }
  }
}
 

/**
 * Send binary frame to each WebSocket session.
 *
 * @param bytes byte array to send.
 */
public void broadcast(final byte[] bytes) {
  for (final Session session : sessions) {
    try {
      session.getRemote().sendBytes(ByteBuffer.wrap(bytes));
    } catch (final IOException e) {
      LOG.warn("Failed to send binary to remote session {}.", session.getRemoteAddress().toString());
    }
  }
}