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

使用 Go-Zero 创建一个健壮的 WebSocket 对话服务并实现 JWT 认证的基本步骤如下：

1. 安装 Go-Zero

首先，确保你已经安装了 Go-Zero。可以通过以下命令安装：

go get -u github.com/zeromicro/go-zero

2. 初始化项目

使用 go-zero 命令行工具来初始化项目。

go-ctl new websocket-chat
cd websocket-chat

3. 编写 JWT 中间件

创建一个文件 jwt.go 用于实现 JWT 的生成和验证：

package middleware

import (
	"net/http"
	"time"

	"github.com/dgrijalva/jwt-go"
	"github.com/zeromicro/go-zero/rest/httpx"
)

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)
	}
}

4. 创建 WebSocket 服务

在 main.go 文件中创建 WebSocket 服务：

package main

import (
    "net/http"
    "github.com/gorilla/websocket"
    "github.com/zeromicro/go-zero/rest/httpx"
    "your_project/middleware" // 替换为你的项目路径
)

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

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

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

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

func main() {

	http.HandleFunc("/ws", middleware.JwtMiddleware(handleConnections))

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

5. 用户登录和获取 Token

你可能还需要一个登录接口，允许用户通过用户名和密码进行身份验证并获取 JWT。例如，你可以在 main.go 中添加如下代码：

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

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

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

然后将该处理程序添加到路由中：

http.HandleFunc("/login", loginHandler)

6. 测试 WebSocket 服务

启动应用程序并测试。使用 Postman 或类似工具请求 /login 接口以获得 JWT，并在请求 /ws 时将其包含在 Authorization 标头中。

请注意，以上代码是一个简化示例，具体生产环境应考虑更复杂的错误处理、日志记录以及安全性等因素。同时，请确保引入必要的依赖项，例如 github.com/gorilla/websocket 和 github.com/dgrijalva/jwt-go。

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