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

使用 Go-Zero 框架构建一个高并发的 WebSocket 对话服务，并包含 JWT 认证的示例代码如下。首先确保你已经安装了 Go-Zero 和其他必要的依赖。

1. 创建项目

创建一个新的 Go 项目：

mkdir websocket-chat
cd websocket-chat
go mod init websocket-chat

2. 安装依赖

使用以下命令安装所需依赖：

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

3. 编写 JWT 认证相关代码

创建一个 auth.go 文件，处理 JWT 的生成和验证：

package main

import (
	"fmt"
	"time"

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

var secretKey = []byte("your_secret_key")

func GenerateToken(username string) (string, error) {
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
		"username": username,
		"exp":      time.Now().Add(time.Hour * 24).Unix(),
	})

	return token.SignedString(secretKey)
}

func ParseToken(tokenStr string) (string, error) {
	token, err := jwt.Parse(tokenStr, func(token *jwt.Token) (interface{}, error) {
		return secretKey, nil
	})
	if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid {
		return claims["username"].(string), nil
	}
	return "", err
}

4. 创建 WebSocket 服务

在 main.go 中编写 WebSocket 服务逻辑：

package main

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

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

func handleConnection(w http.ResponseWriter, r *http.Request) {
	token := r.Header.Get("Authorization")
	username, err := ParseToken(token)
	if err != nil {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		logx.Error("Failed to upgrade connection:", err)
		return
	}
	defer conn.Close()

	logx.Infof("%s connected", username)

	for {
		var msg string
		err := conn.ReadMessage(&msg)
        if err != nil {
            logx.Error("Read message error:", err)
            break
        }
		
        logx.Infof("Received message from %s: %s", username, msg)

        // Echo the message back to the client.
        if err = conn.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
            logx.Error("Write message error:", err)
            break 
        }
    }
}

func main() {
	http.HandleFunc("/ws", handleConnection)
	logx.Info("Server started at :8080")
	if err := http.ListenAndServe(":8080", nil); err != nil {
	    logx.Fatal(err)
    }
}

5. 启动服务器

在终端中运行你的 Go 应用程序：

go run main.go auth.go

6. 测试 WebSocket 服务

可以使用一个简单的 HTML 客户端来测试这个服务，比如下面的代码，可以保存在 index.html 中:

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>WebSocket Chat</title>
</head>
<body>
    <input id="messageInput" type="text"/>
    <button onclick="sendMessage()">Send</button>

    <script>
        const token = "YOUR_JWT_TOKEN"; // Replace with a valid JWT token.
        const socket = new WebSocket(`ws://localhost:8080/ws`, [], { headers: { Authorization: token } });

        socket.onopen = function () {
            console.log('WebSocket connection established');
        };

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

        function sendMessage() {
            const input = document.getElementById('messageInput');
            socket.send(input.value);
            input.value = '';
        }
    </script>
</body>
</html>

确保将 YOUR_JWT_TOKEN 替换为有效的JWT令牌。打开浏览器访问此 HTML 文件并开始发送消息。