星河

醉后不知天在水,满船清梦压星河

Golang 简易WS服务 - 客户端

mingzaily / 2023-02-21


客户端代码较为简单


type SocketClient struct {
	host        string
	isClosed    chan bool
	log         *logger.Log // 自定义的log包
}

type wsMessage struct {
	Type int         `json:"type"`
	Data interface{} `json:"data"`
}

func NewSocketClient(configPath, logPath, host string) *SocketClient {
	w := io.MultiWriter(os.Stdout)
	if logPath != "" {
		f, err := util.OpenFile(logPath)
		if err != nil {
			panic(err)
		}
		w = io.MultiWriter(os.Stdout, f)
	}

	return &SocketClient{
		host:        host,
		isClosed:    make(chan bool),
		log:         logger.New(w, logger.LINFO, log.LstdFlags|log.Lmsgprefix),
	}
}

func (s *SocketClient) Start() {
	// 设置在线
	ClientStatus = ClientStatusOnline

	// 连接服务器
	u := url.URL{Scheme: "ws", Host: s.host, Path: "/socket"}
	s.log.Info("正在连接到", u.String())

	conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		s.log.ErrorF("连接到服务器错误: %v", err)
	}
	defer conn.Close()

	s.log.Info("已连接到服务器")

	interrupt := make(chan os.Signal, 1)

	// 监听中断信号
	signal.Notify(interrupt, syscall.SIGKILL, syscall.SIGINT, syscall.SIGTERM, os.Kill)
	// 心跳包
	go s.pingHandler(conn)
	// 接收消息
	go s.receiveHandler(conn)

	for {
		select {
		case <-interrupt:
			s.log.Info("收到SIGINT中断信号,正在关闭ws连接。。。")
			_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
			select {
			case <-s.isClosed:
				s.log.Info("WS链接已关闭,退出中。。。")
			case <-time.After(time.Duration(1) * time.Second):
				s.log.Info("关闭ws链接超时,退出中。。。")
			}
			return
		case <-s.isClosed:
			s.log.Info("ws链接已关闭,退出中。。。")
			return
		}
	}
}

func (s *SocketClient) pingHandler(conn *websocket.Conn) {
	ticker := time.NewTicker(pingTime)
	for {
		select {
		case <-ticker.C:
			status := strconv.FormatInt(int64(ClientStatus), 10)
			err := conn.WriteMessage(websocket.PingMessage, []byte(status))
			if err != nil {
				s.log.ErrorF("发送心跳包错误: %v", err)
				return
			}
		case <-s.isClosed:
			return
		}
	}
}

func (s *SocketClient) receiveHandler(ws *websocket.Conn) {
	defer close(s.isClosed)
	for {
		messageType, message, err := ws.ReadMessage()
		if err != nil {
			s.log.ErrorF("读取消息 %v", err)
			return
		}
		switch messageType {
		case websocket.TextMessage:
			var textMessage *wsMessage

			err = json.Unmarshal(message, &textMessage)
			if err != nil {
				s.log.ErrorF("消息格式错误: %v", err)
				break
			}

			switch textMessage.Type {
			        ...
			}
		}
	}
}