feat(auth): implemented SecurityTokenService to handle token exchange (#250)

Signed-off-by: Marc Nuri <[email protected]>

Author: manusa

pkg/http/sts.go

@@ -0,0 +1,43 @@
+package http
+
+import (
+	"context"
+
+	"github.com/coreos/go-oidc/v3/oidc"
+	"golang.org/x/oauth2"
+	"golang.org/x/oauth2/google/externalaccount"
+)
+
+type staticSubjectTokenSupplier struct {
+	token string
+}
+
+func (s *staticSubjectTokenSupplier) SubjectToken(_ context.Context, _ externalaccount.SupplierOptions) (string, error) {
+	return s.token, nil
+}
+
+var _ externalaccount.SubjectTokenSupplier = &staticSubjectTokenSupplier{}
+
+type SecurityTokenService struct {
+	*oidc.Provider
+	ClientId                string
+	ClientSecret            string
+	ExternalAccountAudience string
+	ExternalAccountScopes   []string
+}
+
+func (sts *SecurityTokenService) ExternalAccountTokenExchange(ctx context.Context, originalToken *oauth2.Token) (*oauth2.Token, error) {
+	ts, err := externalaccount.NewTokenSource(ctx, externalaccount.Config{
+		TokenURL:             sts.Endpoint().TokenURL,
+		ClientID:             sts.ClientId,
+		ClientSecret:         sts.ClientSecret,
+		Audience:             sts.ExternalAccountAudience,
+		SubjectTokenType:     "urn:ietf:params:oauth:token-type:access_token",
+		SubjectTokenSupplier: &staticSubjectTokenSupplier{token: originalToken.AccessToken},
+		Scopes:               sts.ExternalAccountScopes,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return ts.Token()
+}

pkg/http/sts_test.go

@@ -0,0 +1,123 @@
+package http
+
+import (
+	"encoding/base64"
+	"fmt"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/containers/kubernetes-mcp-server/internal/test"
+	"github.com/coreos/go-oidc/v3/oidc"
+	"golang.org/x/oauth2"
+)
+
+func TestExternalAccountTokenExchange(t *testing.T) {
+	mockServer := test.NewMockServer()
+	authServer := mockServer.Config().Host
+	var tokenExchangeRequest *http.Request
+	mockServer.Handle(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		if req.URL.Path == "/.well-known/openid-configuration" {
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = fmt.Fprintf(w, `{
+				"issuer": "%s",
+				"authorization_endpoint": "https://mock-oidc-provider/authorize",
+				"token_endpoint": "%s/token"
+			}`, authServer, authServer)
+			return
+		}
+		if req.URL.Path == "/token" {
+			tokenExchangeRequest = req
+			_ = tokenExchangeRequest.ParseForm()
+			if tokenExchangeRequest.PostForm.Get("subject_token") != "the-original-access-token" {
+				http.Error(w, "Invalid subject_token", http.StatusUnauthorized)
+				return
+			}
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`{"access_token":"exchanged-access-token","token_type":"Bearer","expires_in":253402297199}`))
+			return
+		}
+	}))
+	t.Cleanup(mockServer.Close)
+	provider, err := oidc.NewProvider(t.Context(), authServer)
+	if err != nil {
+		t.Fatalf("oidc.NewProvider() error = %v; want nil", err)
+	}
+	// With missing Token Source information
+	_, err = (&SecurityTokenService{Provider: provider}).ExternalAccountTokenExchange(t.Context(), &oauth2.Token{})
+	t.Run("ExternalAccountTokenExchange with missing token source returns error", func(t *testing.T) {
+		if err == nil {
+			t.Fatalf("ExternalAccountTokenExchange() error = nil; want error")
+		}
+		if !strings.Contains(err.Error(), "must be set") {
+			t.Errorf("ExternalAccountTokenExchange() error = %v; want missing required field", err)
+		}
+	})
+	// With valid Token Source information
+	sts := SecurityTokenService{
+		Provider:                provider,
+		ClientId:                "test-client-id",
+		ClientSecret:            "test-client-secret",
+		ExternalAccountAudience: "test-audience",
+		ExternalAccountScopes:   []string{"test-scope"},
+	}
+	// With Invalid token
+	_, err = sts.ExternalAccountTokenExchange(t.Context(), &oauth2.Token{
+		AccessToken: "invalid-access-token",
+		TokenType:   "Bearer",
+	})
+	t.Run("ExternalAccountTokenExchange with invalid token returns error", func(t *testing.T) {
+		if err == nil {
+			t.Fatalf("ExternalAccountTokenExchange() error = nil; want error")
+		}
+		if !strings.Contains(err.Error(), "status code 401: Invalid subject_token") {
+			t.Errorf("ExternalAccountTokenExchange() error = %v; want invalid_grant: Invalid subject_token", err)
+		}
+	})
+	// With Valid token
+	exchangeToken, err := sts.ExternalAccountTokenExchange(t.Context(), &oauth2.Token{
+		AccessToken: "the-original-access-token",
+		TokenType:   "Bearer",
+	})
+	t.Run("ExternalAccountTokenExchange with valid token returns new token", func(t *testing.T) {
+		if err != nil {
+			t.Errorf("ExternalAccountTokenExchange() error = %v; want nil", err)
+		}
+		if exchangeToken == nil {
+			t.Fatal("ExternalAccountTokenExchange() = nil; want token")
+		}
+		if exchangeToken.AccessToken != "exchanged-access-token" {
+			t.Errorf("exchangeToken.AccessToken = %s; want exchanged-access-token", exchangeToken.AccessToken)
+		}
+	})
+	t.Run("ExternalAccountTokenExchange with valid token sends POST request", func(t *testing.T) {
+		if tokenExchangeRequest == nil {
+			t.Fatal("tokenExchangeRequest is nil; want request")
+		}
+		if tokenExchangeRequest.Method != "POST" {
+			t.Errorf("tokenExchangeRequest.Method = %s; want POST", tokenExchangeRequest.Method)
+		}
+	})
+	t.Run("ExternalAccountTokenExchange with valid token has correct form data", func(t *testing.T) {
+		if tokenExchangeRequest.Header.Get("Content-Type") != "application/x-www-form-urlencoded" {
+			t.Errorf("tokenExchangeRequest.Content-Type = %s; want application/x-www-form-urlencoded", tokenExchangeRequest.Header.Get("Content-Type"))
+		}
+		if tokenExchangeRequest.PostForm.Get("audience") != "test-audience" {
+			t.Errorf("tokenExchangeRequest.PostForm[audience] = %s; want test-audience", tokenExchangeRequest.PostForm.Get("audience"))
+		}
+		if tokenExchangeRequest.PostForm.Get("subject_token_type") != "urn:ietf:params:oauth:token-type:access_token" {
+			t.Errorf("tokenExchangeRequest.PostForm[subject_token_type] = %s; want urn:ietf:params:oauth:token-type:access_token", tokenExchangeRequest.PostForm.Get("subject_token_type"))
+		}
+		if tokenExchangeRequest.PostForm.Get("subject_token") != "the-original-access-token" {
+			t.Errorf("tokenExchangeRequest.PostForm[subject_token] = %s; want the-original-access-token", tokenExchangeRequest.PostForm.Get("subject_token"))
+		}
+		if len(tokenExchangeRequest.PostForm["scope"]) == 0 || tokenExchangeRequest.PostForm["scope"][0] != "test-scope" {
+			t.Errorf("tokenExchangeRequest.PostForm[scope] = %v; want [test-scope]", tokenExchangeRequest.PostForm["scope"])
+		}
+	})
+	t.Run("ExternalAccountTokenExchange with valid token sends correct client credentials header", func(t *testing.T) {
+		if tokenExchangeRequest.Header.Get("Authorization") != "Basic "+base64.StdEncoding.EncodeToString([]byte("test-client-id:test-client-secret")) {
+			t.Errorf("tokenExchangeRequest.Header[Authorization] = %s; want Basic base64(test-client-id:test-client-secret)", tokenExchangeRequest.Header.Get("Authorization"))
+		}
+	})
+}