Merge pull request #59 from grafana/add-request-headers

Add request headers

pkg/client/client.go

@@ -14,31 +14,50 @@ import (
 	"github.com/grafana/k6build/pkg/api"
 )
 
+const (
+	defaultAuthType = "Bearer"
+)
+
 var (
 	ErrBuildFailed   = errors.New("build failed")   //nolint:revive
 	ErrRequestFailed = errors.New("request failed") //nolint:revive
 )
 
 // BuildServiceClientConfig defines the configuration for accessing a remote build service
 type BuildServiceClientConfig struct {
+	// URL to build service
 	URL string
+	// Authorization credentials passed in the Authorization: <type> <credentials> header
+	// See AuthorizationType
+	Authorization string
+	// AuthorizationType type of credentials in the Authorization: <type> <credentials> header
+	// For example, "Bearer", "Token", "Basic". Defaults to "Bearer"
+	AuthorizationType string
+	// Headers custom request headers
+	Headers map[string]string
 }
 
 // NewBuildServiceClient returns a new client for a remote build service
 func NewBuildServiceClient(config BuildServiceClientConfig) (k6build.BuildService, error) {
 	return &BuildClient{
-		srv: config.URL,
+		srvURL:   config.URL,
+		auth:     config.Authorization,
+		authType: config.AuthorizationType,
+		headers:  config.Headers,
 	}, nil
 }
 
 // BuildClient defines a client of a build service
 type BuildClient struct {
-	srv string
+	srvURL   string
+	authType string
+	auth     string
+	headers  map[string]string
 }
 
 // Build request building an artidact to a build service
 func (r *BuildClient) Build(
-	_ context.Context,
+	ctx context.Context,
 	platform string,
 	k6Constrains string,
 	deps []k6build.Dependency,
@@ -54,12 +73,33 @@ func (r *BuildClient) Build(
 		return k6build.Artifact{}, fmt.Errorf("%w: %w", ErrRequestFailed, err)
 	}
 
-	url, err := url.Parse(r.srv)
+	url, err := url.Parse(r.srvURL)
 	if err != nil {
 		return k6build.Artifact{}, fmt.Errorf("invalid server %w", err)
 	}
 	url.Path = "/build/"
-	resp, err := http.Post(url.String(), "application/json", marshaled) //nolint:noctx
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url.String(), marshaled)
+	if err != nil {
+		return k6build.Artifact{}, fmt.Errorf("%w: %w", ErrRequestFailed, err)
+	}
+	req.Header.Add("Content-Type", "application/json")
+
+	// add authorization header "Authorization: <type> <auth>"
+	if r.auth != "" {
+		authType := r.authType
+		if authType == "" {
+			authType = defaultAuthType
+		}
+		req.Header.Add("Authorization", fmt.Sprintf("%s %s", authType, r.auth))
+	}
+
+	// add custom headers
+	for h, v := range r.headers {
+		req.Header.Add(h, v)
+	}
+
+	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return k6build.Artifact{}, fmt.Errorf("%w: %w", ErrRequestFailed, err)
 	}

pkg/client/client_test.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -14,31 +15,83 @@ import (
 )
 
 type testSrv struct {
-	status   int
-	response api.BuildResponse
+	handlers []requestHandler
+}
+
+// process request and return a boolean indicating if request should be passed to the next handler in the chain
+type requestHandler func(w http.ResponseWriter, r *http.Request) bool
+
+func withValidateRequest() requestHandler {
+	return func(w http.ResponseWriter, r *http.Request) bool {
+		req := api.BuildRequest{}
+		err := json.NewDecoder(r.Body).Decode(&req)
+		if err != nil {
+			w.WriteHeader(http.StatusBadRequest)
+			return false
+		}
+
+		if req.K6Constrains == "" || req.Platform == "" || len(req.Dependencies) == 0 {
+			w.WriteHeader(http.StatusBadRequest)
+			return false
+		}
+
+		return true
+	}
+}
+
+func withAuthorizationCheck(authType string, auth string) requestHandler {
+	return func(w http.ResponseWriter, r *http.Request) bool {
+		authHeader := fmt.Sprintf("%s %s", authType, auth)
+		if r.Header.Get("Authorization") != authHeader {
+			w.WriteHeader(http.StatusUnauthorized)
+			return false
+		}
+		return true
+	}
+}
+
+func withHeadersCheck(headers map[string]string) requestHandler {
+	return func(w http.ResponseWriter, r *http.Request) bool {
+		for h, v := range headers {
+			if r.Header.Get(h) != v {
+				w.WriteHeader(http.StatusBadRequest)
+				return false
+			}
+		}
+		return true
+	}
+}
+
+func withResponse(status int, response api.BuildResponse) requestHandler {
+	return func(w http.ResponseWriter, _ *http.Request) bool {
+		resp := &bytes.Buffer{}
+		err := json.NewEncoder(resp).Encode(response)
+		if err != nil {
+			panic("unexpected error encoding response")
+		}
+
+		w.WriteHeader(status)
+		_, _ = w.Write(resp.Bytes())
+
+		return false
+	}
 }
 
 func (t testSrv) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	w.Header().Add("Content-Type", "application/json")
 
