well-goknown/vendor/github.com/gobwas/ws/wsutil/upgrader.go

69 lines
1.5 KiB
Go

package wsutil
import (
"bufio"
"bytes"
"io"
"io/ioutil"
"net/http"
"github.com/gobwas/ws"
)
// DebugUpgrader is a wrapper around ws.Upgrader. It tracks I/O of a
// WebSocket handshake.
//
// Note that it must not be used in production applications that requires
// Upgrade() to be efficient.
type DebugUpgrader struct {
// Upgrader contains upgrade to WebSocket options.
Upgrader ws.Upgrader
// OnRequest and OnResponse are the callbacks that will be called with the
// HTTP request and response respectively.
OnRequest, OnResponse func([]byte)
}
// Upgrade calls Upgrade() on underlying ws.Upgrader and tracks I/O on conn.
func (d *DebugUpgrader) Upgrade(conn io.ReadWriter) (hs ws.Handshake, err error) {
var (
// Take the Reader and Writer parts from conn to be probably replaced
// below.
r io.Reader = conn
w io.Writer = conn
)
if onRequest := d.OnRequest; onRequest != nil {
var buf bytes.Buffer
// First, we must read the entire request.
req, err := http.ReadRequest(bufio.NewReader(
io.TeeReader(conn, &buf),
))
if err == nil {
// Fulfill the buffer with the response body.
io.Copy(ioutil.Discard, req.Body)
req.Body.Close()
}
onRequest(buf.Bytes())
r = io.MultiReader(
&buf, conn,
)
}
if onResponse := d.OnResponse; onResponse != nil {
var buf bytes.Buffer
// Intercept the response stream written by the Upgrade().
w = io.MultiWriter(
conn, &buf,
)
defer func() {
onResponse(buf.Bytes())
}()
}
return d.Upgrader.Upgrade(struct {
io.Reader
io.Writer
}{r, w})
}