Java Code Examples for org.springframework.web.socket.WebSocketSession#getAttributes()

@Override
public void afterConnectionEstablished(WebSocketSession session) throws Exception {
    Map<String, Object> attributes = session.getAttributes();
    NodeModel nodeModel = (NodeModel) attributes.get("nodeInfo");
    UserModel userInfo = (UserModel) attributes.get("userInfo");
    String dataValue = (String) attributes.get(dataParName);
    String userName = UserModel.getOptUserName(userInfo);
    userName = URLUtil.encode(userName);
    if (nodeModel != null) {
        String url = NodeForward.getSocketUrl(nodeModel, nodeUrl);
        url = StrUtil.format(url, dataValue, userName);
        // 连接节点
        ProxySession proxySession = new ProxySession(url, session);
        session.getAttributes().put("proxySession", proxySession);
    }
    session.sendMessage(new TextMessage(StrUtil.format("欢迎加入:{} 会话id:{} ", userInfo.getName(), session.getId())));
}
 

@Override
protected void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
    if (operateLogController == null) {
        operateLogController = SpringUtil.getBean(OperateLogController.class);
    }
    String msg = message.getPayload();
    Map<String, Object> attributes = session.getAttributes();
    ProxySession proxySession = (ProxySession) attributes.get("proxySession");
    JSONObject json = JSONObject.parseObject(msg);
    String op = json.getString("op");
    ConsoleCommandOp consoleCommandOp = ConsoleCommandOp.valueOf(op);
    if (proxySession != null) {
        this.handleTextMessage(attributes, proxySession, json, consoleCommandOp);
    } else {
        this.handleTextMessage(attributes, session, json, consoleCommandOp);
    }
}
 

/**
 * 封装每个session中需要管理的参数
 * 例如：websockets做分页时，需要pageNum.和pageSize
 *
 * @param webSessionPayload
 * @return
 */
public static WebSocketSession buildWebSocketSession(WebSocketSession webSocketSession, WebSessionPayload webSessionPayload) {
    acquireUserIdAndSetSession(webSocketSession, webSessionPayload);
    Map<String, Object> attributes = webSocketSession.getAttributes();
    attributes.put(WARNING_PAGE_NUM, webSessionPayload.getWarningPage().getPageNum());
    attributes.put(WARNING_PAGE_SIZE, webSessionPayload.getWarningPage().getPageSize());
    attributes.put(MESSAGE_PAGE_NUM, webSessionPayload.getMessagePage().getPageNum());
    attributes.put(MESSAGE_PAGE_SIZE, webSessionPayload.getMessagePage().getPageSize());
    attributes.put(MESSAGE_TYPE, webSessionPayload.getMessageTypes());
    return webSocketSession;
}
 

public static <T> Optional<T> attributesFetch(WebSocketSession session, String key) {
    Map<String, Object> attributes = session.getAttributes();
    Object obj = attributes.get(key);
    if (obj == null) {
        return Optional.empty();
    }
    T t = (T) obj;
    return Optional.of(t);
}
 

/**
 * session关闭后清除内存中的相关session
 *
 * @param session
 * @param status
 */
@Override
public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
    Map<String, Object> attributes = session.getAttributes();
    Object id = attributes.get(USER_ID);
    if (nonNull(id)) {
        Integer userId = (int) id;
        GlobalSession.removeSession(userId, session);
    }
}
 

@Override
public void destroy(WebSocketSession session) {
    try {
        if (session.isOpen()) {
            session.close();
        }
    } catch (IOException ignored) {
    }
    Map<String, Object> attributes = session.getAttributes();
    ProxySession proxySession = (ProxySession) attributes.get("proxySession");
    if (proxySession != null) {
        proxySession.close();
    }
}
 

private PluginWebsocketSessionRef toRef(WebSocketSession session) throws IOException {
    URI sessionUri = session.getUri();
    String path = sessionUri.getPath();
    path = path.substring(WebSocketConfiguration.WS_PLUGIN_PREFIX.length());
    if (path.length() == 0) {
        throw new IllegalArgumentException("URL should contain plugin token!");
    }
    String[] pathElements = path.split("/");
    String pluginToken = pathElements[0];
    // TODO: cache
    PluginMetaData pluginMd = pluginService.findPluginByApiToken(pluginToken);
    if (pluginMd == null) {
        throw new InvalidParameterException("Can't find plugin with specified token!");
    } else {
        SecurityUser currentUser = (SecurityUser) session.getAttributes().get(WebSocketConfiguration.WS_SECURITY_USER_ATTRIBUTE);
        TenantId tenantId = currentUser.getTenantId();
        CustomerId customerId = currentUser.getCustomerId();
        if (PluginApiController.validatePluginAccess(pluginMd, tenantId, customerId)) {
            PluginApiCallSecurityContext securityCtx = new PluginApiCallSecurityContext(pluginMd.getTenantId(), pluginMd.getId(), tenantId,
                    currentUser.getCustomerId());
            return new BasicPluginWebsocketSessionRef(UUID.randomUUID().toString(), securityCtx, session.getUri(), session.getAttributes(),
                    session.getLocalAddress(), session.getRemoteAddress());
        } else {
            throw new SecurityException("Current user is not allowed to use this plugin!");
        }
    }
}
 

