Merge branch 'sh-custom-error-override' into 'master'

Support custom error page override via X-GitLab-Custom-Error

See merge request gitlab-org/gitlab-workhorse!300

internal/staticpages/error_pages.go

@@ -53,7 +53,9 @@ func (s *errorPageResponseWriter) WriteHeader(status int) {
 
 	s.status = status
 
-	if 400 <= s.status && s.status <= 599 {
+	if 400 <= s.status && s.status <= 599 &&
+		s.rw.Header().Get("X-GitLab-Custom-Error") == "" &&
+		s.rw.Header().Get("Content-Type") != "application/json" {
 		errorPageFile := filepath.Join(s.path, fmt.Sprintf("%d.html", s.status))
 
 		// check if custom error page exists, serve this page instead

internal/staticpages/error_pages_test.go

@@ -78,3 +78,68 @@ func TestIfErrorPageIsIgnoredInDevelopment(t *testing.T) {
 	testhelper.AssertResponseCode(t, w, 500)
 	testhelper.AssertResponseBody(t, w, serverError)
 }
+
+func TestIfErrorPageIsIgnoredIfCustomError(t *testing.T) {
+	dir, err := ioutil.TempDir("", "error_page")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(dir)
+
+	errorPage := "ERROR"
+	ioutil.WriteFile(filepath.Join(dir, "500.html"), []byte(errorPage), 0600)
+
+	w := httptest.NewRecorder()
+	serverError := "Interesting Server Error"
+	h := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Add("X-GitLab-Custom-Error", "1")
+		w.WriteHeader(500)
+		fmt.Fprint(w, serverError)
+	})
+	st := &Static{dir}
+	st.ErrorPagesUnless(false, h).ServeHTTP(w, nil)
+	w.Flush()
+	testhelper.AssertResponseCode(t, w, 500)
+	testhelper.AssertResponseBody(t, w, serverError)
+}
+
+func TestErrorPageInterceptedByContentType(t *testing.T) {
+	testCases := []struct {
+		contentType string
+		intercepted bool
+	}{
+		{contentType: "application/json", intercepted: false},
+		{contentType: "text/plain", intercepted: true},
+		{contentType: "text/html", intercepted: true},
+		{contentType: "", intercepted: true},
+	}
+
+	for _, tc := range testCases {
+		dir, err := ioutil.TempDir("", "error_page")
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer os.RemoveAll(dir)
+
+		errorPage := "ERROR"
+		ioutil.WriteFile(filepath.Join(dir, "500.html"), []byte(errorPage), 0600)
+
+		w := httptest.NewRecorder()
+		serverError := "Interesting Server Error"
+		h := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			w.Header().Add("Content-Type", tc.contentType)
+			w.WriteHeader(500)
+			fmt.Fprint(w, serverError)
+		})
+		st := &Static{dir}
+		st.ErrorPagesUnless(false, h).ServeHTTP(w, nil)
+		w.Flush()
+		testhelper.AssertResponseCode(t, w, 500)
+
+		if tc.intercepted {
+			testhelper.AssertResponseBody(t, w, errorPage)
+		} else {
+			testhelper.AssertResponseBody(t, w, serverError)
+		}
+	}
+}