使用 Golang、Fiber/Websockets 实现简单的群聊功能

本文基于 Golang + WebSocket 实现简单 “群组 “聊天功能，许多其他文章都使用一个简单的聊天示例来探讨 WebSocket。但是，有时一个简单的聊天示例是不够的。

背景

我想在项目中实现一个功能，在客户端之间实时同步数据。我有以下几种情况。

虽然不是所有客户都能收到消息，但只有熟悉某个组/组织/团队的客户/用户才需要收到这条消息。
发送者/发起者不需要收到消息。

了解 Golang websocket 代码

代码的起点是main函数。使用 Fiber.New() 创建一个新的Fiber应用程序，将所有端点、Web 套接字和配置注册到应用程序对象中。

import (
 "flag"
 "github.com/gofiber/contrib/websocket"
 "github.com/gofiber/fiber/v2"
 "log"
)
//...... Other codes in between
func main() {
 app := fiber.New()

使用中间件升级到 Web-socket

app.Use("/ws", func(c *fiber.Ctx) error {
  if websocket.IsWebSocketUpgrade(c) {
   c.Locals("allowed", true)
   // Your authentication process goes here. Get the Token from header and validate it
   // Extract the claims from the token and set them to the Locals
   // This is because you cannot access headers in the websocket.Conn object below
   c.Locals("GROUP", string(c.Request().Header.Peek("GROUP")))
   c.Locals("USER", string(c.Request().Header.Peek("USER")))
   return c.Next()
  }
  return fiber.ErrUpgradeRequired
 })

该中间件会检查传入的请求是否用于 WebSocket 升级。如果为真，它就会设置一个本地标志，并允许请求继续进行。在这里，你可以将所有验证和身份验证逻辑放在一起，以便继续进行，例如验证访问令牌和设置本地标志供以后使用。

在我的例子中，我将解析访问令牌中的请求，并将其设置为 GROUP 和 USER 键中的 Locals。(出于演示目的，我在这里跳过了令牌解析）。

结构定义

type mappedConns map[string]map[string]*websocket.Conn // Modified type

type ClientObject struct {
 GROUP string
 USER  string
 conn  *websocket.Conn
}

type BroadcastObject struct {
 MSG  string
 FROM ClientObject
}

var clients = make(mappedConns) // Initialized as a nested map
var register = make(chan ClientObject)
var broadcast = make(chan BroadcastObject)
var unregister = make(chan ClientObject)

在我的使用案例中，我有一个与特定组织相关联的用户池。因此，我想将它们存储为下面的示例。因此相应地定义了 mappedConns。我还初始化了一组用于注册、广播和取消注册连接的通道。

{ 
    "ORG1": { 
        "User1": ConnectionObject (), 
        "User2" : ConnectionObject (), 
        "User3" : ConnectionObject (), 
    }, 
    "ORG2": { 
        "User3": ConnectionObject (), 
        "User4" : ConnectionObject (), 
        "User5" : ConnectionObject () 
    } 
}

Websocket 路由处理

app.Get( "/ws" , websocket.New( func (c *websocket.Conn) { 
    clientObj := ClientObject{ 
       GROUP: c.Locals( "GROUP" ).( string ), 
       USER: c.Locals( "USER " ).( string ), 
       conn: c, 
    } 
    defer  func () { 
       unregister <- clientObj 
       c.Close() 
    }()

该路由必须处理连接的注册、广播和取消注册。上面的代码块显示了取消注册连接的清理过程。这里使用的”<-“是通过通道发送/获取数据的语法。

业务逻辑

我定义了一个函数 socketHandler()，它是一个 goroutine，即在主程序继续执行时，它将在后台独立运行。

在 socketHandler()函数中，一个开关情况处理传入的信息/注册/取消注册到代码中定义的通道（注册、广播、取消注册）。

case message := <-broadcast:
   for org, users := range clients {
    if org == message.FROM.GROUP {
     for user, conn := range users {
      if org != message.FROM.GROUP || user != message.FROM.USER {
       if err := conn.WriteMessage(websocket.TextMessage, []byte(message.MSG)); err != nil {
        log.Println("write error:", err)
        removeClient(org, user) // Update client removal
        conn.WriteMessage(websocket.CloseMessage, []byte{})
        conn.Close()
       }
      }
     }
    }
   }

对于每条广播信息，代码都会在客户端列表中循环。如果消息的 GROUP 值与客户端的 GROUP 相匹配。它就会发送给 GROUP 中的所有用户，但与 FROM 用户匹配的用户除外。

是的，我知道不使用 for 循环会更好，相反，应该使用更快的“dictionary search”。希望你能在代码中进行优化。