用于用户通知的带有 gorilla/websocket 的 Websocket

Question

我正在尝试使用 gorilla/websocket 包实现 websocket。我按照这个例子，它工作：https://github.com/gorilla/websocket/tree/master/examples/chat 但是，我正在尝试根据 user_id 向特定用户发送消息。

JS（vue.js）代码（简体）：

 data() {
        return {
            user_id: 2,
            username: "test",
            userWebSocket: ""
        }
    },
 connectUserToWebSocket() {
        if (window.WebSocket) {
            var loc = window.location;
            var uri = 'ws:';

            if (loc.protocol === 'https:') {
                uri = 'wss:';
            }
            uri += '//localhost:8000/api/ws';


            this.userWebSocket = new WebSocket(uri)

            this.userWebSocket.onopen = function() {
                console.log('Connected')
            }

            this.userWebSocket.onmessage = function(evt) {
                alert(evt.data)
            }
        }
    }

还有 main.go（简化版）：

import (
    "github.com/gorilla/mux"
    "github.com/urfave/negroni"
    "test/notifications"
)

func Router() http.Handler {
  r := mux.NewRouter()

  hub := notifications.NewHub()
  go hub.Run()

  r.HandleFunc("/", controller.Root).Methods("GET")

  //r.HandleFunc("/api/ws", notifications.HandleConnections)
  r.HandleFunc("/api/ws", func(w http.ResponseWriter, r *http.Request) {
    notifications.ServeWs(hub, w, r)
  })

还有 notification.go： 包裹通知

import (
  "net/http"
  "log"
  "github.com/gorilla/websocket"
  "time"
  "bytes"
  "fmt"
)

const (
  // Time allowed to write a message to the peer.
  writeWait = 10 * time.Second

  // Time allowed to read the next pong message from the peer.
  pongWait = 60 * time.Second

  // Send pings to peer with this period. Must be less than pongWait.
  pingPeriod = (pongWait * 9) / 10

  // Maximum message size allowed from peer.
  maxMessageSize = 512
)

var (
  newline = []byte{'\n'}
  space   = []byte{' '}
)

var upgrader = websocket.Upgrader{
 ReadBufferSize:  1024,
 WriteBufferSize: 1024,
 CheckOrigin: func(r *http.Request) bool {
    return true
 },
}


type Client struct {
  hub *Hub

  Id int

  // The websocket connection.
  conn *websocket.Conn

  // Buffered channel of outbound messages.
  send chan []byte
}

type Hub struct {
  // Registered clients.
  clients map[*Client]bool

  users map[int]*Client

  // Inbound messages from the clients.
  broadcast chan []byte

  // Register requests from the clients.
  register chan *Client

  // Unregister requests from clients.
  unregister chan *Client
}

func NewHub() *Hub {
  return &Hub{
    broadcast:  make(chan []byte),
    register:   make(chan *Client),
    unregister: make(chan *Client),
    clients:    make(map[*Client]bool),
    users:      make(map[int]*Client),
 }
}

func (h *Hub) Run() {
 for {
    select {
    case client := <-h.register:
        h.clients[client] = true
    case client := <-h.unregister:
        if _, ok := h.clients[client]; ok {
            delete(h.clients, client)
            close(client.send)
        }
    case message := <-h.broadcast:
        for client := range h.clients {
            select {
            case client.send <- message:
            default:
                close(client.send)
                delete(h.clients, client)
            }
        }
    }
 }
}

// readPump pumps messages from the websocket connection to the hub.
//
// The application runs readPump in a per-connection goroutine. The    application
// ensures that there is at most one reader on a connection by executing all
// reads from this goroutine.
func (c *Client) readPump() {
 defer func() {
    c.hub.unregister <- c
    c.conn.Close()
}()
c.conn.SetReadLimit(maxMessageSize)
c.conn.SetReadDeadline(time.Now().Add(pongWait))
c.conn.SetPongHandler(func(string) error {     c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
for {
    _, message, err := c.conn.ReadMessage()
    if err != nil {
        if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
            log.Printf("error: %v", err)
        }
        break
    }
    message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))
    c.hub.broadcast <- message
 }
}

