proxy_test.go 3.72 KB
Newer Older
1 2 3 4 5 6
package main

import (
	"bytes"
	"fmt"
	"io"
7
	"net"
8 9 10
	"net/http"
	"net/http/httptest"
	"regexp"
11
	"strings"
12
	"testing"
13
	"time"
Jacob Vosmaer's avatar
Jacob Vosmaer committed
14 15 16 17 18

	"gitlab.com/gitlab-org/gitlab-workhorse/internal/badgateway"
	"gitlab.com/gitlab-org/gitlab-workhorse/internal/helper"
	"gitlab.com/gitlab-org/gitlab-workhorse/internal/proxy"
	"gitlab.com/gitlab-org/gitlab-workhorse/internal/testhelper"
19
	"gitlab.com/gitlab-org/gitlab-workhorse/internal/upstream/roundtripper"
20 21
)

22 23
const testVersion = "123"

24
func newProxy(url string, rt http.RoundTripper) *proxy.Proxy {
25
	parsedURL := helper.URLMustParse(url)
26
	if rt == nil {
27
		rt = roundtripper.NewTestBackendRoundTripper(parsedURL)
28
	}
29
	return proxy.NewProxy(parsedURL, testVersion, rt)
30 31
}

32
func TestProxyRequest(t *testing.T) {
33
	ts := testhelper.TestServerWithHandler(regexp.MustCompile(`/url/path\z`), func(w http.ResponseWriter, r *http.Request) {
34 35 36 37 38 39 40 41
		if r.Method != "POST" {
			t.Fatal("Expected POST request")
		}

		if r.Header.Get("Custom-Header") != "test" {
			t.Fatal("Missing custom header")
		}

42 43 44 45
		if h := r.Header.Get("Gitlab-Workhorse"); h != testVersion {
			t.Fatalf("Missing GitLab-Workhorse header: want %q, got %q", testVersion, h)
		}

Jacob Vosmaer's avatar
Jacob Vosmaer committed
46 47
		if h := r.Header.Get("Gitlab-Workhorse-Proxy-Start"); !strings.HasPrefix(h, "1") {
			t.Fatalf("Expect Gitlab-Workhorse-Proxy-Start to start with 1, got %q", h)
48 49
		}

50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
		var body bytes.Buffer
		io.Copy(&body, r.Body)
		if body.String() != "REQUEST" {
			t.Fatal("Expected REQUEST in request body")
		}

		w.Header().Set("Custom-Response-Header", "test")
		w.WriteHeader(202)
		fmt.Fprint(w, "RESPONSE")
	})

	httpRequest, err := http.NewRequest("POST", ts.URL+"/url/path", bytes.NewBufferString("REQUEST"))
	if err != nil {
		t.Fatal(err)
	}
	httpRequest.Header.Set("Custom-Header", "test")

67
	w := httptest.NewRecorder()
68
	newProxy(ts.URL, nil).ServeHTTP(w, httpRequest)
69 70
	testhelper.AssertResponseCode(t, w, 202)
	testhelper.AssertResponseBody(t, w, "RESPONSE")
71

72 73
	if w.Header().Get("Custom-Response-Header") != "test" {
		t.Fatal("Expected custom response header")
74
	}
75
}
76

77 78 79 80 81 82 83 84
func TestProxyError(t *testing.T) {
	httpRequest, err := http.NewRequest("POST", "/url/path", bytes.NewBufferString("REQUEST"))
	if err != nil {
		t.Fatal(err)
	}
	httpRequest.Header.Set("Custom-Header", "test")

	w := httptest.NewRecorder()
85
	newProxy("http://localhost:655575/", nil).ServeHTTP(w, httpRequest)
86
	testhelper.AssertResponseCode(t, w, 502)
87
	testhelper.AssertResponseBodyRegexp(t, w, regexp.MustCompile("dial tcp:.*invalid port.*"))
88 89 90
}

func TestProxyReadTimeout(t *testing.T) {
91
	ts := testhelper.TestServerWithHandler(nil, func(w http.ResponseWriter, r *http.Request) {
92 93 94 95 96 97
		time.Sleep(time.Minute)
	})

	httpRequest, err := http.NewRequest("POST", "http://localhost/url/path", nil)
	if err != nil {
		t.Fatal(err)
98
	}
99

100 101 102 103 104 105 106 107 108
	rt := badgateway.NewRoundTripper(false, &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial,
		TLSHandshakeTimeout:   10 * time.Second,
		ResponseHeaderTimeout: time.Millisecond,
	})
109

110
	p := newProxy(ts.URL, rt)
111
	w := httptest.NewRecorder()
112
	p.ServeHTTP(w, httpRequest)
113
	testhelper.AssertResponseCode(t, w, 502)
114
	testhelper.AssertResponseBody(t, w, "GitLab is not responding")
115 116 117
}

func TestProxyHandlerTimeout(t *testing.T) {
118
	ts := testhelper.TestServerWithHandler(nil,
119 120 121 122 123 124 125 126 127 128 129
		http.TimeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			time.Sleep(time.Second)
		}), time.Millisecond, "Request took too long").ServeHTTP,
	)

	httpRequest, err := http.NewRequest("POST", "http://localhost/url/path", nil)
	if err != nil {
		t.Fatal(err)
	}

	w := httptest.NewRecorder()
130
	newProxy(ts.URL, nil).ServeHTTP(w, httpRequest)
131 132
	testhelper.AssertResponseCode(t, w, 503)
	testhelper.AssertResponseBody(t, w, "Request took too long")
133
}