net/proxymux: add a listener mux that can run SOCKS and HTTP on a single socket.

Updates #3248

Signed-off-by: David Anderson <danderson@tailscale.com>

net/proxymux/mux_test.go

@@ -0,0 +1,172 @@
+// Copyright (c) 2021 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package proxymux
+
+import (
+	"fmt"
+	"net"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"testing"
+
+	"tailscale.com/net/socks5"
+)
+
+func TestSplitSOCKSAndHTTP(t *testing.T) {
+	s := mkWorld(t)
+	defer s.Close()
+
+	s.checkURL(s.httpClient, false)
+	s.checkURL(s.socksClient, false)
+}
+
+func TestSplitSOCKSAndHTTPCloseSocks(t *testing.T) {
+	s := mkWorld(t)
+	defer s.Close()
+
+	s.socksListener.Close()
+	s.checkURL(s.httpClient, false)
+	s.checkURL(s.socksClient, true)
+}
+
+func TestSplitSOCKSAndHTTPCloseHTTP(t *testing.T) {
+	s := mkWorld(t)
+	defer s.Close()
+
+	s.httpListener.Close()
+	s.checkURL(s.httpClient, true)
+	s.checkURL(s.socksClient, false)
+}
+
+func TestSplitSOCKSAndHTTPCloseBoth(t *testing.T) {
+	s := mkWorld(t)
+	defer s.Close()
+
+	s.httpListener.Close()
+	s.socksListener.Close()
+	s.checkURL(s.httpClient, true)
+	s.checkURL(s.socksClient, true)
+}
+
+type world struct {
+	t *testing.T
+
+	// targetListener/target is the HTTP server the client wants to
+	// reach. It unconditionally responds with HTTP 418 "I'm a
+	// teapot".
+	targetListener net.Listener
+	target         http.Server
+	targetURL      string
+
+	// httpListener/httpProxy is an HTTP proxy that can proxy to
+	// target.
+	httpListener net.Listener
+	httpProxy    http.Server
+
+	// socksListener/socksProxy is a SOCKS5 proxy that can dial
+	// targetListener.
+	socksListener net.Listener
+	socksProxy    *socks5.Server
+
+	// jointListener is the mux that serves both HTTP and SOCKS5
+	// proxying.
+	jointListener net.Listener
+
+	// httpClient and socksClient are HTTP clients configured to proxy
+	// through httpProxy and socksProxy respectively.
+	httpClient  *http.Client
+	socksClient *http.Client
+}
+
+func (s *world) checkURL(c *http.Client, wantErr bool) {
+	s.t.Helper()
+	resp, err := c.Get(s.targetURL)
+	if wantErr {
+		if err == nil {
+			s.t.Errorf("HTTP request succeeded unexpectedly: got HTTP code %d, wanted failure", resp.StatusCode)
+		}
+	} else if err != nil {
+		s.t.Errorf("HTTP request failed: %v", err)
+	} else if c := resp.StatusCode; c != http.StatusTeapot {
+		s.t.Errorf("unexpected status code: got %d, want %d", c, http.StatusTeapot)
+	}
+}
+
+func (s *world) Close() {
+	s.jointListener.Close()
+	s.socksListener.Close()
+	s.httpProxy.Close()
+	s.httpListener.Close()
+	s.target.Close()
+	s.targetListener.Close()
+}
+
+func mkWorld(t *testing.T) (ret *world) {
+	t.Helper()
+
+	ret = &world{
+		t: t,
+	}
+	var err error
+
+	ret.targetListener, err = net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	ret.target = http.Server{
+		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusTeapot)
+		}),
+	}
+	go ret.target.Serve(ret.targetListener)
+	ret.targetURL = fmt.Sprintf("http://%s/", ret.targetListener.Addr().String())
+
+	ret.jointListener, err = net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	ret.socksListener, ret.httpListener = SplitSOCKSAndHTTP(ret.jointListener)
+
+	httpProxy := http.Server{
+		Handler: httputil.NewSingleHostReverseProxy(&url.URL{
+			Scheme: "http",
+			Host:   ret.targetListener.Addr().String(),
+			Path:   "/",
+		}),
+	}
+	go httpProxy.Serve(ret.httpListener)
+
+	socksProxy := socks5.Server{}
+	go socksProxy.Serve(ret.socksListener)
+
+	ret.httpClient = &http.Client{
+		Transport: &http.Transport{
+			Proxy: func(*http.Request) (*url.URL, error) {
+				return &url.URL{
+					Scheme: "http",
+					Host:   ret.jointListener.Addr().String(),
+					Path:   "/",
+				}, nil
+			},
+			DisableKeepAlives: true, // one connection per request
+		},
+	}
+
+	ret.socksClient = &http.Client{
+		Transport: &http.Transport{
+			Proxy: func(*http.Request) (*url.URL, error) {
+				return &url.URL{
+					Scheme: "socks5",
+					Host:   ret.jointListener.Addr().String(),
+					Path:   "/",
+				}, nil
+			},
+			DisableKeepAlives: true, // one connection per request
+		},
+	}
+
+	return ret
+}