用 go-zero 写一个健壮的websocket对话服务包含jwt认证-Go服务器开发

下面是使用 Go-Zero 创建一个健壮的 WebSocket 对话服务，并包含 JWT 认证的完整示例。

项目结构

首先，创建项目结构如下：

websocket-chat/
│
├── main.go
├── jwt.go
└── go.mod

1. 初始化 Go-Zero 项目

确保你已经安装了 Go-Zero。然后在项目目录中运行以下命令来初始化模块：

go mod init websocket-chat

接下来，安装所需依赖项：

go get -u github.com/zeromicro/go-zero
go get -u github.com/gorilla/websocket
go get -u github.com/dgrijalva/jwt-go

2. 编写 JWT 生成和验证逻辑 (`jwt.go`)

在 jwt.go 中实现 JWT 的生成和验证逻辑：

package main

import (
	"time"

	"github.com/dgrijalva/jwt-go"
)

var jwtSecret = []byte("your_secret_key")

// GenerateToken 生成JWT Token
func GenerateToken(username string) (string, error) {
	claims := jwt.StandardClaims{
		ExpiresAt: time.Now().Add(time.Hour * 24).Unix(),
		Issuer:    username,
	}
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString(jwtSecret)
}

// JwtMiddleware JWT中间件，用于验证Token并解析用户信息。
func JwtMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		token := r.Header.Get("Authorization")
		
		if token == "" {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		tkn, err := jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
			return jwtSecret, nil
		})

		if err != nil || !tkn.Valid {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

        // 如果需要，可以在这里提取用户信息并附加到请求上下文中

		next(w, r)
	}
}

3. 实现 WebSocket 服务 (`main.go`)

在 main.go 文件中实现 WebSocket 服务以及登录功能：

package main

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

type Client struct {
	conn *websocket.Conn
	send chan []byte
}

var clients = make(map[*Client]bool)

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

func handleConnections(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
	    log.Println("Error during connection upgrade:", err)
	    return 
    }
    
	client := &Client{conn: conn, send: make(chan []byte)}
	
	defer func() {
	    conn.Close()
	    delete(clients, client)
    }()
	
	for {
	    _, message, err := conn.ReadMessage()
	    if err != nil {
	        log.Println("Error reading message:", err)
	        break 
        }
	    
        // 广播消息给所有连接的客户端。
	    for c := range clients {
	        select {
	        case c.send <- message:
	        default:
	            close(c.send)
	            delete(clients,c)
	        }
	    }
   }
}

func loginHandler(w http.ResponseWriter, r *http.Request) {
	username := r.FormValue("username")
	password := r.FormValue("password")

	if username == "user" && password == "pass" { // 简单示例，应替换为真实验证逻辑。
	    token, _ := GenerateToken(username)
	    w.Write([]byte(token))
	    return 
    }

	http.Error(w,http.StatusText(http.StatusUnauthorized),http.StatusUnauthorized)
}

func main() {

	http.HandleFunc("/ws", JwtMiddleware(handleConnections))
	http.HandleFunc("/login", loginHandler)

	go func() {
        for client := range clients {
            for msg := range client.send { 
                if err := client.conn.WriteMessage(websocket.TextMessage,msg); err != nil { 
                    log.Println("Error writing message:", err)
                    break 
                }
            }   
        }  
    }()
	
	log.Fatal(http.ListenAndServe(":8080", nil))
}

4. 测试 WebSocket 服务

启动你的应用程序：

go run main.go jwt.go

使用 Postman 登录并获取 JWT

发起 POST 请求到 /login 接口，提供用户名和密码，例如： “` POST /login HTTP/1.1

username=user&password=pass


2. 如果成功，你将收到一个 JWT 响应。

#### 使用 WebSocket 客户端进行连接

通过 WebSocket 客户端（如 JavaScript 或 Postman）连接到 `/ws` 接口，同时添加 `Authorization` 标头以包含你从登录中获得的 JWT，例如：

```javascript
const socket = new WebSocket('ws://localhost:8080/ws', [], { headers: {'Authorization': 'Bearer your_token_here'} });

socket.onopen = function(event) {
    console.log('WebSocket is connected.');
};

socket.onmessage = function(event) {
    console.log('Message from server:', event.data);
};

这样，你就成功搭建了一个简单但健壮的 WebSocket 对话服务，并且添加了 JWT 认证机制。在实际生产环境中，请确保加强安全性、完善错误处理和日志记录等。

内容由零声教学AI助手提供，问题来源于学员提问