需求

每隔几秒主动调用api拉取信息，会对服务器造成很大的压力，特别对于小型应用，只有一台服务器，配置也不高。客户端要实时的接收一些消息，消息推送也提升了用户的体验。很有必要。

思路

• 使用RabbitMQ接收消息，再推送给客户端
• 使用gin http api接收消息，再推送给客户端
• 客户端通过websocket连接到推送服务，然后等待接收消息并处理消息
• 使用jwt-token验证客户端真实性，确定用户身份

实现

设置模块

读取进行目录的config.toml文件，做为配置文件。

• rabbit_conn_str RabbitMQ消息队列配置串，root是用户名，123是密码，@后为主机端口，/test是租户

• jwt_secret jwt使用的加密串，与主服务保证一致

• public_key_path private_key_path https配置，公钥与私钥，不配置不开启https

• http_port http服务端口号

• white_ips gin http 服务的白名单，使用白名单简单验证接口调用者的身份

rabbit_conn_str = "amqp://root:123@localhost:5672/test" #2323
jwt_secret = '123123123123123123123'
public_key_path = "./server.crt"
private_key_path = "./server.key"
http_port = "80"
white_ips = "127.0.0.1"

package main

import (
	"fmt"
	"github.com/BurntSushi/toml"
	"os"
)

var CONFIG = GetConfig("config.toml")

// Config
// 系统配置类
type Config struct {
	RabbitConnStr  string `json:"rabbit_conn_str" toml:"rabbit_conn_str"`
	JwtSecret      string `json:"jwt_secret" toml:"jwt_secret"`
	PublicKeyPath  string `json:"public_key_path" toml:"public_key_path"`
	PrivateKeyPath string `json:"private_key_path" toml:"private_key_path"`
	HttpPort       string `json:"http_port" toml:"http_port"`
	WhiteIps       string `json:"white_ips" toml:"white_ips"`
}

// GetConfig
// 获取配置
func GetConfig(path string) Config {
	content, err := os.ReadFile(path)
	if err != nil {
		panic(err)
	}
	config := Config{}
	mData, err := toml.Decode(string(content), &config)
	if err != nil {
		panic(err)
	}
	fmt.Printf("%s", mData)
	return config
}

日志模块

日志使用滴滴的zap

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"

使用"gopkg.in/natefinch/lumberjack.v2"做日志轮换

MaxSize 每个日志文件的最大大小，单位MB

MaxBackups 最多保留多少文件

MaxAge 最多保留几天

Compress 是否压缩日志文件

Filename 日志文件名，加上日志，加上主机名

	File = &lumberjack.Logger{
		Filename:   "./logs/" + time.Now().Format("20060102150405") + hostname + ".log", // ⽇志⽂件路径
		MaxSize:    500,  // 文件最大大小 500MB
		MaxBackups: 3, // 最多保留3个备份
		MaxAge:     7, // 最多保留多少天
		Compress:   true, // 是否压缩 默认不压缩
	}

	FileW := zapcore.AddSync(File)
	encoderCfg := zapcore.EncoderConfig{
		MessageKey:     "msg", // msg key 日志主体内容字段名msg
		LevelKey:       "level", // 日志级别字段名level
		NameKey:        "logger", // 日志名
		EncodeLevel:    zapcore.LowercaseLevelEncoder,
		EncodeTime:     zapcore.ISO8601TimeEncoder,
		EncodeDuration: zapcore.StringDurationEncoder,
	}
	core := zapcore.NewCore(zapcore.NewJSONEncoder(encoderCfg), FileW, zap.DebugLevel)
	LOG = zap.New(core).WithOptions() // options...

关闭日志时先写入磁盘再关闭日志再关闭日志文件

	err := LOG.Sync() // 同步文件到磁盘
	if err != nil {
		fmt.Printf("log sync err %v\n", err)
	}
	if File != nil {
		err = File.Close() // 关闭日志文件
		if err != nil {
			fmt.Printf("sync err %s", err.Error())
		}
	}

jwt 模块

直接上代码。 主要功能： 生成token 与验证token

