一个简单的 gin websocket 例子

2020-11-04 tech go 5 mins 1829 字

这个例子是我在开发时候写的测试demo，框架基于 Gin。

基本功能非常简单：

当客户端连接上的时候，发送一个自定义的id，
当收到客户端的json信息时，基于不同的返回码，打印日志。

代码包含两个文件，一个是ws的通用文件ws.go，一个是逻辑处理文件kelu.go，在这里做个记录。

依赖包

github.com/gin-gonic/gin
github.com/gorilla/websocket

ws.go

package ws

import (
  "github.com/gorilla/websocket"
  "net/http"
  "github.com/gin-gonic/gin"
  "fmt"
  "github.com/kelu.org/base/pkg/app"
)

// WebSocket 更新用
var keluWsUpgrader = websocket.Upgrader{
  ReadBufferSize:  1024,
  WriteBufferSize: 1024,
  CheckOrigin:     checkOrigin,
}

func checkOrigin(r *http.Request) bool {
  return true
}

// 路由入口
func ApiWs(c *gin.Context) {
  ws, err := keluWsUpgrader.Upgrade(c.Writer, c.Request, nil)
  if err != nil {
    return
  }

  wsid, err := KeluWsOpen(ws)
  fmt.Println("kelu open socket: ",wsid)
  if err != nil {
    return
  }

  defer ws.Close()

  for {
    //读取ws中的数据
    receivedData := app.Response{}

    // メッセージ読み込み
    err = ws.ReadJSON(&receivedData)
    if err != nil {
      fmt.Printf("KeluWs reading ws error: %v", err)
      break
    }

    switch receivedData.Code {
    case http.StatusAccepted:
      // some code ... ...
      break;
    case http.StatusOK:
      // some code ... ...
      break;

    default:
      ws.WriteJSON(KeluWsBad(wsid,receivedData.Data))
    }

  }

  KeluWsDone(wsid)
  return
}

kelu.go

func KeluWsOpen(ws *websocket.Conn) (int64, error) {
  wsid := time.Now().UnixNano()
  welcome := app.Response{http.StatusAccepted, "ok", "ok", map[string]interface{}{"wsid": wsid}}
  err := ws.WriteJSON(welcome)
  if err != nil {
    fmt.Printf("open websocket error while writing message to client: %v", err)
    ws.Close()
    return 0, err
  }

  return wsid, nil
}

func KeluWsBad(wsid int64, data interface{}) *app.Response {
  resultData := app.Response{http.StatusBadRequest, strconv.FormatInt(wsid, 10), "", data}
  return &resultData
}

func KeluWsDone(wsid int64){
  // some code ... ...
}