Websocket proxy sniffer

Guide to write a websocket proxy sniffer in golang

Add a sniffer to your websocket proxy

Since Go 1.12 was released, setting up a reverse websocket proxy is as easy as writing a few lines of code.

Thanks to httputil!

func main() {
	u, err := url.Parse("http://backend:8080")
	if err != nil {
		log.Fatal(err)
	}
	proxy := httputil.NewSingleHostReverseProxy(u)
	http.ListenAndServe(":8080", proxy)
}

Internally this works with a hijacked http.ResponseWriter before switching from http to websocket protocol. If you are curious search for http.Hijacker here src/net/http/httputil/reverseproxy.go

So how could we forward websocket requests and at the same time add a sniffer to that connection?

Sniffer API

Let's start by defining our API as a combination of 2 elements:

  • A callback that will be called every time a request has been hijacked. It will return the request that has been hijacked and two readers corresponding to data received and sent on the hijacked connection.
  • A Sniffer function that will wrap an existing http.Handler with our callback.
type OnHijacked func(r *http.Request, in, out io.Reader)
func Sniffer(h http.Handler, callback OnHijacked) http.Handler 

Now that our API is defined is time to build the code that satisfies it.

Tee-ing a net.Conn

First we will create a net.Conn implementation that will forward any reads or writes to a pair of io.Writer. Same concept that is implemented by io.TeeReader applied to a net.Conn.

func TeeConn(conn net.Conn, in, out io.Writer) net.Conn {
	return &teeConn{
		Conn:   conn,
		reader: io.TeeReader(conn, in),
		writer: io.MultiWriter(conn, out),
	}
}

type teeConn struct {
	net.Conn
	reader io.Reader
	writer io.Writer
}

func (c *teeConn) Read(p []byte) (n int, err error) {
	return c.reader.Read(p)
}

func (c *teeConn) Write(p []byte) (n int, err error) {
	return c.writer.Write(p)
}

Invoke our OnHijacked callback

The next step will be to create a wrapper around an http.ResponseWriter that will invoke our OnHijacked callback whenever a Hijack() is successfully done.

func CallbackHijacker(w http.ResponseWriter, r *http.Request,
	cb OnHijacked) http.ResponseWriter {
	if h, ok := w.(http.Hijacker); ok {
		w = &callbackHijacker{
			ResponseWriter: w,
			hijacker:       h,
			request:        r,
			callback:       cb,
		}
	}
	return w
}

type callbackHijacker struct {
	http.ResponseWriter
	hijacker http.Hijacker
	request  *http.Request
	callback OnHijacked
}

In order to do this we have to override Hijack() with our custom version.

Time to use our awesome TeeConn function! But, wait a moment… our callback works with io.Readers

type OnHijacked func() (r http.Request, in, out io.Reader)

and TeeConn works with io.Writer 😟

func TeeConn(conn net.Conn, in, out io.Writer) net.Conn

No problem at all! This is a perfect task for io.Pipe(). As the doc says, it can be used to connect code expecting an io.Reader with code expecting an io.Writer.

func (h *callbackHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	conn, buf, err := h.hijacker.Hijack()
	if err != nil {
		return conn, buf, err
	}
	// use `io.Pipe` to connect code expecting an `io.Reader` with code
	// expecting an `io.Writer`
	rIn, wIn := io.Pipe()
	rOut, wOut := io.Pipe()

	// invoke callback
	h.callback(h.request, rIn, rOut)

	// return wrapped conn
	return TeeConn(conn, wIn, wOut), buf, nil
}

Our Sniffer wrapper

Finally we will create a wrapper around http.Handler (our initially defined Sniffer) that will invoke CallbackHijacker every time ServeHTTP() is called.

func Sniffer(h http.Handler, callback OnHijacked) http.Handler {
	return &sniffer{
		handler:  h,
		callback: callback,
	}
}

type sniffer struct {
	handler  http.Handler
	callback OnHijacked
}

func (s *sniffer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w = CallbackHijacker(w, r, s.callback)
	s.handler.ServeHTTP(w, r)
}

Final look

This is how our sniffer will look like applied to the reverse proxy and connected to echo.websocket.org.

We've added Director setting to rewrite requests scheme and host before sending them.

func main() {
	u, err := url.Parse("http://echo.websocket.org")
	if err != nil {
		log.Fatal(err)
	}
	proxy := httputil.NewSingleHostReverseProxy(u)
	proxy.Director = func(r *http.Request) {
		r.URL.Scheme = u.Scheme
		r.URL.Host = u.Host
		r.Host = u.Host
	}
	handler := Sniffer(proxy, func(r *http.Request, in, out io.Reader) {
		go readLoop(r, in, "<")
		go readLoop(r, out, ">")
	})
	log.Println("listening on :8080")
	http.ListenAndServe(":8080", handler)
}

func readLoop(req *http.Request, r io.Reader, dir string) {
	data := make([]byte, 1024)
	for {
		n, err := r.Read(data)
		if err != nil {
			log.Println(err)
			return
		}
		log.Printf("%s %s: %x\n", dir, req.RemoteAddr, data[:n])
	}
}

Source code available in this repository: https://github.com/igolaizola/websocket-proxy-sniffer

Avatar
Iñigo Garcia Olaizola
Trigopher

software engineer, gopher, triathlete, husband & father.