github.com/simonmittag/ws@v1.1.0-rc.5.0.20210419231947-82b846128245/wsutil/upgrader.go (about)

     1  package wsutil
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"io"
     7  	"io/ioutil"
     8  	"net/http"
     9  
    10  	"github.com/simonmittag/ws"
    11  )
    12  
    13  // DebugUpgrader is a wrapper around ws.Upgrader. It tracks I/O of a
    14  // WebSocket handshake.
    15  //
    16  // Note that it must not be used in production applications that requires
    17  // Upgrade() to be efficient.
    18  type DebugUpgrader struct {
    19  	// Upgrader contains upgrade to WebSocket options.
    20  	Upgrader ws.Upgrader
    21  
    22  	// OnRequest and OnResponse are the callbacks that will be called with the
    23  	// HTTP request and response respectively.
    24  	OnRequest, OnResponse func([]byte)
    25  }
    26  
    27  // Upgrade calls Upgrade() on underlying ws.Upgrader and tracks I/O on conn.
    28  func (d *DebugUpgrader) Upgrade(conn io.ReadWriter) (hs ws.Handshake, err error) {
    29  	var (
    30  		// Take the Reader and Writer parts from conn to be probably replaced
    31  		// below.
    32  		r io.Reader = conn
    33  		w io.Writer = conn
    34  	)
    35  	if onRequest := d.OnRequest; onRequest != nil {
    36  		var buf bytes.Buffer
    37  		// First, we must read the entire request.
    38  		req, err := http.ReadRequest(bufio.NewReader(
    39  			io.TeeReader(conn, &buf),
    40  		))
    41  		if err == nil {
    42  			// Fulfill the buffer with the response body.
    43  			io.Copy(ioutil.Discard, req.Body)
    44  			req.Body.Close()
    45  		}
    46  		onRequest(buf.Bytes())
    47  
    48  		r = io.MultiReader(
    49  			&buf, conn,
    50  		)
    51  	}
    52  
    53  	if onResponse := d.OnResponse; onResponse != nil {
    54  		var buf bytes.Buffer
    55  		// Intercept the response stream written by the Upgrade().
    56  		w = io.MultiWriter(
    57  			conn, &buf,
    58  		)
    59  		defer func() {
    60  			onResponse(buf.Bytes())
    61  		}()
    62  	}
    63  
    64  	return d.Upgrader.Upgrade(struct {
    65  		io.Reader
    66  		io.Writer
    67  	}{r, w})
    68  }