increase coverage on serve

nfx · Jul 31, 2022 · 2001ea4 · 2001ea4
1 parent b56e80c
commit 2001ea4
Show file tree

Hide file tree

Showing 5 changed files with 206 additions and 6 deletions.
diff --git a/Makefile b/Makefile
@@ -26,7 +26,7 @@ vendor:
 	go mod vendor
 
 test:
-	go test ./... -coverprofile=coverage.txt -timeout=1m
+	go test ./... -coverprofile=coverage.txt -timeout=30s
 
 coverage: test
 	go tool cover -html=coverage.txt

diff --git a/serve/http_proxy.go b/serve/http_proxy.go
@@ -118,12 +118,12 @@ func (srv *HttpProxyServer) handleConnect(rw http.ResponseWriter, r *http.Reques
 		rw.WriteHeader(501)
 		return
 	}
-	rw.WriteHeader(200)
 	src, _, err := hijacker.Hijack()
 	if err != nil {
 		http.Error(rw, err.Error(), 471)
 		return
 	}
+	rw.WriteHeader(200)
 	// TODO: figure out context propagation
 	go func() {
 		ssl, err := srv.handleHandshake(src, r.URL.Host)

diff --git a/serve/http_proxy_test.go b/serve/http_proxy_test.go
@@ -0,0 +1,167 @@
+package serve
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHttpProxyServer_ListenAndServe_NoConf(t *testing.T) {
+	err := (&HttpProxyServer{}).ListenAndServe()
+	assert.EqualError(t, err, "listener is not configured")
+}
+
+type staticRT func(*http.Request) (*http.Response, error)
+
+func (s staticRT) RoundTrip(r *http.Request) (*http.Response, error) {
+	return s(r)
+}
+
+var dummy = &http.Request{
+	URL: &url.URL{
+		Scheme: "http",
+		Host:   "localhost",
+	},
+}
+
+func TestHttpProxyServer_handleSimpleHttp_RoundTripFailed(t *testing.T) {
+	w := httptest.NewRecorder()
+
+	(&HttpProxyServer{
+		transport: staticRT(func(r *http.Request) (*http.Response, error) {
+			return nil, fmt.Errorf("nope")
+		}),
+	}).handleSimpleHttp(w, dummy)
+
+	resp := w.Result()
+	raw, _ := io.ReadAll(resp.Body)
+
+	assert.Equal(t, 470, resp.StatusCode)
+	assert.Equal(t, "nope\n", string(raw))
+}
+
+func TestHttpProxyServer_handleSimpleHttp_ForwardEmptyBody(t *testing.T) {
+	w := httptest.NewRecorder()
+
+	(&HttpProxyServer{
+		transport: staticRT(func(r *http.Request) (*http.Response, error) {
+			return &http.Response{
+				StatusCode: 201,
+				Header: http.Header{
+					"X-Test": []string{"yes"},
+				},
+			}, nil
+		}),
+	}).handleSimpleHttp(w, dummy)
+
+	resp := w.Result()
+	raw, _ := io.ReadAll(resp.Body)
+
+	// empty body is not nil, but zero-byte
+	assert.Equal(t, "", string(raw))
+	assert.Equal(t, 201, resp.StatusCode)
+	assert.Equal(t, "yes", resp.Header.Get("X-Test"))
+}
+
+type failingReader string
+
+func (f failingReader) Read(p []byte) (n int, err error) {
+	return 0, fmt.Errorf("%s", f)
+}
+
+func TestHttpProxyServer_handleSimpleHttp_FailingReader(t *testing.T) {
+	w := httptest.NewRecorder()
+
+	(&HttpProxyServer{
+		transport: staticRT(func(r *http.Request) (*http.Response, error) {
+			return &http.Response{
+				StatusCode: 201,
+
+				// TODO: make some better reporting to client
+				Body: io.NopCloser(failingReader("nope")),
+			}, nil
+		}),
+	}).handleSimpleHttp(w, dummy)
+
+	resp := w.Result()
+	raw, _ := io.ReadAll(resp.Body)
+
+	assert.Equal(t, "", string(raw))
+	assert.Equal(t, 201, resp.StatusCode)
+}
+
+func TestHttpProxyServer_handleConnect_NoHijack(t *testing.T) {
+	w := httptest.NewRecorder()
+
+	(&HttpProxyServer{}).handleConnect(w, dummy)
+
+	resp := w.Result()
+	assert.Equal(t, 501, resp.StatusCode)
+}
+
+type hijackable struct {
+	*httptest.ResponseRecorder
+	*httptest.Server
+
+	server    net.Conn
+	client    net.Conn
+	hijackErr error
+}
+
+func newHijackable() *hijackable {
+	server, client := net.Pipe()
+	srv := httptest.NewTLSServer(http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(210)
+		}))
+	return &hijackable{
+		ResponseRecorder: httptest.NewRecorder(),
+		Server:           srv,
+		server:           server,
+		client:           client,
+	}
+}
+
+func (w hijackable) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if w.hijackErr != nil {
+		return nil, nil, w.hijackErr
+	}
+	conn, err := net.Dial("tcp", w.Listener.Addr().String())
+	return conn, nil, err
+}
+
+func TestHttpProxyServer_handleConnect_CannotHijack(t *testing.T) {
+	w := hijackable{
+		ResponseRecorder: httptest.NewRecorder(),
+		hijackErr:        fmt.Errorf("nope"),
+	}
+
+	(&HttpProxyServer{}).handleConnect(w, dummy)
+
+	resp := w.Result()
+	assert.Equal(t, 471, resp.StatusCode)
+}
+
+func TestHttpProxyServer_handleConnect_Hijacked(t *testing.T) {
+	t.Skip("unfinished, pick up later")
+	w := newHijackable()
+
+	ca, _ := NewCA()
+	(&HttpProxyServer{
+		ca: ca,
+	}).handleConnect(w, dummy)
+
+	resp := w.Result()
+	assert.Equal(t, 200, resp.StatusCode)
+
+	var raw []byte
+	_, err := w.server.Read(raw)
+	assert.NoError(t, err)
+}
diff --git a/serve/mitm_proxy.go b/serve/mitm_proxy.go
@@ -2,6 +2,7 @@ package serve
 
 import (
 	"fmt"
+	"net"
 	"net/http"
 	"net/url"
 	"time"
@@ -27,9 +28,12 @@ func NewMitmProxyServer(pool *pool.Pool, ca certWrapper) *MitmProxyServer {
 	}
 }
 
