Add option to set success and error pages for when calling AcquireInteractiveOption

wayneforrest · wayneforrest · commit acc88e1115e6 · 2024-06-12T22:08:51.000+12:00
diff --git a/apps/internal/local/server.go b/apps/internal/local/server.go
@@ -28,7 +28,7 @@ var okPage = []byte(`
 </html>
 `)
 
-const failPage = `
+var failPage = []byte(`
 <!DOCTYPE html>
 <html>
 <head>
@@ -40,7 +40,7 @@ const failPage = `
 	<p>Error details: error %s error_description: %s</p>
 </body>
 </html>
-`
+`)
 
 // Result is the result from the redirect.
 type Result struct {
@@ -60,7 +60,7 @@ type Server struct {
 }
 
 // New creates a local HTTP server and starts it.
-func New(reqState string, port int) (*Server, error) {
+func New(reqState string, port int, successPage []byte, errorPage []byte) (*Server, error) {
 	var l net.Listener
 	var err error
 	var portStr string
@@ -84,6 +84,14 @@ func New(reqState string, port int) (*Server, error) {
 		return nil, err
 	}
 
+	if len(successPage) > 0 {
+		okPage = successPage
+	}
+
+	if len(errorPage) > 0 {
+		failPage = errorPage
+	}
+
 	serv := &Server{
 		Addr:     fmt.Sprintf("http://localhost:%s", portStr),
 		s:        &http.Server{Addr: "localhost:0", ReadHeaderTimeout: time.Second},
@@ -145,7 +153,7 @@ func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
 		desc := html.EscapeString(q.Get("error_description"))
 		// Note: It is a little weird we handle some errors by not going to the failPage. If they all should,
 		// change this to s.error() and make s.error() write the failPage instead of an error code.
-		_, _ = w.Write([]byte(fmt.Sprintf(failPage, headerErr, desc)))
+		_, _ = w.Write([]byte(fmt.Sprintf(string(failPage), headerErr, desc)))
 		s.putResult(Result{Err: fmt.Errorf(desc)})
 		return
 	}
diff --git a/apps/public/public.go b/apps/public/public.go
@@ -522,13 +522,59 @@ type interactiveAuthOptions struct {
 	claims, domainHint, loginHint, redirectURI, tenantID string
 	openURL                                              func(url string) error
 	authnScheme                                          AuthenticationScheme
+	successPage                                          []byte
+	errorPage                                            []byte
 }
 
 // AcquireInteractiveOption is implemented by options for AcquireTokenInteractive
 type AcquireInteractiveOption interface {
 	acquireInteractiveOption()
 }
 
+func WithSuccessPage(successPage []byte) interface {
+	AcquireInteractiveOption
+	options.CallOption
+} {
+	return struct {
+		AcquireInteractiveOption
+		options.CallOption
+	}{
+		CallOption: options.NewCallOption(
+			func(a any) error {
+				switch t := a.(type) {
+				case *interactiveAuthOptions:
+					t.successPage = successPage
+				default:
+					return fmt.Errorf("unexpected options type %T", a)
+				}
+				return nil
+			},
+		),
+	}
+}
+
+func WithErrorPage(errorPage []byte) interface {
+	AcquireInteractiveOption
+	options.CallOption
+} {
+	return struct {
+		AcquireInteractiveOption
+		options.CallOption
+	}{
+		CallOption: options.NewCallOption(
+			func(a any) error {
+				switch t := a.(type) {
+				case *interactiveAuthOptions:
+					t.errorPage = errorPage
+				default:
+					return fmt.Errorf("unexpected options type %T", a)
+				}
+				return nil
+			},
+		),
+	}
+}
+
 // WithLoginHint pre-populates the login prompt with a username.
 func WithLoginHint(username string) interface {
 	AcquireInteractiveOption
@@ -671,7 +717,7 @@ func (pca Client) AcquireTokenInteractive(ctx context.Context, scopes []string,
 	if o.authnScheme != nil {
 		authParams.AuthnScheme = o.authnScheme
 	}
-	res, err := pca.browserLogin(ctx, redirectURL, authParams, o.openURL)
+	res, err := pca.browserLogin(ctx, redirectURL, authParams, o.openURL, o.successPage, o.successPage)
 	if err != nil {
 		return AuthResult{}, err
 	}
@@ -709,13 +755,13 @@ func parsePort(u *url.URL) (int, error) {
 }
 
 // browserLogin calls openURL and waits for a user to log in
-func (pca Client) browserLogin(ctx context.Context, redirectURI *url.URL, params authority.AuthParams, openURL func(string) error) (interactiveAuthResult, error) {
+func (pca Client) browserLogin(ctx context.Context, redirectURI *url.URL, params authority.AuthParams, openURL func(string) error, successPage []byte, errorPage []byte) (interactiveAuthResult, error) {
 	// start local redirect server so login can call us back
 	port, err := parsePort(redirectURI)
 	if err != nil {
 		return interactiveAuthResult{}, err
 	}
-	srv, err := local.New(params.State, port)
+	srv, err := local.New(params.State, port, successPage, errorPage)
 	if err != nil {
 		return interactiveAuthResult{}, err
 	}