@Override
public void afterConnectionEstablished(WebSocketSession session) throws Exception {
  sessionId = session.getId();

  Map<String, Object> attributes = session.getAttributes();
  version = (String) attributes.get(VERSION);
  roadName = (String) attributes.get(ROAD_NAME);
  streamName = (String) attributes.get(STREAM_NAME);
  DefaultOffset defaultOffset = (DefaultOffset) attributes.get(DEFAULT_OFFSET);
  @SuppressWarnings("unchecked")
  Set<Sensitivity> grants = (Set<Sensitivity>) attributes.get(GRANTS);
  Authentication authentication = (Authentication) session.getPrincipal();

  RoadConsumer consumer;
  metrics = metricsFactory.create(roadName, streamName);
  try {
    authorisation.checkAuthorisation(authentication, roadName, grants);
    consumer = consumerFactory.create(roadName, streamName, defaultOffset);
    MessageFunction messageFunction = messageFunctionFactory.create(roadName, grants);
    EventSender sender = bytes -> sendEvent(session, bytes);
    service = serviceFactory.create(version, consumer, messageFunction, sender, metrics);
  } catch (UnknownRoadException e) {
    metrics.markRoadNotFound();
    session.close();
    throw e;
  }

  Scheduler scheduler = Schedulers.newSingle(
      String.format("offramp[v%s,%s.%s:%s]", version, roadName, streamName, sessionId));
  metrics.incrementActiveConnections();

  Mono<Void> mono = Mono
      .fromRunnable(service)
      .subscribeOn(scheduler)
      .doOnError(t -> true, t -> {
        log.error("Road: {}, stream: {}, sessionId: {} - Error in OfframpService", roadName, streamName, sessionId, t);
      })
      .doOnSuccess(s -> {
        log.info("Road: {}, stream: {}, sessionId: {} - OfframpService Completed", roadName, streamName, sessionId);
      })
      .doOnTerminate(() -> {
        disposables.dispose();
        metrics.decrementActiveConnections();
      }).then();

  disposables.add(() -> close(service));
  disposables.add(() -> close(consumer));
  disposables.add(() -> close(() -> session.close(SERVER_ERROR)));
  disposables.add(scheduler);
  disposables.add(mono.subscribe());

  log.info("Road: {}, stream: {}, sessionId: {} - Connection established with defaultOffset: {}", roadName,
      streamName, sessionId, defaultOffset);
  metrics.markConnectionEstablished();
}
 

/**
 * Handle STOMP messages going back out to WebSocket clients.
 */
@Override
@SuppressWarnings("unchecked")
public void handleMessageToClient(WebSocketSession session, Message<?> message) {
	if (!(message.getPayload() instanceof byte[])) {
		if (logger.isErrorEnabled()) {
			logger.error("Expected byte[] payload. Ignoring " + message + ".");
		}
		return;
	}

	StompHeaderAccessor accessor = getStompHeaderAccessor(message);
	StompCommand command = accessor.getCommand();

	if (StompCommand.MESSAGE.equals(command)) {
		if (accessor.getSubscriptionId() == null && logger.isWarnEnabled()) {
			logger.warn("No STOMP \"subscription\" header in " + message);
		}
		String origDestination = accessor.getFirstNativeHeader(SimpMessageHeaderAccessor.ORIGINAL_DESTINATION);
		if (origDestination != null) {
			accessor = toMutableAccessor(accessor, message);
			accessor.removeNativeHeader(SimpMessageHeaderAccessor.ORIGINAL_DESTINATION);
			accessor.setDestination(origDestination);
		}
	}
	else if (StompCommand.CONNECTED.equals(command)) {
		this.stats.incrementConnectedCount();
		accessor = afterStompSessionConnected(message, accessor, session);
		if (this.eventPublisher != null) {
			try {
				SimpAttributes simpAttributes = new SimpAttributes(session.getId(), session.getAttributes());
				SimpAttributesContextHolder.setAttributes(simpAttributes);
				Principal user = getUser(session);
				publishEvent(this.eventPublisher, new SessionConnectedEvent(this, (Message<byte[]>) message, user));
			}
			finally {
				SimpAttributesContextHolder.resetAttributes();
			}
		}
	}

	byte[] payload = (byte[]) message.getPayload();
	if (StompCommand.ERROR.equals(command) && getErrorHandler() != null) {
		Message<byte[]> errorMessage = getErrorHandler().handleErrorMessageToClient((Message<byte[]>) message);
		if (errorMessage != null) {
			accessor = MessageHeaderAccessor.getAccessor(errorMessage, StompHeaderAccessor.class);
			Assert.state(accessor != null, "No StompHeaderAccessor");
			payload = errorMessage.getPayload();
		}
	}
	sendToClient(session, accessor, payload);
}
 

/**
 * Handle STOMP messages going back out to WebSocket clients.
 */