package main

import (
	"errors"
	"github.com/dgrijalva/jwt-go"
	"time"
)

//GetNewJwtToken 創建一個新的jwt,
func GetNewJwtToken(hs256Key []byte, userClaims *UserClaims) (string, error) {
	jwtToken := jwt.New(jwt.SigningMethodHS256)
	jwtToken.Claims = userClaims
	return jwtToken.SignedString(hs256Key) //TODO: 這個換成rsa加密
}

//CheckJwtToken 檢查jwt是否有效
func CheckJwtToken(hs256Key []byte, jwtString string) (userClaims UserClaims, err error) {
	_, err = jwt.ParseWithClaims(jwtString, &userClaims, func(_ *jwt.Token) (i interface{}, e error) {
		return hs256Key, nil
	})
	return
}

//ErrTimeExp 同一類型的多個錯誤實例
var ErrTimeExp = errors.New("jwt is expiration")
var ErrIatTime = errors.New("jwt in bad format,iat>=exp")

//UserClaims payload的內容，如果是app角色，則appid為空。,web端默認web
type UserClaims struct {
	UserId string `json:"user_id"` //用戶id
	Iat    int64  `json:"iat"`     //發放時間
	Exp    int64  `json:"exp"`     //過期時間
	Jti    string `json:"jti"`     //token 的id ,唯一標識
	Iss    string `json:"iss"`     //簽發者 是sys_user_control
}

//Valid 是否有效，如果沒效，則入出錯誤
func (u UserClaims) Valid() error {
	if time.Now().Unix() > u.Exp {
		return ErrTimeExp
	}
	if u.Iat > u.Exp {
		return ErrIatTime
	}
	return nil
}

websocket模块

用户通过websocket连接到推送服务

UserConn 代表一个连接到服务的客户端，包括用户Id,客户端类型等

type UserConn struct {
	UserId     string          // 用户id
	ClientType string          // android ios web manage-web wx ali 等 一端只能有一个登录
	RemoteAddr string          // 远程地址
	Token      string          // 用户token
	EnterTime  time.Time       // 进来的时间
	Conn       *websocket.Conn // ws conn 与客户端通信用
}

ConnCollect代表通过websocket连接到推送服务的客户端集合。通过userId Map聚合一个有户的不同客户端的连接，再通过token查找到某个连接，继而由这个连接投递消息。

// ConnCollect 集合
type ConnCollect struct {
   // userId => token => collect
   Lock       sync.RWMutex
   Collection map[string]*struct { // userId
      Lock      sync.RWMutex
      UserConns map[string]*UserConn // token
   }
}

添加用户连接

func (co *ConnCollect) AdUserConn(userConn *UserConn) {
	co.Lock.Lock()
	defer co.Lock.Unlock()

	userConns, ok := co.Collection[userConn.UserId]
	if ok {
		userConns.Lock.Lock()
		defer userConns.Lock.Unlock()
		conn, ok := userConns.UserConns[userConn.Token]
		if ok {
			err := conn.Conn.Close()
			if err != nil {
				LOG.Info("close client")
			}
			delete(userConns.UserConns, userConn.Token) // 删除token
		}
		userConns.UserConns[userConn.Token] = userConn // 添加token
	} else {
		co.Collection[userConn.UserId] = &struct {
			Lock      sync.RWMutex
			UserConns map[string]*UserConn
		}{Lock: sync.RWMutex{}, UserConns: map[string]*UserConn{userConn.Token: userConn}}
	}
}

删除用户连接

func (co *ConnCollect) RemoveUserConn(userConn *UserConn) {
	co.Lock.Lock()
	defer co.Lock.Unlock()

	userConns, ok := co.Collection[userConn.UserId]
	if ok {
		userConns.Lock.Lock()
		defer userConns.Lock.Unlock()
		userConnItem, ok := userConns.UserConns[userConn.Token]
		if ok {
			if userConnItem == userConn {
				err := userConn.Conn.Close()
				if err != nil {
					LOG.Info("close client")
				}
				delete(userConns.UserConns, userConn.Token) // 删除token
			}
		}
	} else {
		LOG.Warn("this user conn all ready removed", zap.Any("userConn", userConn))
	}
}

