golang fiber/fasthttp 中使用websocket

这里websocket package用的是 gws , 只需要构造标准库的*http.Request和http.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())
}

本作品采用《CC 协议》，转载必须注明作者和本文链接