这里websocket package用的是 [gws](https://github.com/lxzan/gws) , 只需要构造标准库的`*http.Request`和`http.ResponseWriter`并实现`http.Hijacker`, 便能复用代码升级`websocket`协议, 完整代码如下:
```go
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())
}
```
有疑问加站长微信联系（非本文作者）
