.

main_test.go

@@ -460,8 +460,17 @@ func sha1s(data []byte) string {
 	return fmt.Sprintf("%x", sha1.Sum(data))
 }
 
-func download(t *testing.T, url string) (*http.Response, []byte) {
-	resp, err := http.Get(url)
+func download(t *testing.T, url string, h http.Header) (*http.Response, []byte) {
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// copy header to request
+	for k,v := range h {
+		req.Header[k] = v
+	}
+	client := &http.Client{}
+	resp, err := client.Do(req)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -473,65 +482,87 @@ func download(t *testing.T, url string) (*http.Response, []byte) {
 	return resp, body
 }
 
-func TestAllowedBlobDownload(t *testing.T) {
-	// Prepare test server and backend
-	ts := testAuthServer(nil, 200, gitOkBody(t))
-	defer ts.Close()
-	ws := startWorkhorseServer(ts.URL)
-	defer ws.Close()
+// Context for downloading & verifying paths under URL prefix
+type DownloadContext struct {
+	t *testing.T
+	urlPrefix string
+	Header http.Header
+}
+
+func DownloadContextNew(t *testing.T, urlPrefix string) *DownloadContext {
+	h := make(http.Header)
+	return &DownloadContext{t, urlPrefix, h}
+}
+
+func (dl DownloadContext) downloadRaw(path string) (*http.Response, []byte) {
+	return download(dl.t, dl.urlPrefix + path, dl.Header)
+}
 
-	downloadRaw := func (refpath string) (*http.Response, []byte) {
-		return download(t, fmt.Sprintf("%s/%s/raw/%s", ws.URL, testProject, refpath))
+// download `path` and expect content sha1 to be `expectSha1`
+func (dl DownloadContext) ExpectSha1(path, expectSha1 string) {
+	resp, out := dl.downloadRaw(path)
+	if resp.StatusCode != 200 {
+		dl.t.Fatalf("Unexpected status code (expected 200, got %v)", resp.StatusCode)
 	}
-	downloadAndExpectSha1 := func (refpath, expectSha1 string) {
-		_, out := downloadRaw(refpath)
 	outSha1 := sha1s(out)
 	if outSha1 != expectSha1 {
-			t.Fatal("Unexpected content in blob download")
-		}
-	}
-	downloadAndExpect := func (refpath, expect string) {
-		downloadAndExpectSha1(refpath, sha1s([]byte(expect)))
+		dl.t.Fatal("Unexpected content in blob download")
 	}
+}
+
+// download `path` and expect content data to be `expect`
+func (dl DownloadContext) Expect(path, expect string) {
+	dl.ExpectSha1(path, sha1s([]byte(expect)))
+}
 
-	downloadAndExpect("5f923865/README.md", "testme\n======\n\nSample repo for testing gitlab features\n")
-	downloadAndExpectSha1("5f923865/files/ruby/popen.rb", "68990cc20fa74383358797a27967fa2b45d7d8f6")
-	downloadAndExpectSha1("874797c3/files/ruby/popen.rb", "4c266708f2bfd7ca3fed3f7ec74253f92ff3fe73")
-	resp, _ := downloadRaw("master/non-existing-file")
-	if resp.StatusCode != 404 {
-		t.Fatalf("Unexpected status code (expected 404, got %v)", resp.StatusCode)
+// download `path` and expect HTTP status code to be `code`
+func (dl DownloadContext) ExpectCode(path string, code int) {
+	resp, _ := dl.downloadRaw(path)
+	if resp.StatusCode != code {
+		dl.t.Fatalf("Unexpected status code (expected %v, got %v)", code, resp.StatusCode)
 	}
 }
 
+func TestBlobDownload(t *testing.T) {
+	// Prepare test server and "all-ok" auth backend
+	ts := testAuthServer(nil, 200, gitOkBody(t))
+	defer ts.Close()
+	ws := startWorkhorseServer(ts.URL)
+	defer ws.Close()
+	dl := DownloadContextNew(t, fmt.Sprintf("%s/%s/raw", ws.URL, testProject))
+
+	dl.Expect("/5f923865/README.md", "testme\n======\n\nSample repo for testing gitlab features\n")
+	dl.ExpectSha1("/5f923865/README.md", "5f7af35c185a9e5face2f4afb6d7c4f00328d04c")
+	dl.ExpectSha1("/5f923865/files/ruby/popen.rb", "68990cc20fa74383358797a27967fa2b45d7d8f6")
+	dl.ExpectSha1("/874797c3/files/ruby/popen.rb", "4c266708f2bfd7ca3fed3f7ec74253f92ff3fe73")
+	dl.ExpectCode("/master/non-existing-file", 404)
+}
+
 func TestDeniedBlobDownload(t *testing.T) {
-	// Prepare test server and backend
+	// Prepare test server and "all-deny" auth backend
 	ts := testAuthServer(nil, 403, "Access denied")
 	defer ts.Close()
 	ws := startWorkhorseServer(ts.URL)
 	defer ws.Close()
+	dl := DownloadContextNew(t, fmt.Sprintf("%s/%s/raw", ws.URL, testProject))
 
-	downloadAndExpect := func (refpath string, code int) {
-		resp, _ := download(t, fmt.Sprintf("%s/%s/raw/%s", ws.URL, testProject, refpath))
-		if resp.StatusCode != code {
-			t.Fatalf("Unexpected status code (expected %v, got %v)", code, resp.StatusCode)
-		}
-	}
-
-	downloadAndExpect("5f923865/README.md", 403)
-	downloadAndExpect("5f923865/files/ruby/popen.rb", 403)
-	downloadAndExpect("874797c3/files/ruby/popen.rb", 403)
-	downloadAndExpect("master/non-existing-file", 403)
+	dl.ExpectCode("/5f923865/README.md", 403)
+	dl.ExpectCode("/5f923865/files/ruby/popen.rb", 403)
+	dl.ExpectCode("/874797c3/files/ruby/popen.rb", 403)
+	dl.ExpectCode("/master/non-existing-file", 403)
 }
 
 func TestPrivateBlobDownload(t *testing.T) {
+	// Prepare test server and auth backend:
 	// access is ok if token is provided either via query or via header
-	ts := testServerWithHandler(url, func(w, http.ResponseWriter, r *http.Request) {
+	ts := testServerWithHandler(nil, func(w http.ResponseWriter, r *http.Request) {
 		log.Println("UPSTREAM", r.Method, r.URL)
 		token_ok1 := r.URL.Query().Get("aaa_token") == "TOKEN-4AAA"
 		token_ok2 := r.Header.Get("BBB-TOKEN") == "TOKEN-4BBB"
 		if !(token_ok1 || token_ok2) {
 			w.WriteHeader(403)
-			fmt.Fprintf("w", "Access denied")
+			fmt.Fprintf(w, "Access denied")
+			return
 		}
 
 		data, err := json.Marshal(gitOkBody(t))
@@ -541,10 +572,20 @@ func TestPrivateBlobDownload(t *testing.T) {
 
 		w.WriteHeader(200)
 		w.Write(data)
-	}
+	})
 	defer ts.Close()
 	ws := startWorkhorseServer(ts.URL)
 	defer ws.Close()
-
-	download(t, "%s/%s/raw/5f923865/README.md
+	dl := DownloadContextNew(t, fmt.Sprintf("%s/%s/raw", ws.URL, testProject))
+
+	dl.ExpectCode("/5f923865/README.md", 403)
+	dl.ExpectCode("/5f923865/README.md?bbb_token=TOKEN-4BBB", 403)
+	dl.ExpectCode("/5f923865/README.md?aaa_token=TOKEN-4AAA", 200)
+	dl.ExpectSha1("/5f923865/README.md?aaa_token=TOKEN-4AAA", "5f7af35c185a9e5face2f4afb6d7c4f00328d04c")
+
+	dl.Header.Add("AAA-TOKEN", "TOKEN-4AAA")
+	dl.ExpectCode("/5f923865/README.md", 403)
+	dl.Header.Add("BBB-TOKEN", "TOKEN-4BBB")
+	dl.ExpectCode("/5f923865/README.md", 200)
+	dl.ExpectSha1("/5f923865/README.md", "5f7af35c185a9e5face2f4afb6d7c4f00328d04c")
 }