websocket 用戶綁定 角色推送

需求：後臺管理賬戶在接收到訂閱消息時，主動推送提示到前臺。訂閱消息可參考 spring-boot redis發佈訂閱

依賴：pom.xml

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>com.lsj</groupId>
    <artifactId>websocket</artifactId>
    <version>0.0.1-SNAPSHOT</version>
    <packaging>jar</packaging>

    <name>websocket</name>
    <description>Demo project for Spring Boot</description>

    <parent>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-parent</artifactId>
        <version>1.5.6.RELEASE</version>
        <relativePath/> <!-- lookup parent from repository -->
    </parent>

    <properties>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
        <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
        <java.version>1.8</java.version>
    </properties>

    <dependencies>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-thymeleaf</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-websocket</artifactId>
        </dependency>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-test</artifactId>
            <scope>test</scope>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <plugin>
                <groupId>org.springframework.boot</groupId>
                <artifactId>spring-boot-maven-plugin</artifactId>
            </plugin>
        </plugins>
    </build>


</project>

配置：config

package com.lsj.websocket.config;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

@Configuration
public class WebSocketConfig {
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }

}

服務：server

package com.lsj.websocket.server;


import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * websocket
 */
@Component
@ServerEndpoint(value = "/websocket/{userId}/{ip}/{isAdmin}")
public class WebSocketServer {
    // 日誌
    private static final Logger log = LoggerFactory.getLogger(WebSocketServer.class);
    // 在線人數
    private static int onlineCount = 0;
    // 在線人員註冊 一個用戶可以多瀏覽器登錄
    private static ConcurrentHashMap<String, Map<String,WebSocketServer>> webSocketSet = new ConcurrentHashMap<>();
    // 在線管理員註冊 一個用戶可以多瀏覽器登錄
    private static ConcurrentHashMap<String, Map<String,WebSocketServer>> adminWebSocketSet = new ConcurrentHashMap<>();
    //與某個客戶端的連接會話，需要通過它來給客戶端發送數據
    private Session session;
    // 用戶ID
    private String userId = "";

    /**
     * 打開連接回調方法
     * @param userId 用戶ID
     * @param ip IP
     * @param isAdmin 是否管理員
     * @param session 客戶端session
     */
    @OnOpen
    public void onOpen(@PathParam(value = "userId") String userId,@PathParam("ip") String ip,
                       @PathParam("isAdmin") Integer isAdmin,Session session) {

        this.session = session;
        this.userId = userId;

        if (isAdmin == 0){
            // 非管理員註冊
            Map<String, WebSocketServer> map = webSocketSet.get(userId);
            map = null == map ? new HashMap<String, WebSocketServer>() : map;
            map.put(ip,this);
            webSocketSet.put(userId,map);

        }else {
            // 管理員註冊
            Map<String,WebSocketServer> map = adminWebSocketSet.get(userId);
            map = null == map ? new HashMap<String, WebSocketServer>() : map;
            map.put(ip,this);
            adminWebSocketSet.put(userId,map);
        }
        //在線數加1
        addOnlineCount();
        log.info("有一連接關閉！當前在線人數爲" + getOnlineCount());
    }

    /**
     * 關閉連接回調方法
     * @param userId 用戶ID
     * @param ip IP
     * @param isAdmin 是否管理員
     */
    @OnClose
    public void onClose(@PathParam(value = "userId") String userId,@PathParam("ip") String ip,
                        @PathParam("isAdmin") Integer isAdmin) {
        if (isAdmin == 0){
            Map<String, WebSocketServer> map = webSocketSet.get(userId);
            if (null != map && null != map.get(ip)){
                map.remove(ip);
            }
        }else {
            Map<String, WebSocketServer> map = adminWebSocketSet.get(userId);
            if (null != map && null != map.get(ip)){
                map.remove(ip);
            }
        }
        subOnlineCount();           //在線數減1
        log.info("有一連接關閉！當前在線人數爲" + getOnlineCount());
    }