关闭所有连接，准备退出

func (co *ConnCollect) CloseAllForExit() {
	var closeConn = func(userConn *struct {
		Lock      sync.RWMutex
		UserConns map[string]*UserConn
	}) {
		userConn.Lock.Lock()
		defer userConn.Lock.Unlock()
		for _, conn := range userConn.UserConns {
			if err := conn.Conn.Close(); err != nil {
				LOG.Error("close all err:" + err.Error())
			}
		}
	}
	co.Lock.Lock()
	defer co.Lock.Unlock()
	for _, collect := range co.Collection {
		closeConn(collect)
	}
}

webSocket Handle 处理 

// WsHandle
// websocket 服务
type WsHandle struct {
   upgrade         websocket.Upgrader
   UserConnCollect ConnCollect // userId -> user
   Cfg             *Config
}

创建websocket handle

func NewWsHandle() *WsHandle {
	return &WsHandle{
		upgrade: websocket.Upgrader{
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
			// 解决跨域问题
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		}, // http 协议升级
		UserConnCollect: ConnCollect{
			Lock: sync.RWMutex{},
			Collection: map[string]*struct {
				Lock      sync.RWMutex
				UserConns map[string]*UserConn
			}{},
		}, // 客户端连接集合
		Cfg: &CONFIG, // 配置文件
	}
}

// 清除所有连接准备退出
func (handle *WsHandle) CloseAllForExit() {
	handle.UserConnCollect.CloseAllForExit()
}

发送消息

// 发送从http api或rabbitmq接收到的消息 
// userId 要发送给的人
// msg 消息本体内容
func (handle *WsHandle) SendMessage(userId string, msg []byte) error {
	// 读锁
    handle.UserConnCollect.Lock.RLock()
	defer handle.UserConnCollect.Lock.RUnlock()
    // 根据用户id查找用户连接,可能不止一个
	userConns, ok := handle.UserConnCollect.Collection[userId]
	if ok {
        // 用户连接读锁
		userConns.Lock.RLock()
		defer userConns.Lock.RUnlock()
        // 没有连接直接返回
		if len(userConns.UserConns) == 0 {
			return errors.New("no this user")
		}
		for _, conn := range userConns.UserConns {
            // 通过websocket写入消息到客户端
			err := conn.Conn.WriteMessage(websocket.TextMessage, msg)
			if err != nil {
				LOG.Error("send message err", zap.String("user_id", userId), zap.String("msg", string(msg)), zap.String("remote_addr", conn.RemoteAddr))
				err := conn.Conn.Close() // 发送错误主动关闭链接
				if err != nil {
					LOG.Error("close conn err " + err.Error())
				}
			}
		}
		return nil
	} else {
		return errors.New("no this user")
	}
}

通过http api 直接向某个用户主动投递消息

// 发送短信
func (handle *WsHandle) handleSendMessage(c *gin.Context) {

	// 获取用户id 对比ip白名单
	ip := c.ClientIP()
	LOG.Info("ip", zap.String("ip", ip))
	contain := strings.Contains(CONFIG.WhiteIps, ip)
	if !contain || ip == "" {
		LOG.Error("http 发送接口 客户不在ip白名单中", zap.String("ip", ip))
		c.AbortWithStatusJSON(401, map[string]string{"code": "401", "message": "no auth"})
		return
	}
    // 获取表单参数
	userId := c.PostForm("user_id")
	message := c.PostForm("message")
	if userId == "" || message == "" {
		LOG.Error("http 发送接口 参数错误", zap.String("user_id", userId), zap.String("message", message))
		c.AbortWithStatusJSON(400, map[string]string{"code": "400", "message": "参数错误"})
		return
	}
    // 根据参数发送消息
	err := handle.SendMessage(userId, []byte(message))
	if err != nil {
		if errors.Is(err, errors.New("no this user")) {
			c.JSON(200, map[string]string{"code": "201", "message": "no this user"})
		} else {
			LOG.Error("http 推送消息失败", zap.String("user_id", userId), zap.String("message", message), zap.Error(err))
			c.AbortWithStatusJSON(400, map[string]string{"code": "400", "message": "推送失败"})
			return
		}
	}
	LOG.Error("http 推送消息成功", zap.String("user_id", userId), zap.String("message", message))
    // 反回成功发送
	c.JSON(200, map[string]string{"code": "200", "message": "success"})
}

