From 1e1fbf66b807da89dd17373e03e0b8be4caa2f11 Mon Sep 17 00:00:00 2001 From: yesoreyeram <153843+yesoreyeram@users.noreply.github.com> Date: Tue, 4 Jun 2024 21:00:51 +0100 Subject: [PATCH 1/2] added custom header --- clientcredentials/clientcredentials.go | 10 ++++++++++ jwt/jwt.go | 10 ++++++++++ token.go | 19 ++++++++++++++++++- token_test.go | 2 ++ 4 files changed, 40 insertions(+), 1 deletion(-) diff --git a/clientcredentials/clientcredentials.go b/clientcredentials/clientcredentials.go index 2459d069f..4d0d098bf 100644 --- a/clientcredentials/clientcredentials.go +++ b/clientcredentials/clientcredentials.go @@ -51,6 +51,12 @@ type Config struct { // authStyleCache caches which auth style to use when Endpoint.AuthStyle is // the zero value (AuthStyleAutoDetect). authStyleCache internal.LazyAuthStyleCache + + // CustomTokenHeaderKey - If set, Instead of `Authorization` this will be used as header key + CustomTokenHeaderKey string + + // CustomTokenPrefix - If set, Instead of `Bearer ` prefix / token type retrieved, this will be used + CustomTokenPrefix string } // Token uses client credentials to retrieve a token. @@ -120,5 +126,9 @@ func (c *tokenSource) Token() (*oauth2.Token, error) { RefreshToken: tk.RefreshToken, Expiry: tk.Expiry, } + if c.conf != nil { + t.CustomTokenHeaderKey = c.conf.CustomTokenHeaderKey + t.CustomTokenPrefix = c.conf.CustomTokenPrefix + } return t.WithExtra(tk.Raw), nil } diff --git a/jwt/jwt.go b/jwt/jwt.go index b2bf18298..6cf563905 100644 --- a/jwt/jwt.go +++ b/jwt/jwt.go @@ -74,6 +74,12 @@ type Config struct { // UseIDToken optionally specifies whether ID token should be used instead // of access token when the server returns both. UseIDToken bool + + // CustomTokenHeaderKey - If set, Instead of `Authorization` this will be used as header key + CustomTokenHeaderKey string + + // CustomTokenPrefix - If set, Instead of `Bearer ` prefix / token type retrieved, this will be used + CustomTokenPrefix string } // TokenSource returns a JWT TokenSource using the configuration @@ -181,5 +187,9 @@ func (js jwtSource) Token() (*oauth2.Token, error) { } token.AccessToken = tokenRes.IDToken } + if js.conf != nil { + token.CustomTokenHeaderKey = js.conf.CustomTokenHeaderKey + token.CustomTokenPrefix = js.conf.CustomTokenPrefix + } return token, nil } diff --git a/token.go b/token.go index 5bbb33217..071ded26d 100644 --- a/token.go +++ b/token.go @@ -57,10 +57,19 @@ type Token struct { // expired, by subtracting from Expiry. If zero, defaultExpiryDelta // is used. expiryDelta time.Duration + + // CustomTokenHeaderKey - If set, Instead of `Authorization` this will be used as header key + CustomTokenHeaderKey string + + // CustomTokenPrefix - If set, Instead of `Bearer ` prefix / token type retrieved, this will be used + CustomTokenPrefix string } // Type returns t.TokenType if non-empty, else "Bearer". func (t *Token) Type() string { + if t.CustomTokenPrefix != "" { + return "" + } if strings.EqualFold(t.TokenType, "bearer") { return "Bearer" } @@ -82,7 +91,15 @@ func (t *Token) Type() string { // This method is unnecessary when using Transport or an HTTP Client // returned by this package. func (t *Token) SetAuthHeader(r *http.Request) { - r.Header.Set("Authorization", t.Type()+" "+t.AccessToken) + headerKey := "Authorization" + if strings.TrimSpace(t.CustomTokenHeaderKey) != "" { + headerKey = t.CustomTokenHeaderKey + } + headerValue := strings.TrimSpace(strings.Join([]string{t.Type(), t.AccessToken}, " ")) + if t.CustomTokenPrefix != "" { + headerValue = strings.TrimSpace(strings.Join([]string{t.CustomTokenPrefix, headerValue}, "")) + } + r.Header.Set(headerKey, headerValue) } // WithExtra returns a new Token that's a clone of t, but using the diff --git a/token_test.go b/token_test.go index 0d8c7df86..f2cac281a 100644 --- a/token_test.go +++ b/token_test.go @@ -71,6 +71,8 @@ func TestTokenTypeMethod(t *testing.T) { {name: "mac", tok: &Token{TokenType: "mac"}, want: "MAC"}, {name: "mac-caps", tok: &Token{TokenType: "MAC"}, want: "MAC"}, {name: "mac-mixed_case", tok: &Token{TokenType: "mAc"}, want: "MAC"}, + {name: "clear type", tok: &Token{TokenType: "Bearer", CustomTokenPrefix: " "}, want: ""}, + {name: "custom type", tok: &Token{TokenType: "mAC", CustomTokenPrefix: "something"}, want: ""}, } for _, tc := range cases { if got, want := tc.tok.Type(), tc.want; got != want { From 7bba1e2e45441cb4c2e045aa966ea7a6047cf467 Mon Sep 17 00:00:00 2001 From: yesoreyeram <153843+yesoreyeram@users.noreply.github.com> Date: Mon, 3 Feb 2025 15:54:29 +0000 Subject: [PATCH 2/2] cleanup --- clientcredentials/clientcredentials.go | 6 ++---- jwt/jwt.go | 6 ++---- token.go | 6 +++--- 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/clientcredentials/clientcredentials.go b/clientcredentials/clientcredentials.go index 0b439992b..0d5a6a078 100644 --- a/clientcredentials/clientcredentials.go +++ b/clientcredentials/clientcredentials.go @@ -126,9 +126,7 @@ func (c *tokenSource) Token() (*oauth2.Token, error) { RefreshToken: tk.RefreshToken, Expiry: tk.Expiry, } - if c.conf != nil { - t.CustomTokenHeaderKey = c.conf.CustomTokenHeaderKey - t.CustomTokenPrefix = c.conf.CustomTokenPrefix - } + t.CustomTokenHeaderKey = c.conf.CustomTokenHeaderKey + t.CustomTokenPrefix = c.conf.CustomTokenPrefix return t.WithExtra(tk.Raw), nil } diff --git a/jwt/jwt.go b/jwt/jwt.go index 6cf563905..42719c481 100644 --- a/jwt/jwt.go +++ b/jwt/jwt.go @@ -187,9 +187,7 @@ func (js jwtSource) Token() (*oauth2.Token, error) { } token.AccessToken = tokenRes.IDToken } - if js.conf != nil { - token.CustomTokenHeaderKey = js.conf.CustomTokenHeaderKey - token.CustomTokenPrefix = js.conf.CustomTokenPrefix - } + token.CustomTokenHeaderKey = js.conf.CustomTokenHeaderKey + token.CustomTokenPrefix = js.conf.CustomTokenPrefix return token, nil } diff --git a/token.go b/token.go index c155b2250..d973578de 100644 --- a/token.go +++ b/token.go @@ -102,11 +102,11 @@ func (t *Token) SetAuthHeader(r *http.Request) { if strings.TrimSpace(t.CustomTokenHeaderKey) != "" { headerKey = t.CustomTokenHeaderKey } - headerValue := strings.TrimSpace(strings.Join([]string{t.Type(), t.AccessToken}, " ")) + headerValue := t.Type() + " " + t.AccessToken if t.CustomTokenPrefix != "" { - headerValue = strings.TrimSpace(strings.Join([]string{t.CustomTokenPrefix, headerValue}, "")) + headerValue = t.CustomTokenPrefix + t.AccessToken } - r.Header.Set(headerKey, headerValue) + r.Header.Set(headerKey, strings.TrimSpace(headerValue)) } // WithExtra returns a new Token that's a clone of t, but using the