    /**
     * 收到客戶端消息後調用的方法
     *
     * @param message 客戶端發送過來的消息*/
    @OnMessage
    public void onMessage(String message, Session session) {
        log.info("來自客戶端的消息:" + message);
        //可以自己約定字符串內容，比如 內容|0 表示信息羣發，內容|X 表示信息發給id爲X的用戶
        try {
            sendtoAll(message);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 發生錯誤回調方法
     * @param session 連接會話
     * @param error 錯誤
     */
    @OnError
    public void onError(Session session, Throwable error) {
        log.error("發生錯誤");
        error.printStackTrace();
    }


    /**
     * 發送信息
     * @param message 消息
     * @throws IOException 異常
     */
    public void sendMessage(String message) throws IOException {
        this.session.getBasicRemote().sendText(message);
    }

    /**
     * 發送信息給指定ID用戶，如果用戶不在線則返回不在線信息給自己
     * @param message 消息
     * @param sendUserId 用戶ID
     * @throws IOException 異常
     */
    public void sendToUser(String message,String sendUserId) throws IOException {
        // 是否爲非管理員
        if (webSocketSet.get(sendUserId) != null) {
            Map<String, WebSocketServer> map = webSocketSet.get(sendUserId);
            for (WebSocketServer wss : map.values()){
                wss.sendMessage(message);
            }
        }
        // 發送給所有管理員
        for (Map<String,WebSocketServer> map : adminWebSocketSet.values()) {
            for (WebSocketServer wss : map.values()){
                wss.sendMessage(message);
            }
        }
    }

    /**
     * 發送信息給所有人
     * @param message 消息
     * @throws IOException 異常
     */
    public void sendtoAll(String message) throws IOException {
        // 發送給所有非管理員用戶
        for (Map<String,WebSocketServer> map : webSocketSet.values()) {
            for (WebSocketServer wss : map.values()){
                wss.sendMessage(message);
            }
        }
        // 發送給所有管理員用戶
        for (Map<String,WebSocketServer> map : adminWebSocketSet.values()) {
            for (WebSocketServer wss : map.values()){
                wss.sendMessage(message);
            }
        }
    }


    public static synchronized int getOnlineCount() {
        return onlineCount;
    }

    public static synchronized void addOnlineCount() {
        WebSocketServer.onlineCount++;
    }

    public static synchronized void subOnlineCount() {
        WebSocketServer.onlineCount--;
    }

}

頁面：

<!DOCTYPE HTML>
<html>
<head>
    <title>WebSocket</title>
</head>

<body>
Welcome<br/>
<input id="text" type="text" /><button onclick="send()">Send</button>    <button onclick="closeWebSocket()">Close</button>
<div id="message">
</div>
</body>

<script type="text/javascript">
    var websocket = null;

    //判斷當前瀏覽器是否支持WebSocket
    if('WebSocket' in window){
        websocket = new WebSocket("ws://localhost:8880/websocket/2/234/1");
    }
    else{
        alert('Not support websocket')
    }

    //連接發生錯誤的回調方法
    websocket.onerror = function(){
        setMessageInnerHTML("error");
    };

    //連接成功建立的回調方法
    websocket.onopen = function(event){
        setMessageInnerHTML("open");
    }

    //接收到消息的回調方法
    websocket.onmessage = function(event){
        setMessageInnerHTML(event.data);
    }

    //連接關閉的回調方法
    websocket.onclose = function(){
        setMessageInnerHTML("close");
    }

    //監聽窗口關閉事件，當窗口關閉時，主動去關閉websocket連接，防止連接還沒斷開就關閉窗口，server端會拋異常。
    window.onbeforeunload = function(){
        websocket.close();
    }

    //將消息顯示在網頁上
    function setMessageInnerHTML(innerHTML){
        document.getElementById('message').innerHTML += innerHTML + '<br/>';
    }

    //關閉連接
    function closeWebSocket(){
        websocket.close();
    }

    //發送消息
    function send(){
        var message = document.getElementById('text').value;
        websocket.send(message);
    }
</script>
</html>

主動發送測試：

package com.lsj.websocket;

import com.lsj.websocket.server.WebSocketServer;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.RequestMapping;

import javax.annotation.Resource;
import java.io.IOException;

@Controller
@RequestMapping("/")
public class SendTest {

    @Resource
    WebSocketServer webSocketServer;


    @RequestMapping("/send")
    public void send(){
        try {
            webSocketServer.sendToUser("普通用戶信息，管理員收得到 ","1");
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @RequestMapping("/send2")
    public void send2(){
        try {
            webSocketServer.sendToUser("管理員專屬信息 ","2");
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

}

總結：主要在調用new WebSocket()時傳入userId、ip、isAdmin 進行用戶頁面session綁定。實現 相同用戶多瀏覽器、管理員角色發送。