处理websocket新进连接

func (handle *WsHandle) handle(c *gin.Context) {
	LOG.Info("新进ws客户", zap.String("remote_addr", c.Request.RemoteAddr))
    // 由http升级到websocket
	wsConn, err := handle.upgrade.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		LOG.Info("http协议升级失败:{}", zap.Error(err))
		return
	}
    // 退出连接，关闭连接
	defer func() {
		err2 := wsConn.Close()
		if err2 != nil {
			LOG.Error("connection exit", zap.Error(err2))
		} else {
			LOG.Info("ws退出链接")
		}
	}()
	// 等待客户端发送来的认证消息，即jwt串
	msgAuth := MsgAuth{}
	for {
		mt, msg, err := wsConn.ReadMessage()
		if err != nil {
			LOG.Info("read err while Wait for auth message", zap.Error(err), zap.String("remote_addr", c.Request.RemoteAddr))
			return
		}
        // ping->pong 连接测试
		if mt == websocket.PingMessage {
			err := wsConn.WriteMessage(websocket.PongMessage, []byte("pong"))
			if err != nil {
				return
			}
			continue
		}
        // 非text消息，跳过
		if mt != websocket.TextMessage {
			LOG.Info("不是有用的消息", zap.ByteString("mt", msg))
			continue
		}
        // 成功接收到认证信息
		err = json.Unmarshal(msg, &msgAuth)
		if err != nil {
			LOG.Error("not a init message", zap.String("msg", string(msg)))
			return
		} else {
			LOG.Info("ws 成功获得认证用户信息", zap.Any("msg_auth", msgAuth))
			break
		}
	}
	// 检查token是否有效
	claim, err := CheckJwtToken([]byte(handle.Cfg.JwtSecret), msgAuth.Token)
	if err != nil {
		LOG.Error("ws token验证失败", zap.Any("token", msgAuth.Token))
		return
	}
    // 返回成功认证消息
	err = wsConn.WriteMessage(websocket.TextMessage, []byte("{\n    \"data\": null,\n    \"messageType\": \"登录成功\"\n}\n"))
	if err != nil {
		LOG.Error("send message error", zap.Error(err))
		return
	}
    // 构建用户连接对象 添加到连接集中
	userConn := &UserConn{
		RemoteAddr: c.Request.RemoteAddr,
		UserId:     claim.UserId,
		ClientType: msgAuth.ClientType,
		Token:      msgAuth.Token,
		EnterTime:  time.Now(),
		Conn:       wsConn,
	}
	LOG.Info("用户认证成功,进入等待消息阶段", zap.Any("user_conn", userConn), zap.String("userId", userConn.UserId), zap.String("token", userConn.Token))
	handle.UserConnCollect.AdUserConn(userConn)
    // 退出时的收尾工作
	defer func() {
		handle.UserConnCollect.RemoveUserConn(userConn)
		LOG.Info("userConn defer 退出 已经删除连接信息")
	}()
    // for 循环接收从客户端发送来的消息，如果连接关闭则会主动退出循环
	for {
		mt, message, err := wsConn.ReadMessage()
		if err != nil {
			LOG.Error("read conn return err", zap.Error(err))
			break
		} else {
			LOG.Info("read message from client", zap.String("content", string(message)), zap.Int("type", mt))
            // ping->pong 测试连接有效性
			if mt == websocket.PingMessage {
				err := wsConn.WriteMessage(websocket.PongMessage, []byte("pong"))
				if err != nil {
					break
				}
			}
		}
	}

	LOG.Info("链接退出", zap.String("token", userConn.Token), zap.String("user_id", userConn.UserId))
}

