Working websocket proxy implementation in Golang using httputil package

wsproxy.go

// Starts a websocket proxy
func StartWebSocketProxy() error {
	// Create Reverse Proxy instance
	reverseProxy := httputil.NewSingleHostReverseProxy(&url.URL{})

	reverseProxy.ModifyResponse = func(response *http.Response) error {
		if response.StatusCode == 101 {
			log.Println("Received websocket upgrade response")
		}
		return nil
	}

	// Target websocket host URL
	targetWSEndpoint := "ws://localhost:8080/listener"
	targetWSURL, err := url.Parse(targetWSEndpoint)
	if err != nil {
		return err
	}

	// Setup proxy director
	reverseProxy.Director = func(request *http.Request) {
		// Gives `unsupported protocol scheme "ws"` if not overridden with "http"
		request.URL.Scheme = "http"
		request.URL.Host = targetWSURL.Host
		request.URL.Path = targetWSURL.Path
		request.URL.RawPath = targetWSURL.RawPath
	}

	// Wire Up reverse proxy instance to web server handler route
	mux := http.NewServeMux()
	mux.HandleFunc("/proxy", func(writer http.ResponseWriter, request *http.Request) {
		reverseProxy.ServeHTTP(writer, request)
	})

	// Start Web server
	err = http.ListenAndServe(":8081", mux)
	if err != nil {
		return fmt.Errorf("error listening proxy: %s", err)
	}

	return nil
}

wsproxy_test.go

func TestWebsocketEndpointWithProxy(t *testing.T) {
	wsEndpoint := "ws://localhost:8081/proxy"
	conn, err := websocket.Dial(wsEndpoint, "", "http://localhost")
	assert.Nil(t, err)
	_, err = conn.Write([]byte("text message"))
	assert.Nil(t, err)
	_, err = conn.Write([]byte{2, 5, 100})
	assert.Nil(t, err)
	_, err = conn.Write(append([]byte("text message"), []byte{2, 3}...))
	assert.Nil(t, err)
  
	conn.PayloadType = websocket.BinaryFrame
	_, err = conn.Write(append([]byte("text message"), []byte{2, 3, 4}...))
	assert.Nil(t, err)
}