From cf59d9bd6d2ff9f3426e66d26dd82df7dc482bf5 Mon Sep 17 00:00:00 2001
From: Daniel Cormier <dcormier@users.noreply.github.com>
Date: Thu, 20 Jun 2024 15:42:06 -0400
Subject: [PATCH] Support passing a state value to `AuthCodeURL` on public
 clients

Related to #485.
---
 apps/public/public.go      | 28 ++++++++++++++++++++++++++--
 apps/public/public_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/apps/public/public.go b/apps/public/public.go
index 392e5e43..067f18c5 100644
--- a/apps/public/public.go
+++ b/apps/public/public.go
@@ -142,7 +142,7 @@ func New(clientID string, options ...Option) (Client, error) {
 
 // authCodeURLOptions contains options for AuthCodeURL
 type authCodeURLOptions struct {
-	claims, loginHint, tenantID, domainHint string
+	claims, loginHint, tenantID, domainHint, state string
 }
 
 // AuthCodeURLOption is implemented by options for AuthCodeURL
@@ -152,7 +152,7 @@ type AuthCodeURLOption interface {
 
 // AuthCodeURL creates a URL used to acquire an authorization code.
 //
-// Options: [WithClaims], [WithDomainHint], [WithLoginHint], [WithTenantID]
+// Options: [WithClaims], [WithDomainHint], [WithLoginHint], [WithTenantID], [WithState]
 func (pca Client) AuthCodeURL(ctx context.Context, clientID, redirectURI string, scopes []string, opts ...AuthCodeURLOption) (string, error) {
 	o := authCodeURLOptions{}
 	if err := options.ApplyOptions(&o, opts); err != nil {
@@ -162,12 +162,36 @@ func (pca Client) AuthCodeURL(ctx context.Context, clientID, redirectURI string,
 	if err != nil {
 		return "", err
 	}
+	ap.State = o.state
 	ap.Claims = o.claims
 	ap.LoginHint = o.loginHint
 	ap.DomainHint = o.domainHint
 	return pca.base.AuthCodeURL(ctx, clientID, redirectURI, scopes, ap)
 }
 
+// WithState adds a user-generated state to the request.
+func WithState(state string) interface {
+	AuthCodeURLOption
+	options.CallOption
+} {
+	return struct {
+		AuthCodeURLOption
+		options.CallOption
+	}{
+		CallOption: options.NewCallOption(
+			func(a any) error {
+				switch t := a.(type) {
+				case *authCodeURLOptions:
+					t.state = state
+				default:
+					return fmt.Errorf("unexpected options type %T", a)
+				}
+				return nil
+			},
+		),
+	}
+}
+
 // WithClaims sets additional claims to request for the token, such as those required by conditional access policies.
 // Use this option when Azure AD returned a claims challenge for a prior request. The argument must be decoded.
 // This option is valid for any token acquisition method.
diff --git a/apps/public/public_test.go b/apps/public/public_test.go
index c0fb9b33..2e4c5970 100644
--- a/apps/public/public_test.go
+++ b/apps/public/public_test.go
@@ -856,6 +856,44 @@ func TestWithDomainHint(t *testing.T) {
 	}
 }
 
+func TestWithState(t *testing.T) {
+	state := "abc-123-secure-string"
+	client, err := New("client-id")
+	if err != nil {
+		t.Fatal(err)
+	}
+	client.base.Token.AccessTokens = &fake.AccessTokens{}
+	client.base.Token.Authority = &fake.Authority{}
+	client.base.Token.Resolver = &fake.ResolveEndpoints{}
+	for _, expectHint := range []bool{true, false} {
+		t.Run(fmt.Sprint(expectHint), func(t *testing.T) {
+			var urlOpts []AuthCodeURLOption
+			if expectHint {
+				urlOpts = append(urlOpts, WithState(state))
+			}
+			u, err := client.AuthCodeURL(context.Background(), "id", "https://localhost", tokenScope, urlOpts...)
+			if err != nil {
+				t.Fatal(err)
+			}
+			parsed, err := url.Parse(u)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !parsed.Query().Has("state") {
+				if !expectHint {
+					return
+				}
+				t.Fatal("expected a state")
+			} else if !expectHint {
+				t.Fatal("expected no state")
+			}
+			if actual := parsed.Query()["state"]; len(actual) != 1 || actual[0] != state {
+				t.Fatalf(`unexpected state "%v"`, actual)
+			}
+		})
+	}
+}
+
 func TestWithAuthenticationScheme(t *testing.T) {
 	clientInfo := base64.RawStdEncoding.EncodeToString([]byte(`{"uid":"uid","utid":"utid"}`))
 	lmo, tenant := "login.microsoftonline.com", "tenant"