tailscale/net/proxymux/mux_test.go

173 lines
3.9 KiB
Go
Raw Permalink Normal View History

// 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
}