认证消息结构

// MsgAuth
// 用户认证注册消息
// 用于初始化用户信息
type MsgAuth struct {
	ClientType string `json:"client_type"`
	Token      string `json:"token"`
}

从RabbitMQ接收来自服务器的消息，并查找用户推送消息

package main

import (
	jsoniter "github.com/json-iterator/go" // json库 
	"github.com/streadway/amqp" // rabbitmq 客户端
	"go.uber.org/zap" // zap 日志
	"sync" // 锁与同步
)

type RabbitMq struct {
	Mq       *amqp.Connection // rabbitmq 连接
	wsHandle *WsHandle // websocket 服务对象
	Wait     sync.WaitGroup // wait
}

// newRabbitMq 连接函数 连接至rabbit mq
func newRabbitMq(conStr string, wsHandle *WsHandle) RabbitMq {
	conn, err := amqp.Dial(conStr)
	if err != nil {
		panic(err)
	}
	return RabbitMq{Mq: conn, wsHandle: wsHandle}
}

// BeginReceive 开始 消息队列开始接收消息
func (mq *RabbitMq) BeginReceive() {

	LOG.Info("开始接收消息")
	defer mq.Wait.Done()

    // 创建一个rabbit mq通道
	ch, err := mq.Mq.Channel()
	if err != nil {
		panic(err)
	}
	defer ch.Close()
    // 设置这个通道一次只能处理一个消息
    // 在消息未被手动确认时不会有新消息被投递来
	err = ch.Qos(1, 0, false)
	if err != nil {
		panic(err)
	}
    // 交换机
	err = ch.ExchangeDeclare(
		"exchange_test", "direct", true, false, false, false, make(map[string]interface{}),
	)
	if err != nil {
		panic(err)
	}
    // rabbitmq中的队列
	q, err := ch.QueueDeclare(
		"PushClientMessage", // name
		true,                // durable
		false,               // delete when unused
		false,               // exclusive
		false,               // no-Wait
		nil,                 // arguments
	)
	if err != nil {
		panic(err)
	}
    // 绑定交换机与队列 交换机收到消息后会流到队列，再由队列发到消费者
	err = ch.QueueBind("PushClientMessage", "PushClientMessage", "exchange_test", false, nil)
	if err != nil {
		panic(err)
	}
	// 定义一个消费者 一个队列可以有多个消费者，共同消费这个队列中的消息
	msgs, err := ch.Consume(
		q.Name, // queue
		"one",  // consumer 消费者名称
		false,  // auto-ack 不自动确认 手动确认
		false,  // exclusive
		false,  // no-local
		false,  // no-Wait
		nil,    // args
	)
	if err != nil {
		panic(err)
	}
    // 循环接收处理消息
	for {
		select {
		case msg, ok := <-msgs:
			if ok {
                // 获取到从消息队列下发的消息
				msgStr := string(msg.Body)
				LOG.Info("从消息队列中取到消息", zap.String("msg", msgStr))
                // 解释出这个消息要发送给那个用户
				userId := jsoniter.Get(msg.Body, "userId").ToString()
				if userId != "" {
					// 发送消息
					LOG.Info("准备投递消息", zap.String("userId", userId))
                    // 发送消息到用户
					err := mq.wsHandle.SendMessage(userId, msg.Body)
					if err != nil {
						LOG.Error("消息投递到客户端失败", zap.Error(err), zap.Any("msg", msg.Body))              
                        // 未找到用户的情况下，驳回消息不重新投递
						if err.Error() == "no this user" {
							err := msg.Reject(false)
							if err != nil {
								LOG.Error("消息驳回失败", zap.Error(err))
							} else {
								LOG.Info("消息驳回不重新入队", zap.Any("msg", msg.Body))
							}
						} else {
                            // 其他错误驳回后，再次投递处理
							err := msg.Reject(true)
							if err != nil {
								LOG.Error("消息驳回失败", zap.Error(err))
							} else {
								LOG.Info("消息驳回重新入队", zap.Any("msg", msg.Body))
							}
						}
					} else {
                        // 确认此条消息消费成功
						err := msg.Ack(false)
						if err != nil {
							LOG.Error("消息确认失败", zap.Error(err))
						} else {
							LOG.Info("消息确认成功并成功处理", zap.Any("msg", msg.Body))
						}
					}
				} else {
                    // 下发的消息中没有用户id，无法推送，驳回后不重新投递
					err := msg.Reject(false)
					if err != nil {
						LOG.Error("消息驳回失败", zap.Error(err))
					} else {
						LOG.Info("消息驳回不重新入队 未找到客户连接", zap.Error(err))
					}
				}
			} else {
				LOG.Info("消息队列关闭,准备退出")
				return // return 退出for循环
			}
		}
	}
}

