# websocket-module

**Repository Path**: dreamwithouttrace/websocket-module

## Basic Information

- **Project Name**: websocket-module
- **Description**: 嘻嘻嘻嘻嘻嘻嘻嘻嘻嘻
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-08-03
- **Last Updated**: 2026-02-04

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# WebSocket Module

[![Go Version](https://img.shields.io/badge/Go-1.24-blue)](https://go.dev/)
[![License](https://img.shields.io/badge/License-MIT-green)](LICENSE)

这是一个高性能、可复用的 WebSocket 服务器模块，提供了完整的 WebSocket 连接管理、消息处理、用户认证、gRPC 服务等功能。

## 特性

- **高性能**: 使用分片连接池管理大量并发连接
- **安全认证**: 支持 Token 认证
- **实时统计**: 提供在线用户统计和监控（按性别统计）
- **消息确认**: 支持消息 ACK 机制，确保消息可靠传递
- **限流保护**: 内置请求限流机制
- **模块化设计**: 清晰的模块分离，易于扩展和定制
- **完整 API**: 提供 RESTful API 进行服务器管理
- **gRPC 支持**: 内置 gRPC 服务，支持远程调用
- **连接管理**: 支持连接踢出、在线状态查询等功能
- **心跳保活**: 自动心跳检测和超时踢出

## 目录结构

```
websocket-module/
├── core/              # 核心接口和类型定义
│   ├── types.go       # 核心类型（User、Connection、Config等）
│   └── connection.go   # WebSocket连接实现
├── transport/         # 传输层管理（连接池、分片管理）
│   └── transport.go   # 传输管理器实现
├── handlers/          # 处理器
│   ├── handshake.go   # WebSocket握手处理
│   └── api.go        # RESTful API处理
├── middleware/        # 中间件
│   └── middleware.go  # 认证和限流中间件
├── protocol/          # 消息协议定义
│   └── types.go      # 消息格式定义
├── service/           # gRPC 服务定义
│   ├── handler.proto  # gRPC proto定义
│   ├── handler.pb.go # 生成的处理文件
│   ├── handler_grpc.pb.go # gRPC生成的存根文件
│   └── service.go    # gRPC服务实现
├── monitoring/        # 监控配置
│   └── prometheus.yml # Prometheus配置
├── examples/          # 使用示例
│   ├── basic/         # 基础使用示例
│   ├── custom_user/   # 自定义用户实现示例
│   └── grpc_client/   # gRPC客户端示例
├── server.go          # 主服务器
├── go.mod            # Go模块文件
└── README.md          # 文档
```

## 快速开始

### 1. 安装依赖

```bash
go mod download
```

### 2. 基本使用

```go
package main

import (
	"log"
	"time"

	"gitee.com/dreamwithouttrace/websocket-module/core"
	"gitee.com/dreamwithouttrace/websocket-module/protocol"
	"gitee.com/dreamwithouttrace/websocket-module"

	"github.com/go-redis/redis"
	"gorm.io/gorm"
)

// User 实现 core.User 接口
type User struct {
	ID             int64
	UniqueID       string
	Nickname       string
	Avatar         string
	Sex            string // "1"=男性, "2"=女性, 其他=无性别
	SocketID       string
	IsOnline       int
	IsFaker        int
	IsApp          int
	IsWx           int
	LastLoginTime  int64
	LastInnerIP    string
	Score          int
	Birthday       string
	DisabledEndTime int64
	AppID          int64
}

func (u *User) GetID() int64                    { return u.ID }
func (u *User) GetUniqueID() string            { return u.UniqueID }
func (u *User) GetNickname() string            { return u.Nickname }
func (u *User) GetAvatar() string              { return u.Avatar }
func (u *User) GetSex() string                 { return u.Sex }
func (u *User) GetSocketID() string            { return u.SocketID }
func (u *User) GetIsOnline() int                { return u.IsOnline }
func (u *User) SetIsOnline(status int)         { u.IsOnline = status }
func (u *User) GetIsFaker() int                { return u.IsFaker }
func (u *User) GetIsApp() int                   { return u.IsApp }
func (u *User) GetIsWx() int                    { return u.IsWx }
func (u *User) LifecycleDelay()                 {}
func (u *User) SetSocketID(socketID string)     { u.SocketID = socketID }
func (u *User) GetLastLoginTime() int64        { return u.LastLoginTime }
func (u *User) SetLastLoginTime(t int64)       { u.LastLoginTime = t }
func (u *User) GetLastInnerIP() string          { return u.LastInnerIP }
func (u *User) SetLastInnerIP(ip string)       { u.LastInnerIP = ip }
func (u *User) GetScore() int                  { return u.Score }
func (u *User) SetScore(score int)             { u.Score = score }
func (u *User) GetBirthday() string             { return u.Birthday }
func (u *user) SetBirthday(birthday string)    { u.Birthday = birthday }
func (u *user) GetDisabledEndTime() int64      { return u.DisabledEndTime }
func (u *user) GetAppID() int64                { return u.AppID }

func main() {
	// 创建配置
	config := core.DefaultConfig()
	config.Host = "0.0.0.0"
	config.Port = 8080
	config.GrpcHost = "0.0.0.0"
	config.GrpcPort = 8001
	config.TokenName = "token"
	config.SharedCount = 1000 // 连接分片数量

	// 创建Redis客户端
	redisClient := redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "",
		DB:       0,
	})

	// 创建数据库连接（根据你的实际情况初始化）
	var db *gorm.DB // 这里需要根据你的数据库类型初始化

	// 定义用户查询函数
	queryFunc := func(db *gorm.DB, token string) (core.User, error) {
		// 根据token从数据库查询用户
		// 这里实现你的用户查询逻辑
		user := &User{
			ID:       1,
			UniqueID: "user_1",
			Nickname: "测试用户",
			Sex:      "1",
		}
		return user, nil
	}

	// 创建WebSocket服务器
	server := websocket.NewServer(config, redisClient, db, queryFunc)

	// 设置回调函数
	server.SetOnConnect(func(conn core.Connection) error {
		log.Printf("用户连接: %s", conn.GetUserID())
		return nil
	})

	server.SetOnMessage(func(conn core.Connection, user core.User, role interface{}, message *protocol.Message) {
		log.Printf("收到消息: 类型=%s, 用户ID=%d", message.Type, message.UserId)
		// 处理消息逻辑
	})

	server.SetOnDisconnect(func(conn core.Connection, user core.User, role interface{}) {
		log.Printf("用户断开连接: %s", conn.GetUserID())
	})

	// 启动服务器
	log.Println("启动WebSocket服务器...")
	if err := server.Start(); err != nil {
		log.Fatalf("服务器启动失败: %v", err)
	}

	// 阻塞等待
	select {}
}
```

### 3. 自定义配置

```go
config := &core.Config{
	Host:              "0.0.0.0",
	Port:              8080,
	GrpcHost:          "0.0.0.0",
	GrpcPort:          8001,
	TokenName:         "token",
	ReadTimeout:       5 * time.Second,
	WriteTimeout:      5 * time.Second,
	MaxHeaderBytes:    1 << 20,
	SharedCount:       1000,        // 连接分片数量
	EnableCompression: true,
	ReadBufferSize:    4096,
	WriteBufferSize:   4096,
}
```

## API 接口

### WebSocket 连接

- `GET /ws?token=xxx` - 标准 WebSocket 连接（需要认证）

### 管理 API

| 方法 | 路径 | 描述 |
|------|------|------|
| POST | `/kick?userId=123` | 踢出指定用户 |
| GET | `/online/lists` | 获取在线用户列表 |
| GET | `/get/online?id=123` | 获取指定用户信息 |
| GET | `/check/content?keyword=xxx` | 检查内容（GET方式） |
| POST | `/check/content` | 检查内容（POST方式） |
| GET | `/statistics` | 获取统计信息 |

### 统计信息响应示例

```json
{
    "code": 0,
    "msg": "success",
    "data": {
        "total": 100,
        "boy": 50,
        "girl": 40,
        "noSex": 10
    }
}
```

## gRPC 服务

服务器同时提供 gRPC 服务，默认监听在 `GrpcHost:GrpcPort`。

### gRPC 方法

| 方法 | 描述 |
|------|------|
| `Kick(KickRequest) returns (KickResponse)` | 踢出用户 |
| `Exists(ExistsRequest) returns (ExistsResponse)` | 检查用户是否在线 |
| `Push(MessageRequest) returns (MessageResponse)` | 推送消息 |
| `Update(UpdateRequest) returns (UpdateResponse)` | 更新用户信息 |

### gRPC 消息类型

```protobuf
message KickRequest {
  int64 id = 1;
}

message ExistsRequest {
  int64 id = 1;
}

message MessageRequest {
  int64 userId = 1;
  Message data = 2;
}

message UpdateRequest {
  int64 id = 1;
}
```

## 核心接口

### User 接口

```go
type User interface {
	// GetID 账号唯一ID
	GetID() int64
	// GetUniqueID 账号唯一ID
	GetUniqueID() string
	// GetNickname 昵称
	GetNickname() string
	// GetAvatar 头像
	GetAvatar() string
	// GetSex 性别
	GetSex() string // "1"=男性, "2"=女性, 其他=无性别
	// GetSocketID 节点编号
	GetSocketID() string
	// GetIsOnline 是否在线
	GetIsOnline() int
	// SetIsOnline 设置在线状态
	SetIsOnline(status int)
	// GetIsFaker 是否假用户
	GetIsFaker() int
	// GetIsApp 是否APP用户
	GetIsApp() int
	// GetIsWx 是否微信用户
	GetIsWx() int
	// LifecycleDelay Token生命周期刷新
	LifecycleDelay()
	// SetSocketID 设置节点信息
	SetSocketID(socketID string)
	// GetLastLoginTime 最后登录时间
	GetLastLoginTime() int64
	// SetLastLoginTime 设置最后登录时间
	SetLastLoginTime(time int64)
	// GetLastInnerIP 最后接入时的IP
	GetLastInnerIP() string
	// SetLastInnerIP 设置接入IP
	SetLastInnerIP(ip string)
	// GetScore 评分
	GetScore() int
	// SetScore 设置评分
	SetScore(score int)
	// GetBirthday 生日
	GetBirthday() string
	// SetBirthday 生日设置
	SetBirthday(birthday string)
	// GetDisabledEndTime 禁用结束时间
	GetDisabledEndTime() int64
	// GetAppID 所属应用
	GetAppID() int64
}
```

### Connection 接口

```go
type Connection interface {
	// GetUser 获取 User 用户信息
	GetUser() User
	// SetUser 设置 User 用户信息
	SetUser(user User)
	// GetChannelID 当前链接编号
	GetChannelID() string
	// GetUserID User 用户ID
	GetUserID() string
	// GetIP 链接时的IP
	GetIP() string
	// WriteMessage 给客户端推送信息
	WriteMessage(data *protocol.Message) error
	// WriteByteMessage 给客户端推送byte字节信息
	WriteByteMessage(data []byte) error
	// Close 关闭链接
	Close() error
	// IsConnected 是否链接
	IsConnected() bool
	// SetConnected 设置是否链接
	SetConnected(connected bool)
	// GetWaits 发送后待确认的信息
	GetWaits() *Waits
	// GetSendNo 消息最新的编号
	GetSendNo() int64
	// IncrementSendNo 自增消息编号
	IncrementSendNo()
	// Kick 踢用户下线
	Kick()
	// Ping 查看链接是否接通
	Ping()
	// GetRole 链接所绑定的角色信息
	GetRole() interface{}
	// SetRole 为链接绑定角色信息
	SetRole(role interface{})
	// Pong Pong信息
	Pong()
	// GetUniqueID 获取User用户唯一ID
	GetUniqueID() string
	// SetReadDeadline 设置超时时间
	SetReadDeadline(second int) error
	// ConnectionDurationThisTime 其他
	ConnectionDurationThisTime() *ClientInfo
}
```

## 消息格式

### 协议消息结构

```go
type Message struct {
	Type      string                 `json:"type"`      // 消息类型
	ID        string                 `json:"id"`        // 消息ID（唯一标识符）
	Timestamp int64                  `json:"timestamp"` // 时间戳
	IsGroup   bool                   `json:"is_group"`  // 是否群组消息
	NeedAck   bool                   `json:"need_ack"`  // 是否需要ACK确认
	UserId    int64                  `json:"user_id"`   // 用户ID
	Channel   string                 `json:"channel"`   // 频道ID
	SendInfo  map[string]interface{} `json:"send_info,omitempty"` // 发送者信息
	Data      interface{}            `json:"data"`      // 消息数据
}
```

### 客户端发送消息示例

```json
{
    "type": "chat",
    "id": "msg_123456",
    "timestamp": 1640995200,
    "need_ack": true,
    "user_id": 123,
    "channel": "channel123",
    "is_group": false,
    "data": {
        "message": "Hello World",
        "roomId": "room1"
    }
}
```

### 服务器推送消息示例

```json
{
    "type": "message",
    "id": "msg_789012",
    "timestamp": 1640995200,
    "need_ack": false,
    "user_id": 0,
    "channel": "channel123",
    "is_group": false,
    "data": {
        "content": "Hello from server",
        "timestamp": 1640995200
    }
}
```

## 中间件

### 认证中间件

认证中间件用于验证用户身份，支持从 Redis 和数据库查询用户信息。

```go
auth := middleware.NewAuthMiddleware(redis, db, queryFunc)
user, err := auth.Process(config, r)
if err != nil {
	http.Error(w, err.Error(), http.StatusUnauthorized)
	return
}
```

### 限流中间件

限流中间件用于限制请求频率。

```go
rateLimit := middleware.NewRateLimitMiddleware(redis, 100, time.Minute)
// 在路由中使用限流中间件
```

## 配置说明

| 配置项 | 类型 | 默认值 | 说明 |
|--------|------|--------|------|
| Host | string | "127.0.0.1" | HTTP服务器监听地址 |
| Port | int | 14101 | HTTP服务器监听端口 |
| GrpcHost | string | - | gRPC服务器监听地址 |
| GrpcPort | int | - | gRPC服务器监听端口 |
| TokenName | string | - | Token参数名称 |
| ReadTimeout | time.Duration | 2s | 读取超时时间 |
| WriteTimeout | time.Duration | 2s | 写入超时时间 |
| MaxHeaderBytes | int | 1MB | 最大请求头大小 |
| SharedCount | int64 | 1000 | 连接分片数量 |
| EnableCompression | bool | true | 是否启用压缩 |
| ReadBufferSize | int | 4096 | 读取缓冲区大小 |
| WriteBufferSize | int | 4096 | 写入缓冲区大小 |

## 扩展开发

### 自定义用户实现

需要实现 `core.User` 接口，示例见 `examples/custom_user/main.go`。

### 自定义中间件

实现 `middleware.Middleware` 接口：

```go
type CustomMiddleware struct {
	// 自定义字段
}

func (m *CustomMiddleware) Process(w http.ResponseWriter, r *http.Request) (core.User, error) {
	// 自定义处理逻辑
	return user, nil
}
```

### 添加自定义路由

```go
server.AddRouter("GET", "/custom", func(w http.ResponseWriter, r *http.Request) {
	// 自定义处理逻辑
})
```

## 性能优化

1. **连接分片**: 使用 `SharedCount` 参数调整分片数量，提高并发性能
2. **缓冲区大小**: 根据消息大小调整 `ReadBufferSize` 和 `WriteBufferSize`
3. **超时设置**: 根据网络环境调整 `ReadTimeout` 和 `WriteTimeout`
4. **压缩**: 对于文本消息，启用 `EnableCompression` 可以提高传输效率

## 监控和调试

### 获取统计信息

```go
total, boy, girl, noSex := server.GetClientsStatic()
log.Printf("在线用户: %d, 男性: %d, 女性: %d, 无性别: %d", total, boy, girl, noSex)
```

### 通过 API 获取统计

```bash
curl http://localhost:8080/statistics
```

## 依赖项

- `github.com/gorilla/websocket` - WebSocket 实现
- `github.com/gorilla/mux` - HTTP 路由
- `github.com/go-redis/redis` - Redis 客户端
- `gorm.io/gorm` - ORM 框架
- `google.golang.org/grpc` - gRPC 框架
- `google.golang.org/protobuf` - Protocol Buffers

## 注意事项

1. **Redis 依赖**: 模块依赖 Redis 进行连接限制和缓存，请确保 Redis 服务可用
2. **数据库接口**: 需要实现数据库接口来获取用户信息，并提供 `queryFunc` 函数
3. **错误处理**: 建议在生产环境中添加完善的错误处理和日志记录
4. **安全考虑**: 在生产环境中，请确保认证逻辑的安全性
5. **gRPC 服务**: 如果不需要 gRPC 服务，可以不设置 `GrpcHost` 和 `GrpcPort`

## 示例代码

详细的使用示例请参考 `examples/` 目录：

- `examples/basic/main.go` - 基本使用示例
- `examples/custom_user/main.go` - 自定义用户实现示例
- `examples/grpc_client/main.go` - gRPC 客户端示例

## 许可证

MIT License