-	// validate request
-	req := api.BuildRequest{}
-	err := json.NewDecoder(r.Body).Decode(&req)
-	if err != nil {
-		w.WriteHeader(http.StatusBadRequest)
-		return
+	// check headers
+	for _, check := range t.handlers {
+		if !check(w, r) {
+			return
+		}
 	}
 
-	// send canned response
+	// by default return ok and an empty artifact
 	resp := &bytes.Buffer{}
-	err = json.NewEncoder(resp).Encode(t.response)
-	if err != nil {
-		// set uncommon status code to signal something unexpected happened
-		w.WriteHeader(http.StatusTeapot)
-		return
-	}
+	_ = json.NewEncoder(resp).Encode(api.BuildResponse{Artifact: k6build.Artifact{}}) //nolint:errchkjson
 
-	w.WriteHeader(t.status)
+	w.WriteHeader(http.StatusOK)
 	_, _ = w.Write(resp.Bytes())
 }
 
@@ -47,35 +100,59 @@ func TestRemote(t *testing.T) {
 
 	testCases := []struct {
 		title     string
-		status    int
-		resp      api.BuildResponse
+		headers   map[string]string
+		auth      string
+		authType  string
+		handlers  []requestHandler
 		expectErr error
 	}{
 		{
-			title:  "normal build",
-			status: http.StatusOK,
-			resp: api.BuildResponse{
-				Error:    "",
-				Artifact: k6build.Artifact{},
+			title: "normal build",
+			handlers: []requestHandler{
+				withValidateRequest(),
+			},
+		},
+		{
+			title: "build request failed",
+			handlers: []requestHandler{
+				withResponse(http.StatusOK, api.BuildResponse{Error: ErrBuildFailed.Error()}),
+			},
+			expectErr: ErrBuildFailed,
+		},
+		{
+			title:    "auth header",
+			auth:     "token",
+			authType: "Bearer",
+			handlers: []requestHandler{
+				withAuthorizationCheck("Bearer", "token"),
+			},
+			expectErr: nil,
+		},
+		{
+			title:    "with default auth type",
+			auth:     "token",
+			authType: "",
+			handlers: []requestHandler{
+				withAuthorizationCheck("Bearer", "token"),
 			},
+			expectErr: nil,
 		},
 		{
-			title:  "request failed",
-			status: http.StatusInternalServerError,
-			resp: api.BuildResponse{
-				Error:    "request failed",
-				Artifact: k6build.Artifact{},
+			title: "failed auth",
+			handlers: []requestHandler{
+				withResponse(http.StatusUnauthorized, api.BuildResponse{Error: "Authorization Required"}),
 			},
 			expectErr: ErrRequestFailed,
 		},
 		{
-			title:  "failed build",
-			status: http.StatusOK,
-			resp: api.BuildResponse{
-				Error:    "failed build",
-				Artifact: k6build.Artifact{},
+			title: "custom headers",
+			headers: map[string]string{
+				"Custom-Header": "Custom-Value",
 			},
-			expectErr: ErrBuildFailed,
+			handlers: []requestHandler{
+				withHeadersCheck(map[string]string{"Custom-Header": "Custom-Value"}),
+			},
+			expectErr: nil,
 		},
 	}
 
@@ -85,13 +162,17 @@ func TestRemote(t *testing.T) {
 			t.Parallel()
 
 			srv := httptest.NewServer(testSrv{
-				status:   tc.status,
-				response: tc.resp,
+				handlers: tc.handlers,
 			})
 
+			defer srv.Close()
+
 			client, err := NewBuildServiceClient(
 				BuildServiceClientConfig{
-					URL: srv.URL,
+					URL:               srv.URL,
+					Headers:           tc.headers,
+					Authorization:     tc.auth,
+					AuthorizationType: tc.authType,
 				},
 			)
 			if err != nil {