+// package-private variable, to simplify tests
+var mitmDefaultAddr = "localhost:8090"
+
 func (mps *MitmProxyServer) Configure(c app.Config) error {
 	// TODO: make sure about private interfaces only!!!
-	mps.Addr = c.StrOr("addr", "localhost:8090")
+	mps.Addr = c.StrOr("addr", mitmDefaultAddr)
 	mps.ReadTimeout = c.DurOr("read_timeout", 15*time.Second)
 	mps.IdleTimeout = c.DurOr("idle_timeout", 15*time.Second)
 	mps.WriteTimeout = c.DurOr("write_timeout", 15*time.Second)
@@ -39,12 +43,16 @@ func (mps *MitmProxyServer) Configure(c app.Config) error {
 
 func (mps *MitmProxyServer) transportProxy() func(*http.Request) (*url.URL, error) {
 	return func(r *http.Request) (*url.URL, error) {
-		if mps.Addr == "" {
-			return nil, fmt.Errorf("mitm is not configured")
+		if mps.listener == nil {
+			return nil, fmt.Errorf("mitm is not initialized")
+		}
+		addr, ok := mps.listener.Addr().(*net.TCPAddr)
+		if !ok {
+			return nil, fmt.Errorf("not a tcp listener: %v", mps.listener.Addr())
 		}
 		return &url.URL{
 			Scheme: "http",
-			Host:   mps.Addr,
+			Host:   addr.String(),
 		}, nil
 	}
 }

diff --git a/serve/mitm_proxy_test.go b/serve/mitm_proxy_test.go
@@ -2,6 +2,8 @@ package serve
 
 import (
 	"crypto/tls"
+	"fmt"
+	"net"
 	"net/http"
 	"testing"
 	"time"
@@ -14,6 +16,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func init() {
+	// listen mitm on a random port each time
+	mitmDefaultAddr = "localhost:0"
+}
+
 func TestMitmWorksForHttp(t *testing.T) {
 	ca, err := NewCA()
 	assert.NoError(t, err)
@@ -75,3 +82,21 @@ func TestMitmWorksForHttps(t *testing.T) {
 	// TODO: make it working properly
 	assert.Equal(t, 200, res.StatusCode)
 }
+
+func TestMitmTransportProxyNoInit(t *testing.T) {
+	_, err := (&MitmProxyServer{}).transportProxy()(nil)
+	assert.EqualError(t, err, "mitm is not initialized")
+}
+
+func TestMitmTransportProxyWrongListener(t *testing.T) {
+	tmp := fmt.Sprintf("%s/x", t.TempDir())
+	conn, err := net.Listen("unix", tmp)
+	assert.NoError(t, err)
+	defer conn.Close()
+	_, err = (&MitmProxyServer{
+		HttpProxyServer: HttpProxyServer{
+			listener: conn,
+		},
+	}).transportProxy()(nil)
+	assert.EqualError(t, err, fmt.Sprintf("not a tcp listener: %s", tmp))
+}