websocket连接样例

一个非常好用的测试websocket连接的网站

websocket连接前端部分

let main = function () {
    let ws = new WebSocket('ws://localhost:8080/ws/1001');

    ws.onopen = function (e) {
        console.log('连接打开');
        // obj为发送的内容，具体详细内容需要与后端交流
        let obj = {
            'toId': 1002,
            'msg': "1002 你最近好吗?"
        }
        // send方法向服务器发送消息
        ws.send(JSON.stringify(obj));
    }

    // 关闭连接时发送
    ws.onclose = function (e) {
        console.log('连接关闭', e.code, e.reason);
    }

    ws.onmessage = function (e) {
        // 当接收到服务器发来的消息触发该函数
        console.log(e);
        let data = e.data;
        console.log(data);
    }
}

websocket连接后端部分

后端部分采用java，数据库使用mongoDB，使用springboot-starter-websocket框架
首先是websocket连接的handler (MessageHandler.java)

package com.jarrett.charroom.websocket;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jarrett.charroom.dao.MessageDao;
import com.jarrett.charroom.pojo.Message;
import com.jarrett.charroom.pojo.UserData;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.concurrent.ConcurrentHashMap;

@Component
public class MessageHandler extends TextWebSocketHandler {

    // 存储所有用户的map映射集合
    private static final ConcurrentHashMap<Long, WebSocketSession> SESSION = new ConcurrentHashMap<>();

    private static final ObjectMapper MAPPER = new ObjectMapper();

    @Autowired
    MessageDao messageDao;

    // 建立websocket连接
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Long uid = (Long) session.getAttributes().get("uid");
        SESSION.put(uid, session);
        System.out.println(uid + "已经建立连接!");
    }

    // 退出websocket连接
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Long uid = (Long) session.getAttributes().get("uid");
        SESSION.remove(uid);
        System.out.println(uid + "退出了连接");
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage textMessage) throws Exception {
        Long uid = (Long) session.getAttributes().get("uid");
        System.out.println(textMessage.getPayload());
        JsonNode jsonNode = MAPPER.readTree(textMessage.getPayload());
        long toId = jsonNode.get("toId").asLong();
        String msg = jsonNode.get("msg").asText();

        Message message = Message.builder().from(UserData.USER_MAP.get(uid)).to(UserData.USER_MAP.get(toId))
                .msg(msg).build();
        message = messageDao.saveMessage(message);
        WebSocketSession toSession = SESSION.get(toId);
        if (toSession != null && toSession.isOpen()) {
            toSession.sendMessage(new TextMessage(MAPPER.writeValueAsString(message)));
            messageDao.updateMessageState(message.getId(), 1);
        }
    }
}

然后是handler的拦截器 (MessageHandlerInterceptor.java)

package com.jarrett.charroom.websocket;

import org.apache.commons.lang3.StringUtils;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

@Component
public class MessageHandlerInterceptor implements HandshakeInterceptor {
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        // URL : /ws/uid
        String path = request.getURI().getPath();
        String[] split = StringUtils.split(path, '/');

        if (split == null || split.length != 2) return false;
        if (!StringUtils.isNumeric(split[1])) {
            return false;
        }
        attributes.put("uid", Long.valueOf(split[1]));
        System.out.println(split[1] + "申请连接");
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {

    }
}

最后是将handler和handler拦截器注册 (WebSocketConfig.java)

package com.jarrett.charroom.websocket;

import com.jarrett.charroom.pojo.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Autowired
    MessageHandler messageHandler;

    @Autowired
    MessageHandlerInterceptor messageHandlerInterceptor;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(this.messageHandler, "/ws/{uid}")
                .setAllowedOrigins("*")
                .addInterceptors(this.messageHandlerInterceptor);
    }
}