Repository: amp-labs/connectors
Pull request #1559 (2025-04-11 19:36:54)

feat: Add a Mode to common.Header to control how headers get applied
This adds a mode field to common.Header so we can have a little more control over how headers get applied to a request. They can be appended (use add), overwritten (use set), or added if unset (check for empty then add). This gives us the ability on the server side to specify how clashing headers are handled.

This also adds a generic modifier so that we can inject header removal logic, if necessary.

This is possibly necessary to fix the Asana and Gong connectors -- I think the Authorization header is being added more than once.

eberle1080: @jlimatampersand that's a fair point about the regressions, I'll remove mode from the providers themselves and add some tests. However, the actual fixes for Asana and Gong will happen in the server repo, not here -- the issue has to do with oauth stuff, not core provider calls.

Patch:
diff --git a/common/header.go b/common/header.go
index fcc71cd3..c4a0ba99 100644
--- a/common/header.go
+++ b/common/header.go
@@ -110,7 +110,7 @@ func (c *headerAuthClient) Do(req *http.Request) (*http.Response, error) {
 	req = req.Clone(req.Context())
 
 	for _, header := range c.headers {
-		req.Header.Add(header.Key, header.Value)
+		header.ApplyToRequest(req)
 	}
 
 	if c.dynamicHeaders != nil {
@@ -120,10 +120,15 @@ func (c *headerAuthClient) Do(req *http.Request) (*http.Response, error) {
 		}
 
 		for _, header := range hdrs {
-			req.Header.Add(header.Key, header.Value)
+			header.ApplyToRequest(req)
 		}
 	}
 
+	modifier, hasModifier := getRequestModifier(req.Context()) //nolint:contextcheck
+	if hasModifier {
+		modifier(req)
+	}
+
 	rsp, err := c.client.Do(req)
 	if err != nil {
 		return rsp, err
diff --git a/common/http.go b/common/http.go
index 53fa5378..0ebbf2d1 100644
--- a/common/http.go
+++ b/common/http.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
+	"net/textproto"
 	"net/url"
 	"strconv"
 	"strings"
@@ -15,14 +16,55 @@ import (
 	"github.com/amp-labs/connectors/common/logging"
 )
 
+// HeaderMode determines how the header should be applied to the request.
+type HeaderMode int
+
+const (
+	// headerModeUnset is the default mode. It appends the header to the request.
+	headerModeUnset = iota
+
+	// HeaderModeAppend appends the header to the request.
+	HeaderModeAppend
+
+	// HeaderModeOverwrite unconditionally overwrites the header in the request.
+	HeaderModeOverwrite
+
+	// HeaderModeSetIfMissing sets the header in the request if it is not already set.
+	HeaderModeSetIfMissing
+)
+
 // Header is a key/value pair that can be added to a request.
 type Header struct {
-	Key   string `json:"key"`
-	Value string `json:"value"`
+	Key   string     `json:"key"`
+	Value string     `json:"value"`
+	Mode  HeaderMode `json:"mode"`
+}
+
+func (h Header) ApplyToRequest(req *http.Request) {
+	switch h.Mode {
+	case HeaderModeOverwrite:
+		req.Header.Set(h.Key, h.Value)
+	case HeaderModeSetIfMissing:
+		if len(req.Header.Values(h.Key)) == 0 {
+			req.Header.Add(h.Key, h.Value)
+		}
+	case HeaderModeAppend:
+		fallthrough
+	case headerModeUnset:
+		fallthrough
+	default:
+		req.Header.Add(h.Key, h.Value)
+	}
 }
 
 func (h Header) equals(other Header) bool {
-	return h.Key == other.Key && h.Value == other.Value
+	return textproto.CanonicalMIMEHeaderKey(h.Key) == textproto.CanonicalMIMEHeaderKey(other.Key) &&
+		h.Value == other.Value &&
+		h.Mode == other.Mode
+}
+
+func (h Header) String() string {
+	return fmt.Sprintf("%s: %s", h.Key, h.Value)
 }
 
 var HeaderFormURLEncoded = Header{ // nolint:gochecknoglobals
@@ -42,6 +84,12 @@ func (h Headers) Has(target Header) bool {
 	return false
 }
 
+func (h Headers) ApplyToRequest(req *http.Request) {
+	for _, header := range h {
+		header.ApplyToRequest(req)
+	}
+}
+
 // ErrorHandler allows the caller to inject their own HTTP error handling logic.
 // All non-2xx responses will be passed to the error handler. If the error handler
 // returns nil, then the error is ignored and the caller is responsible for handling
diff --git a/common/jwt.go b/common/jwt.go
deleted file mode 100644
index 88ab74a5..00000000
--- a/common/jwt.go
+++ /dev/null
@@ -1,15 +0,0 @@
-package common
-
-import (
-	"context"
-)
-
-// NewJwtAuthHTTPClient returns a new http client, with automatic Jwt generation &
-// authentication. Specifically this means that the client will automatically
-// add the Jwt (as a header) to every request, based on the given secret and claims.
-func NewJwtAuthHTTPClient( //nolint:ireturn
-	ctx context.Context,
-	opts ...HeaderAuthClientOption,
-) (AuthenticatedHTTPClient, error) {
-	return NewHeaderAuthHTTPClient(ctx, opts...)
-}
diff --git a/common/oauth.go b/common/oauth.go
index 5de5040e..7d5a1ca7 100644
--- a/common/oauth.go
+++ b/common/oauth.go
@@ -198,6 +198,11 @@ func (t *oauth2Transport) RoundTrip(req *http.Request) (*http.Response, error) {
 	// req.Body is assumed to be closed by the base RoundTripper.
 	reqBodyClosed = true
 
+	modifier, hasModifier := getRequestModifier(req2.Context()) //nolint:contextcheck
+	if hasModifier {
+		modifier(req)
+	}
+
 	rsp, err := t.base().RoundTrip(req2)
 	if err != nil {
 		return rsp, err
diff --git a/common/query.go b/common/query.go
index 9450833b..70d9f697 100644
--- a/common/query.go
+++ b/common/query.go
@@ -108,6 +108,11 @@ func (c *queryParamAuthClient) Do(req *http.Request) (*http.Response, error) {
 
 	req.URL.RawQuery = query.Encode()
 
+	modifier, hasModifier := getRequestModifier(req.Context()) //nolint:contextcheck
+	if hasModifier {
+		modifier(req)
+	}
+
 	rsp, err := c.client.Do(req)
 	if err != nil {
 		return rsp, err
diff --git a/common/request.go b/common/request.go
new file mode 100644
index 00000000..fbd03c63
--- /dev/null
+++ b/common/request.go
@@ -0,0 +1,46 @@
+package common
+
+import (
+	"context"
+	"net/http"
+)
+
+type requestModifierKey string
+
+type RequestModifier func(req *http.Request)
+
+// WithRequestModifier adds a request modifier to the context. The request modifier
+// will be called with the request before it is sent. This allows the caller to
+// inject logic into the request, such as adding/removing headers or modifying the
+// request body. Generally this should be used for debugging purposes only. Although
+// there may be some use cases where it is useful to modify the request before it
+// is sent. It can also double as an injectable observer function.
+func WithRequestModifier(ctx context.Context, modifier RequestModifier) context.Context {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	return context.WithValue(ctx, requestModifierKey("requestModifier"), modifier)
+}
+
+func getRequestModifier(ctx context.Context) (RequestModifier, bool) {
+	if ctx == nil {
+		return nil, false
+	}
+
+	value := ctx.Value(requestModifierKey("requestModifier"))
+	if value == nil {
+		return nil, false
+	}
+
+	modifier, ok := value.(RequestModifier)
+	if !ok {
+		return nil, false
+	}
+
+	if modifier == nil {
+		return nil, false
+	}
+
+	return modifier, true
+}


Test patch:
diff --git a/common/http_test.go b/common/http_test.go
new file mode 100644
index 00000000..df556dd7
--- /dev/null
+++ b/common/http_test.go
@@ -0,0 +1,208 @@
+package common
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type dummyTransport struct{}
+
+func (d *dummyTransport) RoundTrip(request *http.Request) (*http.Response, error) {
+	rsp := &http.Response{
+		StatusCode: http.StatusOK,
+		Status:     http.StatusText(http.StatusOK),
+		Header:     make(http.Header),
+		Body:       io.NopCloser(strings.NewReader("hello")),
+	}
+
+	for k, v1 := range request.Header {
+		for _, v2 := range v1 {
+			rsp.Header.Add(k, v2)
+		}
+	}
+
+	return rsp, nil
+}
+
+func TestModifierInClient(t *testing.T) {
+	t.Parallel()
+
+	modifier, ok := getRequestModifier(context.Background())
+	require.False(t, ok)
+	require.Nil(t, modifier)
+
+	client := &http.Client{
+		Transport: &dummyTransport{},
+	}
+
+	ac, err := NewHeaderAuthHTTPClient(context.Background(), WithHeaderClient(client))
+	require.NoError(t, err)
+
+	ctx := WithRequestModifier(context.Background(), func(req *http.Request) {
+		req.Header.Set("Header", "value")
+	})
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://example.com", nil)
+	require.NoError(t, err)
+
+	rsp, err := ac.Do(req)
+	require.NoError(t, err)
+
+	defer func() {
+		_ = rsp.Body.Close()
+	}()
+
+	vals := rsp.Header.Values("Header")
+	require.Len(t, vals, 1)
+
+	require.Equal(t, "value", vals[0])
+}
+
+func TestModifierStandalone(t *testing.T) {
+	t.Parallel()
+
+	modifier, ok := getRequestModifier(context.Background())
+	require.False(t, ok)
+	require.Nil(t, modifier)
+
+	ctx := WithRequestModifier(context.Background(), func(req *http.Request) {
+		req.Header.Set("Header", "value")
+	})
+
+	modifier, ok = getRequestModifier(ctx)
+	require.True(t, ok)
+
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://example.com", nil)
+	require.NoError(t, err)
+
+	modifier(req)
+
+	require.Equal(t, "value", req.Header.Get("Header"))
+}
+
+func TestApplySetIfMissingEmptyHeadersRequest(t *testing.T) {
+	t.Parallel()
+
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://example.com", nil)
+	require.NoError(t, err)
+
+	headers := Headers{
+		{
+			Key:   "Header",
+			Value: "value2",
+			Mode:  HeaderModeSetIfMissing,
+		},
+	}
+
+	headers.ApplyToRequest(req)
+
+	vals := req.Header.Values("Header")
+	require.Len(t, vals, 1)
+
+	require.Equal(t, "value2", vals[0])
+}
+
+func TestApplySetIfMissingNonEmptyHeadersRequest(t *testing.T) {
+	t.Parallel()
+
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://example.com", nil)
+	require.NoError(t, err)
+
+	req.Header.Add("Header", "value1")
+
+	headers := Headers{
+		{
+			Key:   "Header",
+			Value: "value2",
+			Mode:  HeaderModeSetIfMissing,
+		},
+	}
+
+	headers.ApplyToRequest(req)
+
+	vals := req.Header.Values("Header")
+	require.Len(t, vals, 1)
+
+	require.Equal(t, "value1", vals[0])
+}
+
+func TestApplySetHeadersEmptyRequest(t *testing.T) {
+	t.Parallel()
+
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://example.com", nil)
+	require.NoError(t, err)
+
+	headers := Headers{
+		{
+			Key:   "Header",
+			Value: "value2",
+			Mode:  HeaderModeOverwrite,
+		},
+	}
+
+	headers.ApplyToRequest(req)
+
+	vals := req.Header.Values("Header")
+	require.Len(t, vals, 1)
+
+	require.Equal(t, "value2", vals[0])
+}
+
+func TestApplySetHeadersNonEmptyRequest(t *testing.T) {
+	t.Parallel()
+
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://example.com", nil)
+	require.NoError(t, err)
+
+	req.Header.Add("Header", "value1")
+
+	headers := Headers{
+		{
+			Key:   "Header",
+			Value: "value2",
+			Mode:  HeaderModeOverwrite,
+		},
+	}
+
+	headers.ApplyToRequest(req)
+
+	vals := req.Header.Values("Header")
+	require.Len(t, vals, 1)
+
+	require.Equal(t, "value2", vals[0])
+}
+
+func TestApplyAppendHeadersToRequest(t *testing.T) {
+	t.Parallel()
+
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://example.com", nil)
+	require.NoError(t, err)
+
+	req.Header.Add("Header", "value1")
+
+	headers := Headers{
+		{
+			Key:   "Header",
+			Value: "value2",
+			Mode:  HeaderModeAppend,
+		},
+		{
+			Key:   "Header",
+			Value: "value3",
+		},
+	}
+
+	headers.ApplyToRequest(req)
+
+	vals := req.Header.Values("Header")
+	require.Len(t, vals, 3)
+
+	require.Equal(t, "value1", vals[0])
+	require.Equal(t, "value2", vals[1])
+	require.Equal(t, "value3", vals[2])
+}
diff --git a/test/atlassian/connector.go b/test/atlassian/connector.go
index a1c4b94f..352eb4b3 100644
--- a/test/atlassian/connector.go
+++ b/test/atlassian/connector.go
@@ -45,7 +45,7 @@ func GetAtlassianConnectConnector(ctx context.Context, claims map[string]any) *a
 		common.WithDynamicHeaders(atlassian.JwtTokenGenerator(claims, reader.Get(credscanning.Fields.Secret))),
 	}
 
-	client, err := common.NewJwtAuthHTTPClient(ctx, opts...)
+	client, err := common.NewHeaderAuthHTTPClient(ctx, opts...)
 	if err != nil {
 		panic(err)
 	}