// writePump pumps messages from the hub to the websocket connection.
//
// A goroutine running writePump is started for each connection. The
// application ensures that there is at most one writer to a connection by
// executing all writes from this goroutine.
func (c *Client) writePump() {
 ticker := time.NewTicker(pingPeriod)
 defer func() {
    ticker.Stop()
    c.conn.Close()
 }()
 for {
    select {
    case message, ok := <-c.send:
        c.conn.SetWriteDeadline(time.Now().Add(writeWait))
        if !ok {
            // The hub closed the channel.
            c.conn.WriteMessage(websocket.CloseMessage, []byte{})
            return
        }

        w, err := c.conn.NextWriter(websocket.TextMessage)
        if err != nil {
            return
        }
        w.Write(message)

        // Add queued chat messages to the current websocket message.
        n := len(c.send)
        for i := 0; i < n; i++ {
            w.Write(newline)
            w.Write(<-c.send)
        }

        if err := w.Close(); err != nil {
            return
        }
    case <-ticker.C:
        c.conn.SetWriteDeadline(time.Now().Add(writeWait))
        if err := c.conn.WriteMessage(websocket.PingMessage, nil); err !=    nil {
            return
        }
    }
 }
}

// serveWs handles websocket requests from the peer.
func ServeWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
  conn, err := upgrader.Upgrade(w, r, nil)
  if err != nil {
    fmt.Println(err)
    log.Println(err)
    return
  }
  client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)}
  client.hub.register <- client

  // Allow collection of memory referenced by the caller by doing all work in
  // new goroutines.
  go client.writePump()
  go client.readPump()
}

问题是我不知道如何将 JS 中的 user_id 添加到连接中。我也想保持它的安全。

如您所见，我尝试了以下方法：

type Client struct {
  ...
  Id int
}

还有：

type Hub struct {
  // Registered clients.
  clients map[*Client]bool

  users map[int]*Client
  ...
}

还有：

func NewHub() *Hub {
  return &Hub{
    ...
    clients:    make(map[*Client]bool),
    users:      make(map[int]*Client),
  }
}

但是我还是不知道怎么用user_id打开websocket。

TL:DR 1. 我有工作代码向所有用户发送通知。 2. 我想根据 ID 向特定用户发送通知。

根本没有帮助的类似问题： How to send to only one client and not all clients using Go and gorilla websocket

Answer 1

我试图做同样的事情，发送系统通知。我使用了这个示例并根据需要对其进行了修改：https://github.com/AnupKumarPanwar/Golang-realtime-chat-rooms

当设备连接到 websocket 时，也会发送用户 ID，如下所示： ws://localhost:8081/ws/connect/{userID} 或任何您的网址。您可以做用户名，然后在服务器上查找用户 ID，或者只使用用户名作为“房间 ID”

然后将该用户添加到房间：

// serveWs handles websocket requests from the peer.
func serveWs(w http.ResponseWriter, r *http.Request) {
    vars := mux.Vars(r)

    upgrader.CheckOrigin = func(r *http.Request) bool { return true }
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println(err.Error())
        return
    }
    c := &connection{send: make(chan []byte, 256), ws: ws}
    s := subscription{c, vars["userID"]}
    h.register <- s
    go s.writePump()
    go s.readPump()
}

"s := subscription{c, vars["userID"]}" 是在您的情况下使用用户 ID 创建特定“房间”的方式。

要向那个房间广播消息，只需这样：

func NewNotification(userID string, msg []byte) error {
    m := message{msg, userID}
    h.broadcast <- m
    return nil
}

让您的通知调用该函数向房间广播。

希望这是有道理的！