websocket proxy

中文

Lightweight websocket proxy library

100 lines of code implements a lightweight websocket proxy library, does not depend on other third-party libraries, and supports ws and wss proxies; if you only need a simple websocket traffic proxy function without any modification to the forwarded content, using this library will be very useful.

Features

Extreme performance, almost no performance loss, very low consumption of cpu and memory
Support websocket handshake phase for management and control
Supports setting header headers (cookie, origin, etc.) in the handshake phase
Support ws, wss proxy

Install

go get github.com/pretty66/websocketproxy

Use

import (
    "github.com/pretty66/websocketproxy"
    "net/http"
)

wp, err := websocketproxy.NewProxy("ws://82.157.123.54:9010/ajaxchattest", func(r *http.Request) error {
    // Permission to verify
    r.Header.Set("Cookie", "----")
    // Source of disguise
    r.Header.Set("Origin", "http://82.157.123.54:9010")
    return nil
})
if err != nil {
    t.Fatal()
}
// proxy path
http.HandleFunc("/wsproxy", wp.Proxy)
http.ListenAndServe(":9696", nil)

Test

Run the test file and start listening on the 127.0.0.1:9696 port, and use the online testing tool http://coolaf.com/tool/chattest to connect to the proxy to test the request response

example

Core-code

func (wp *WebsocketProxy) Proxy(writer http.ResponseWriter, request *http.Request) {
    // Check whether it is a Websocket request
	if strings.ToLower(request.Header.Get("Connection")) != "upgrade" ||
		strings.ToLower(request.Header.Get("Upgrade")) != "websocket" {
		_, _ = writer.Write([]byte(`Must be a websocket request`))
		return
	}
    // Hijack connections
	hijacker, ok := writer.(http.Hijacker)
	if !ok {
		return
	}
	conn, _, err := hijacker.Hijack()
	if err != nil {
		return
	}
	defer conn.Close()
    // Clone request, set destination address path
	req := request.Clone(context.TODO())
	req.URL.Path, req.URL.RawPath, req.RequestURI = wp.defaultPath, wp.defaultPath, wp.defaultPath
	req.Host = wp.remoteAddr
    // Handshake before callback
	if wp.beforeHandshake != nil {
		// Add headers, permission authentication + masquerade sources
		err = wp.beforeHandshake(req)
		if err != nil {
			_, _ = writer.Write([]byte(err.Error()))
			return
		}
	}
    // Determine the protocol and select the dialing process
	var remoteConn net.Conn
	switch wp.scheme {
	case WsScheme:
		remoteConn, err = net.Dial("tcp", wp.remoteAddr)
	case WssScheme:
		remoteConn, err = tls.Dial("tcp", wp.remoteAddr, wp.tlsc)
	}
	if err != nil {
		_, _ = writer.Write([]byte(err.Error()))
		return
	}
	defer remoteConn.Close()
	// Sends a handshake packet to the target WebSocket service
	err = req.Write(remoteConn)
	if err != nil {
		wp.logger.Println("remote write err:", err)
		return
	}
    // Traffic transparent transmission
	errChan := make(chan error, 2)
	copyConn := func(a, b net.Conn) {
		_, err := io.Copy(a, b)
		errChan <- err
	}
	go copyConn(conn, remoteConn) // response
	go copyConn(remoteConn, conn) // request
	select {
	case err = <-errChan:
		if err != nil {
			log.Println(err)
		}
	}
}

Thanks for free JetBrains Open Source license

License

websocketproxy is under the Apache 2.0 license. See the LICENSE directory for details.

reatdoom/websocketproxy