RF: headers (#49)

Author: syywu

checks/checks.go

@@ -9,6 +9,7 @@ import (
 
 type Checks struct {
 	Carbon     *Carbon
+	Headers    *Headers
 	IpAddress  *Ip
 	LegacyRank *LegacyRank
 	Rank       *Rank
@@ -22,6 +23,7 @@ func NewChecks() *Checks {
 	}
 	return &Checks{
 		Carbon:     NewCarbon(client),
+		Headers:    NewHeaders(client),
 		IpAddress:  NewIp(NewNetIp()),
 		LegacyRank: NewLegacyRank(legacyrank.NewInMemoryStore()),
 		Rank:       NewRank(client),

checks/headers.go

@@ -0,0 +1,36 @@
+package checks
+
+import (
+	"context"
+	"net/http"
+)
+
+type Headers struct {
+	client *http.Client
+}
+
+func NewHeaders(client *http.Client) *Headers {
+	return &Headers{client: client}
+}
+
+func (h *Headers) List(ctx context.Context, url string) (map[string]string, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := h.client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	responseHeaders := make(map[string]string)
+	for k, v := range resp.Header {
+		for _, s := range v {
+			responseHeaders[k] = s
+		}
+	}
+
+	return responseHeaders, nil
+}

checks/headers_test.go

@@ -0,0 +1,28 @@
+package checks
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/xray-web/web-check-api/testutils"
+)
+
+func TestList(t *testing.T) {
+	t.Parallel()
+
+	c := testutils.MockClient(&http.Response{
+		Header: http.Header{
+			"Cache-Control":    {"private, max-age=0"},
+			"X-Xss-Protection": {"0"},
+		},
+	})
+	h := NewHeaders(c)
+
+	actual, err := h.List(context.Background(), "example.com")
+	assert.NoError(t, err)
+
+	assert.Equal(t, "private, max-age=0", actual["Cache-Control"])
+	assert.Equal(t, "0", actual["X-Xss-Protection"])
+}

handlers/headers.go

@@ -2,31 +2,21 @@ package handlers
 
 import (
 	"net/http"
+
+	"github.com/xray-web/web-check-api/checks"
 )
 
-func HandleGetHeaders() http.Handler {
+func HandleGetHeaders(h *checks.Headers) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		rawURL, err := extractURL(r)
 		if err != nil {
 			JSONError(w, ErrMissingURLParameter, http.StatusBadRequest)
 			return
 		}
 
-		resp, err := http.Get(rawURL.String())
+		headers, err := h.List(r.Context(), rawURL.String())
 		if err != nil {
 			JSONError(w, err, http.StatusInternalServerError)
-			return
-		}
-		defer resp.Body.Close()
-
-		// Copying headers from the response
-		headers := make(map[string]interface{})
-		for key, values := range resp.Header {
-			if len(values) > 1 {
-				headers[key] = values
-			} else {
-				headers[key] = values[0]
-			}
 		}
 
 		JSON(w, headers, http.StatusOK)

handlers/headers_test.go

@@ -1,12 +1,12 @@
 package handlers
 
 import (
-	"encoding/json"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/xray-web/web-check-api/checks"
 )
 
 func TestHandleGetHeaders(t *testing.T) {
@@ -15,45 +15,24 @@ func TestHandleGetHeaders(t *testing.T) {
 	t.Run("url parameter is missing", func(t *testing.T) {
 		req := httptest.NewRequest(http.MethodGet, "/headers", nil)
 		rec := httptest.NewRecorder()
-		HandleGetHeaders().ServeHTTP(rec, req)
+		HandleGetHeaders(nil).ServeHTTP(rec, req)
 
 		assert.Equal(t, http.StatusBadRequest, rec.Code)
 		assert.JSONEq(t, `{"error": "missing URL parameter"}`, rec.Body.String())
 	})
 
 	t.Run("invalid url format", func(t *testing.T) {
-		req := httptest.NewRequest(http.MethodGet, "/headers?url=invalid-url", nil)
-		rec := httptest.NewRecorder()
-		HandleGetHeaders().ServeHTTP(rec, req)
-
-		assert.Equal(t, http.StatusInternalServerError, rec.Code)
-	})
-
-	t.Run("valid url", func(t *testing.T) {
 		mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			w.Header().Set("Content-Type", "application/json")
 			w.Header().Set("X-Custom-Header", "value")
 			w.WriteHeader(http.StatusOK)
 		}))
 		defer mockServer.Close()
 
-		req := httptest.NewRequest(http.MethodGet, "/headers?url="+mockServer.URL, nil)
+		req := httptest.NewRequest(http.MethodGet, "/headers?url=invalid-url", nil)
 		rec := httptest.NewRecorder()
-		HandleGetHeaders().ServeHTTP(rec, req)
-
-		assert.Equal(t, http.StatusOK, rec.Code)
-
-		var responseBody map[string]interface{}
-		err := json.Unmarshal(rec.Body.Bytes(), &responseBody)
-		assert.NoError(t, err)
+		HandleGetHeaders(checks.NewHeaders(mockServer.Client())).ServeHTTP(rec, req)
 
-		expectedHeaders := map[string]interface{}{
-			"Content-Type":    "application/json",
-			"X-Custom-Header": "value",
-		}
-
-		for key, expectedValue := range expectedHeaders {
-			assert.Equal(t, expectedValue, responseBody[key])
-		}
+		assert.Equal(t, http.StatusInternalServerError, rec.Code)
 	})
 }

server/server.go

@@ -40,7 +40,7 @@ func (s *Server) routes() {
 	s.mux.Handle("GET /api/dnssec", handlers.HandleDnsSec())
 	s.mux.Handle("GET /api/firewall", handlers.HandleFirewall())
 	s.mux.Handle("GET /api/get-ip", handlers.HandleGetIP(s.checks.IpAddress))
-	s.mux.Handle("GET /api/headers", handlers.HandleGetHeaders())
+	s.mux.Handle("GET /api/headers", handlers.HandleGetHeaders(s.checks.Headers))
 	s.mux.Handle("GET /api/hsts", handlers.HandleHsts())
 	s.mux.Handle("GET /api/http-security", handlers.HandleHttpSecurity())
 	s.mux.Handle("GET /api/legacy-rank", handlers.HandleLegacyRank(s.checks.LegacyRank))