From 60bd118a9c314eab49b64b799ffe75df52e8966d Mon Sep 17 00:00:00 2001
From: Joshua Casey <joshua.casey@broadcom.com>
Date: Mon, 18 Nov 2024 16:30:07 -0600
Subject: [PATCH] pinniped CLI should print the audit-ID in certain error cases

Co-authored-by: Ryan Richard <richardry@vmware.com>
---
 internal/kubeclient/roundtrip.go |  2 +-
 pkg/oidcclient/login.go          | 47 +++++++++++++++++++++++++++++++-
 test/testlib/env.go              |  2 +-
 3 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/internal/kubeclient/roundtrip.go b/internal/kubeclient/roundtrip.go
index e95c492aa..463d0b853 100644
--- a/internal/kubeclient/roundtrip.go
+++ b/internal/kubeclient/roundtrip.go
@@ -204,7 +204,7 @@ func handleCreateOrUpdate(
 	negotiatedSerializer runtime.NegotiatedSerializer,
 ) (bool, *http.Response, error) {
 	if req.GetBody == nil {
-		return true, nil, fmt.Errorf("unreadible body for request: %#v", middlewareReq) // this should never happen
+		return true, nil, fmt.Errorf("unreadable body for request: %#v", middlewareReq) // this should never happen
 	}
 
 	body, err := req.GetBody()
diff --git a/pkg/oidcclient/login.go b/pkg/oidcclient/login.go
index f1b7e6d71..27f4f7e37 100644
--- a/pkg/oidcclient/login.go
+++ b/pkg/oidcclient/login.go
@@ -33,6 +33,7 @@ import (
 	oidcapi "go.pinniped.dev/generated/latest/apis/supervisor/oidc"
 	"go.pinniped.dev/internal/federationdomain/upstreamprovider"
 	"go.pinniped.dev/internal/httputil/httperr"
+	"go.pinniped.dev/internal/httputil/roundtripper"
 	"go.pinniped.dev/internal/httputil/securityheader"
 	"go.pinniped.dev/internal/net/phttp"
 	"go.pinniped.dev/internal/plog"
@@ -357,6 +358,40 @@ type nopCache struct{}
 func (*nopCache) GetToken(SessionCacheKey) *oidctypes.Token  { return nil }
 func (*nopCache) PutToken(SessionCacheKey, *oidctypes.Token) {}
 
+// maybePrintAuditID will choose to log the auditID when certain failure cases are detected,
+// to give a breadcrumb for an admin to follow.
+// Older Supervisors and other OIDC identity providers may not provide this header.
+func maybePrintAuditID(rt http.RoundTripper) http.RoundTripper {
+	return roundtripper.WrapFunc(rt, func(r *http.Request) (*http.Response, error) {
+		path := r.URL.Path
+		response, responseErr := rt.RoundTrip(r)
+		if response != nil && response.Header.Get("audit-ID") != "" {
+			switch {
+			case response.StatusCode >= http.StatusMultipleChoices && response.StatusCode < http.StatusBadRequest:
+				// failing oauth2/authorize redirects from audit-enabled Supervisors
+
+				location, err := url.Parse(response.Header.Get(httpLocationHeaderName))
+				if err != nil || location.Query().Get("error") != "" {
+					plog.Info("Received auditID for failed request",
+						"path", path,
+						"statusCode", response.StatusCode,
+						"auditID", response.Header.Get("audit-ID"))
+				}
+			case response.StatusCode >= http.StatusBadRequest:
+				// failing discovery, oauth2/authorize, or oauth2/token responses from audit-enabled Supervisors
+
+				plog.Info("Received auditID for failed request",
+					"path", path,
+					"statusCode", response.StatusCode,
+					"auditID", response.Header.Get("audit-ID"))
+			default:
+				// noop
+			}
+		}
+		return response, responseErr
+	})
+}
+
 // Login performs an OAuth2/OIDC authorization code login using a localhost listener.
 func Login(issuer string, clientID string, opts ...Option) (*oidctypes.Token, error) {
 	h := handlerState{
@@ -379,7 +414,15 @@ func Login(issuer string, clientID string, opts ...Option) (*oidctypes.Token, er
 		getEnv:        os.Getenv,
 		listen:        net.Listen,
 		stdinIsTTY:    func() bool { return term.IsTerminal(stdin()) },
-		getProvider:   upstreamoidc.New,
+		getProvider: func(config *oauth2.Config, provider *coreosoidc.Provider, client *http.Client) upstreamprovider.UpstreamOIDCIdentityProviderI {
+			// can't use upstreamoidc.New here since it does not set the Name
+			return &upstreamoidc.ProviderConfig{
+				Name:     issuer, // use the issuer as the Name
+				Config:   config,
+				Provider: provider,
+				Client:   client,
+			}
+		},
 		validateIDToken: func(ctx context.Context, provider *coreosoidc.Provider, audience string, token string) (*coreosoidc.IDToken, error) {
 			return provider.Verifier(&coreosoidc.Config{ClientID: audience}).Verify(ctx, token)
 		},
@@ -393,6 +436,8 @@ func Login(issuer string, clientID string, opts ...Option) (*oidctypes.Token, er
 		}
 	}
 
+	h.httpClient.Transport = maybePrintAuditID(h.httpClient.Transport)
+
 	if h.cliToSendCredentials {
 		if h.loginFlow != "" {
 			return nil, fmt.Errorf("do not use deprecated option WithCLISendingCredentials when using option WithLoginFlow")
diff --git a/test/testlib/env.go b/test/testlib/env.go
index 194097db5..bb6a3e52c 100644
--- a/test/testlib/env.go
+++ b/test/testlib/env.go
@@ -401,7 +401,7 @@ func (e *TestEnv) WithoutCapability(cap Capability) *TestEnv {
 // Please use this sparingly. We would prefer that a test run on every cluster type where it can possibly run, so
 // prefer to run everywhere when possible or use cluster capabilities when needed, rather than looking at the
 // type of cluster to decide to skip a test. However, there are some tests that do not depend on or interact with
-// Kubernetes itself which really only need to run on on a single platform to give us the coverage that we desire.
+// Kubernetes itself which really only need to run on a single platform to give us the coverage that we desire.
 func (e *TestEnv) WithKubeDistribution(distro KubeDistro) *TestEnv {
 	e.t.Helper()
 	if e.KubernetesDistribution != distro {