main函数的样子

func main() {
	wsHandle := NewWsHandle()
	router := gin.Default()
	router.GET("/", func(c *gin.Context) {
		c.String(http.StatusOK, "Welcome Gin Server")
	})
	router.POST("/push", wsHandle.handleSendMessage) // gin push api 服务器主动推送消息api
	router.GET("/ws", wsHandle.handle) // gin websocket api 客户端ws连接路径

	srv := &http.Server{
		Addr:    ":" + CONFIG.HttpPort,
		Handler: router,
	}
	// 开启http服务
	go func() {
        // 确定是否开启https 并开启http和websocket服务
		if CONFIG.PrivateKeyPath != "" {
			// service connections
			LOG.Info("开始tls http服务")
			if err := srv.ListenAndServeTLS(CONFIG.PublicKeyPath, CONFIG.PrivateKeyPath); err != nil && err != http.ErrServerClosed {
				LOG.Error("ws listen err", zap.Error(err))
				panic(err)
			} else {
				LOG.Info("ws listen exist")
			}
		} else {
			// service connections
			LOG.Info("开始非tls http服务")
			if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
				LOG.Error("ws listen err", zap.Error(err))
				panic(err)
			} else {
				LOG.Info("ws listen exist")
			}
		}
	}()
	// 开启消息队列,开始从消息队列接收消息
	mq := newRabbitMq(CONFIG.RabbitConnStr, wsHandle)
	go func() {
		LOG.Info("开始消息队列接收服务")
		mq.Wait.Add(1)
		mq.BeginReceive()
	}()

	// 监听 ctrl+c 退出信号
	quit := make(chan os.Signal)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit // 等待消息以退出程序
	LOG.Info("begin shutdown server")
	// 关闭消息队列服务
	err2 := mq.Mq.Close()
	if err2 != nil {
		LOG.Info("err on mq close", zap.Error(err2))
	}
	mq.Wait.Wait()
	// 关闭http服务
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	wsHandle.CloseAllForExit() // 关闭所有websocket连接
	if err := srv.Shutdown(ctx); err != nil {
		LOG.Info("http ws server shutdown err", zap.Error(err))
	}
	// 关闭日志服务
	LOG.Info("server shutdown finish")
	CloseZapLog()
	fmt.Println("server exit end")
}

docker

在docker中编译

FROM golang:1.17 as goBuilder
WORKDIR /ws
ADD ./ ./
RUN export CGO_ENABLED=0 && export GOPROXY=https://goproxy.cn,direct && go build -v -o wspush

FROM alpine:3.15
WORKDIR /
COPY --from=goBuilder /ws/wspush ./
COPY --from=goBuilder /ws/config.toml ./
RUN chmod u+x ./wspush
EXPOSE 80
CMD ["./wspush"]

在docker中编译

#!/bin/bash

docker run -it -v "$PWD":/app -v "$PWD"/.cache/gopath:/go -v "$PWD"/.cache/gocache:/root/.cache/ golang:1.17 \
           /bin/sh -c "cd /app && export CGO_ENABLED=0 && export GOPROXY=https://goproxy.cn,direct && go build -v"
if [ "$?" -eq 0 ];then
  echo "build ok"
else
  echo "no"
  exit 1
fi

代码

https://download.csdn.net/download/a34ErxV/81270186https://download.csdn.net/download/a34ErxV/81270186