@Override
@SuppressWarnings("unchecked")
public void handleMessageToClient(WebSocketSession session, Message<?> message) {
	if (!(message.getPayload() instanceof byte[])) {
		if (logger.isErrorEnabled()) {
			logger.error("Expected byte[] payload. Ignoring " + message + ".");
		}
		return;
	}

	StompHeaderAccessor accessor = getStompHeaderAccessor(message);
	StompCommand command = accessor.getCommand();

	if (StompCommand.MESSAGE.equals(command)) {
		if (accessor.getSubscriptionId() == null && logger.isWarnEnabled()) {
			logger.warn("No STOMP \"subscription\" header in " + message);
		}
		String origDestination = accessor.getFirstNativeHeader(SimpMessageHeaderAccessor.ORIGINAL_DESTINATION);
		if (origDestination != null) {
			accessor = toMutableAccessor(accessor, message);
			accessor.removeNativeHeader(SimpMessageHeaderAccessor.ORIGINAL_DESTINATION);
			accessor.setDestination(origDestination);
		}
	}
	else if (StompCommand.CONNECTED.equals(command)) {
		this.stats.incrementConnectedCount();
		accessor = afterStompSessionConnected(message, accessor, session);
		if (this.eventPublisher != null) {
			try {
				SimpAttributes simpAttributes = new SimpAttributes(session.getId(), session.getAttributes());
				SimpAttributesContextHolder.setAttributes(simpAttributes);
				Principal user = getUser(session);
				publishEvent(this.eventPublisher, new SessionConnectedEvent(this, (Message<byte[]>) message, user));
			}
			finally {
				SimpAttributesContextHolder.resetAttributes();
			}
		}
	}

	byte[] payload = (byte[]) message.getPayload();
	if (StompCommand.ERROR.equals(command) && getErrorHandler() != null) {
		Message<byte[]> errorMessage = getErrorHandler().handleErrorMessageToClient((Message<byte[]>) message);
		if (errorMessage != null) {
			accessor = MessageHeaderAccessor.getAccessor(errorMessage, StompHeaderAccessor.class);
			Assert.state(accessor != null, "No StompHeaderAccessor");
			payload = errorMessage.getPayload();
		}
	}
	sendToClient(session, accessor, payload);
}
 

/**
 * Handle STOMP messages going back out to WebSocket clients.
 */
@Override
@SuppressWarnings("unchecked")
public void handleMessageToClient(WebSocketSession session, Message<?> message) {
	if (!(message.getPayload() instanceof byte[])) {
		logger.error("Expected byte[] payload. Ignoring " + message + ".");
		return;
	}

	StompHeaderAccessor stompAccessor = getStompHeaderAccessor(message);
	StompCommand command = stompAccessor.getCommand();

	if (StompCommand.MESSAGE.equals(command)) {
		if (stompAccessor.getSubscriptionId() == null) {
			logger.warn("No STOMP \"subscription\" header in " + message);
		}
		String origDestination = stompAccessor.getFirstNativeHeader(SimpMessageHeaderAccessor.ORIGINAL_DESTINATION);
		if (origDestination != null) {
			stompAccessor = toMutableAccessor(stompAccessor, message);
			stompAccessor.removeNativeHeader(SimpMessageHeaderAccessor.ORIGINAL_DESTINATION);
			stompAccessor.setDestination(origDestination);
		}
	}
	else if (StompCommand.CONNECTED.equals(command)) {
		this.stats.incrementConnectedCount();
		stompAccessor = afterStompSessionConnected(message, stompAccessor, session);
		if (this.eventPublisher != null && StompCommand.CONNECTED.equals(command)) {
			try {
				SimpAttributes simpAttributes = new SimpAttributes(session.getId(), session.getAttributes());
				SimpAttributesContextHolder.setAttributes(simpAttributes);
				Principal user = session.getPrincipal();
				publishEvent(new SessionConnectedEvent(this, (Message<byte[]>) message, user));
			}
			finally {
				SimpAttributesContextHolder.resetAttributes();
			}
		}
	}

	byte[] payload = (byte[]) message.getPayload();

	if (StompCommand.ERROR.equals(command) && getErrorHandler() != null) {
		Message<byte[]> errorMessage = getErrorHandler().handleErrorMessageToClient((Message<byte[]>) message);
		stompAccessor = MessageHeaderAccessor.getAccessor(errorMessage, StompHeaderAccessor.class);
		Assert.notNull(stompAccessor, "Expected STOMP headers");
		payload = errorMessage.getPayload();
	}

	sendToClient(session, stompAccessor, payload);
}
 

private String getHttpSessionId(WebSocketSession wsSession) {
	Map<String, Object> attributes = wsSession.getAttributes();
	return SessionRepositoryMessageInterceptor.getSessionId(attributes);
}
 

/**
 * session关闭后清除内存中的相关session
 *
 * @param session
 * @param status
 * @throws Exception
 */
@Override
public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
    Map<String, Object> attributes = session.getAttributes();
    DeploymentWebSessionPayload payload = (DeploymentWebSessionPayload) attributes.get(JOB_NAME);
    GlobalSession.removeSessionByJob(payload, session);
}