在 fiber/fasthttp 中使用websocket

lxzan · 2023-04-25 16:23:49 · 1891 次点击 · 大约8小时之前 开始浏览    置顶
这是一个创建于 2023-04-25 16:23:49 的主题,其中的信息可能已经有所发展或是发生改变。

这里websocket package用的是 gws , 只需要构造标准库的*http.Requesthttp.ResponseWriter并实现http.Hijacker, 便能复用代码升级websocket协议, 完整代码如下:

package main

import (
    "bufio"
    fiber "github.com/gofiber/fiber/v2"
    "github.com/lxzan/gws"
    "log"
    "net"
    "net/http"
    "net/url"
)

func main() {
    upgrader := gws.NewUpgrader(new(WebSocket), nil)

    app := fiber.New(fiber.Config{})

    app.Get("/connect", func(c *fiber.Ctx) error {
        r, err := getRequest(c)
        if err != nil {
            return err
        }
        c.Context().HijackSetNoResponse(true)
        c.Context().Hijack(func(conn net.Conn) {
            w := &responseWriter{
                conn: conn,
                brw:  bufio.NewReadWriter(bufio.NewReader(conn), nil), // gws没有使用到bufio.Writer
            }
            socket, err := upgrader.Accept(w, r)
            if err != nil {
                log.Printf(err.Error())
                return
            }
            socket.Listen()
        })
        return nil
    })

    app.Listen(":3003")
}

func getRequest(c *fiber.Ctx) (*http.Request, error) {
    URL, err := url.Parse(string(c.Request().URI().FullURI()))
    if err != nil {
        return nil, err
    }

    var r = &http.Request{
        Method: c.Method(),
        Header: http.Header{},
        URL:    URL,
    }
    c.Request().Header.VisitAll(func(key, value []byte) {
        r.Header.Set(string(key), string(value))
    })
    return r, nil
}

type responseWriter struct {
    conn net.Conn
    brw  *bufio.ReadWriter
}

func (r *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    return r.conn, r.brw, nil
}

func (r *responseWriter) Header() http.Header { return http.Header{} }

func (r *responseWriter) Write(bytes []byte) (int, error) { return 0, nil }

func (r *responseWriter) WriteHeader(statusCode int) { return }

type WebSocket struct {
    gws.BuiltinEventHandler
}

func (w WebSocket) OnMessage(socket *gws.Conn, message *gws.Message) {
    defer message.Close()
    socket.WriteMessage(message.Opcode, message.Bytes())
}

有疑问加站长微信联系(非本文作者)

入群交流(和以上内容无关):加入Go大咖交流群,或添加微信:liuxiaoyan-s 备注:入群;或加QQ群:692541889

1891 次点击  
加入收藏 微博
暂无回复
添加一条新回复 (您需要 登录 后才能回复 没有账号 ?)
  • 请尽量让自己的回复能够对别人有帮助
  • 支持 Markdown 格式, **粗体**、~~删除线~~、`单行代码`
  • 支持 @ 本站用户;支持表情(输入 : 提示),见 Emoji cheat sheet
  • 图片支持拖拽、截图粘贴等方式上传