Add port reuse functionality

main.go

@@ -141,6 +141,10 @@ var serverHelp = `
 
     --reverse, Allow clients to specify reverse port forwarding remotes
     in addition to normal remotes.
+
+    --reuseport, Allow the server to reuse the listening port. This still 
+    allows the underlying service to be reached (which you need to specify
+    through --backend).
 
     --tls-key, Enables TLS and provides optional path to a PEM-encoded
     TLS private key. When this flag is set, you must also set --tls-cert,
@@ -177,6 +181,7 @@ func server(args []string) {
 	flags.StringVar(&config.Proxy, "backend", "", "")
 	flags.BoolVar(&config.Socks5, "socks5", false, "")
 	flags.BoolVar(&config.Reverse, "reverse", false, "")
+	flags.BoolVar(&config.ReusePort, "reuseport", false, "")
 	flags.StringVar(&config.TLS.Key, "tls-key", "", "")
 	flags.StringVar(&config.TLS.Cert, "tls-cert", "", "")
 	flags.Var(multiFlag{&config.TLS.Domains}, "tls-domain", "")

server/server.go

@@ -28,6 +28,7 @@ type Config struct {
 	Proxy     string
 	Socks5    bool
 	Reverse   bool
+	ReusePort bool
 	KeepAlive time.Duration
 	TLS       TLSConfig
 }
@@ -101,12 +102,19 @@ func NewServer(c *Config) (*Server, error) {
 			return nil, server.Errorf("Missing protocol (%s)", u)
 		}
 		server.reverseProxy = httputil.NewSingleHostReverseProxy(u)
+
 		//always use proxy host
 		server.reverseProxy.Director = func(r *http.Request) {
-			//enforce origin, keep path
-			r.URL.Scheme = u.Scheme
-			r.URL.Host = u.Host
-			r.Host = u.Host
+			if c.ReusePort {
+				// inherit origin
+				r.URL.Host = u.Host // r.URL.Host is our destination - r.Host contains the actual host header
+				r.URL.Scheme = "http" // we only do HTTP backends
+			} else {
+				//enforce origin, keep path
+				r.URL.Scheme = u.Scheme
+				r.URL.Host = u.Host
+				r.Host = u.Host
+			}
 		}
 	}
 	//print when reverse tunnelling is enabled

server/server_handler.go

@@ -16,11 +16,13 @@ import (
 
 // handleClientHandler is the main http websocket handler for the chisel server
 func (s *Server) handleClientHandler(w http.ResponseWriter, r *http.Request) {
-	//websockets upgrade AND has chisel prefix
-	upgrade := strings.ToLower(r.Header.Get("Upgrade"))
+
 	protocol := r.Header.Get("Sec-WebSocket-Protocol")
-	if upgrade == "websocket" && strings.HasPrefix(protocol, "chisel-") {
+	if strings.HasPrefix(protocol, "chisel-") {
 		if protocol == chshare.ProtocolVersion {
+			// force add websocket headers
+			r.Header.Add("Connection", "upgrade")
+			r.Header.Add("Upgrade", "websocket")
 			s.handleWebsocket(w, r)
 			return
 		}

server/server_listen.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/jpillora/chisel/share/settings"
 	"golang.org/x/crypto/acme/autocert"
+	reuseport "github.com/libp2p/go-reuseport"
 )
 
 //TLSConfig enables configures TLS
@@ -43,11 +44,29 @@ func (s *Server) listener(host, port string) (net.Listener, error) {
 			extra = " (WARNING: LetsEncrypt will attempt to connect to your domain on port 443)"
 		}
 	}
-	//tcp listen
-	l, err := net.Listen("tcp", host+":"+port)
+
+	var l net.Listener
+	var err error
+
+	if s.config.ReusePort {
+		if tlsConf != nil {
+			return nil, errors.New("Port reuse/hijack only works for HTTP connections.")
+		}
+
+		if s.config.Proxy == "" {
+			return nil, errors.New("For port reuse, you need to specify backend to redirect non-chisel requests to (http://127.0.0.1 perhaps?)")
+		}
+
+		s.Infof("Attempting port reuse")
+		l, err = reuseport.Listen("tcp", host+":"+port)
+	} else {
+		l, err = net.Listen("tcp", host+":"+port)
+	}
+
 	if err != nil {
 		return nil, err
 	}
+
 	//optionally wrap in tls
 	proto := "http"
 	if tlsConf != nil {