gitlab-org--gitlab-foss/workhorse/proxy_test.go

135 lines
4.3 KiB
Go

package main
import (
"bytes"
"fmt"
"io/ioutil"
"net"
"net/http"
"net/http/httptest"
"regexp"
"testing"
"time"
"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"
"gitlab.com/gitlab-org/gitlab/workhorse/internal/upstream/roundtripper"
"github.com/stretchr/testify/require"
)
const testVersion = "123"
func newProxy(url string, rt http.RoundTripper, opts ...func(*proxy.Proxy)) *proxy.Proxy {
parsedURL := helper.URLMustParse(url)
if rt == nil {
rt = roundtripper.NewTestBackendRoundTripper(parsedURL)
}
return proxy.NewProxy(parsedURL, testVersion, rt, opts...)
}
func TestProxyRequest(t *testing.T) {
ts := testhelper.TestServerWithHandler(regexp.MustCompile(`/url/path\z`), func(w http.ResponseWriter, r *http.Request) {
require.Equal(t, "POST", r.Method, "method")
require.Equal(t, "test", r.Header.Get("Custom-Header"), "custom header")
require.Equal(t, testVersion, r.Header.Get("Gitlab-Workhorse"), "version header")
require.Regexp(
t,
regexp.MustCompile(`\A1`),
r.Header.Get("Gitlab-Workhorse-Proxy-Start"),
"expect Gitlab-Workhorse-Proxy-Start to start with 1",
)
body, err := ioutil.ReadAll(r.Body)
require.NoError(t, err, "read body")
require.Equal(t, "REQUEST", string(body), "body contents")
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"))
require.NoError(t, err)
httpRequest.Header.Set("Custom-Header", "test")
w := httptest.NewRecorder()
newProxy(ts.URL, nil).ServeHTTP(w, httpRequest)
require.Equal(t, 202, w.Code)
testhelper.RequireResponseBody(t, w, "RESPONSE")
require.Equal(t, "test", w.Header().Get("Custom-Response-Header"), "custom response header")
}
func TestProxyWithCustomHeaders(t *testing.T) {
ts := testhelper.TestServerWithHandler(regexp.MustCompile(`/url/path\z`), func(w http.ResponseWriter, r *http.Request) {
require.Equal(t, "value", r.Header.Get("Custom-Header"), "custom proxy header")
require.Equal(t, testVersion, r.Header.Get("Gitlab-Workhorse"), "version header")
_, err := w.Write([]byte(`ok`))
require.NoError(t, err, "write ok response")
})
httpRequest, err := http.NewRequest("POST", ts.URL+"/url/path", nil)
require.NoError(t, err)
w := httptest.NewRecorder()
testProxy := newProxy(ts.URL, nil, proxy.WithCustomHeaders(map[string]string{"Custom-Header": "value"}))
testProxy.ServeHTTP(w, httpRequest)
testhelper.RequireResponseBody(t, w, "ok")
}
func TestProxyError(t *testing.T) {
httpRequest, err := http.NewRequest("POST", "/url/path", bytes.NewBufferString("REQUEST"))
require.NoError(t, err)
httpRequest.Header.Set("Custom-Header", "test")
w := httptest.NewRecorder()
newProxy("http://localhost:655575/", nil).ServeHTTP(w, httpRequest)
require.Equal(t, 502, w.Code)
require.Regexp(t, regexp.MustCompile("dial tcp:.*invalid port.*"), w.Body.String(), "response body")
}
func TestProxyReadTimeout(t *testing.T) {
ts := testhelper.TestServerWithHandler(nil, func(w http.ResponseWriter, r *http.Request) {
time.Sleep(time.Minute)
})
httpRequest, err := http.NewRequest("POST", "http://localhost/url/path", nil)
require.NoError(t, err)
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,
})
p := newProxy(ts.URL, rt)
w := httptest.NewRecorder()
p.ServeHTTP(w, httpRequest)
require.Equal(t, 502, w.Code)
testhelper.RequireResponseBody(t, w, "GitLab is not responding")
}
func TestProxyHandlerTimeout(t *testing.T) {
ts := testhelper.TestServerWithHandler(nil,
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)
require.NoError(t, err)
w := httptest.NewRecorder()
newProxy(ts.URL, nil).ServeHTTP(w, httpRequest)
require.Equal(t, 503, w.Code)
testhelper.RequireResponseBody(t, w, "Request took too long")
}