for openziti/ziti#3496 adds TOTP enrollment support during OIDC auth

CHANGELOG.md

@@ -1,3 +1,9 @@
+# Release notes 1.5.3
+
+## Issues Fixed and Dependency Updates
+
+* github.com/openziti/sdk-golang: [v1.5.2 -> v1.5.3](https://github.com/openziti/sdk-golang/compare/v1.5.2...v1.5.3)
+
 # Release notes 1.5.2
 
 ## What's New

edge-apis/client_edge_client.go

@@ -73,9 +73,10 @@ func NewClientApiClientWithConfig(config *ApiClientConfig) *ClientApiClient {
 
 	newApi := rest_client_api_client.New(transportPool, nil)
 	api := ZitiEdgeClient{
-		ZitiEdgeClient:      newApi,
-		TotpCodeProvider:    config.TotpCodeProvider,
-		ClientTransportPool: transportPool,
+		ZitiEdgeClient:         newApi,
+		TotpCodeProvider:       config.TotpCodeProvider,
+		TotpEnrollmentProvider: config.TotpEnrollmentProvider,
+		ClientTransportPool:    transportPool,
 	}
 	ret.API = &api
 	ret.AuthEnabledApi = &api
@@ -103,9 +104,10 @@ type ZitiEdgeClient struct {
 	versionInfo *rest_model.Version
 	versionOnce sync.Once
 
-	TotpCodeProvider    TotpCodeProvider
-	ClientTransportPool ClientTransportPool
-	OidcRedirectUri     string
+	TotpCodeProvider       TotpCodeProvider
+	TotpEnrollmentProvider TotpEnrollmentProvider
+	ClientTransportPool    ClientTransportPool
+	OidcRedirectUri        string
 }
 
 func (self *ZitiEdgeClient) SetOidcRedirectUri(redirectUri string) {
@@ -182,12 +184,13 @@ func (self *ZitiEdgeClient) legacyAuth(credentials Credentials, configTypes []st
 // oidcAuth performs OIDC OAuth flow based authentication.
 func (self *ZitiEdgeClient) oidcAuth(credentials Credentials, configTypeOverrides []string, httpClient *http.Client) (ApiSession, error) {
 	config := &EdgeOidcAuthConfig{
-		ClientTransportPool: self.ClientTransportPool,
-		Credentials:         credentials,
-		ConfigTypeOverrides: configTypeOverrides,
-		HttpClient:          httpClient,
-		TotpCodeProvider:    self.TotpCodeProvider,
-		RedirectUri:         self.OidcRedirectUri,
+		ClientTransportPool:    self.ClientTransportPool,
+		Credentials:            credentials,
+		ConfigTypeOverrides:    configTypeOverrides,
+		HttpClient:             httpClient,
+		TotpCodeProvider:       self.TotpCodeProvider,
+		TotpEnrollmentProvider: self.TotpEnrollmentProvider,
+		RedirectUri:            self.OidcRedirectUri,
 	}
 
 	return oidcAuth(config)

edge-apis/client_edge_management.go

@@ -64,9 +64,10 @@ func NewManagementApiClientWithConfig(config *ApiClientConfig) *ManagementApiCli
 
 	newApi := rest_management_api_client.New(transportPool, nil)
 	api := ZitiEdgeManagement{
-		ZitiEdgeManagement:  newApi,
-		TotpCodeProvider:    config.TotpCodeProvider,
-		ClientTransportPool: transportPool,
+		ZitiEdgeManagement:     newApi,
+		TotpCodeProvider:       config.TotpCodeProvider,
+		TotpEnrollmentProvider: config.TotpEnrollmentProvider,
+		ClientTransportPool:    transportPool,
 	}
 
 	ret.API = &api
@@ -96,9 +97,10 @@ type ZitiEdgeManagement struct {
 	versionOnce sync.Once
 	versionInfo *rest_model.Version
 
-	TotpCodeProvider    TotpCodeProvider
-	ClientTransportPool ClientTransportPool
-	OidcRedirectUri     string
+	TotpCodeProvider       TotpCodeProvider
+	TotpEnrollmentProvider TotpEnrollmentProvider
+	ClientTransportPool    ClientTransportPool
+	OidcRedirectUri        string
 }
 
 func (self *ZitiEdgeManagement) SetOidcRedirectUri(redirectUri string) {
@@ -177,12 +179,13 @@ func (self *ZitiEdgeManagement) legacyAuth(credentials Credentials, configTypes
 // oidcAuth performs OIDC OAuth flow based authentication.
 func (self *ZitiEdgeManagement) oidcAuth(credentials Credentials, configTypeOverrides []string, httpClient *http.Client) (ApiSession, error) {
 	config := &EdgeOidcAuthConfig{
-		ClientTransportPool: self.ClientTransportPool,
-		Credentials:         credentials,
-		ConfigTypeOverrides: configTypeOverrides,
-		HttpClient:          httpClient,
-		TotpCodeProvider:    self.TotpCodeProvider,
-		RedirectUri:         self.OidcRedirectUri,
+		ClientTransportPool:    self.ClientTransportPool,
+		Credentials:            credentials,
+		ConfigTypeOverrides:    configTypeOverrides,
+		HttpClient:             httpClient,
+		TotpCodeProvider:       self.TotpCodeProvider,
+		TotpEnrollmentProvider: self.TotpEnrollmentProvider,
+		RedirectUri:            self.OidcRedirectUri,
 	}
 	return oidcAuth(config)
 }

edge-apis/clients_shared.go

@@ -76,32 +76,58 @@ type OidcAuthResponses struct {
 	// PrimaryCredentialResponse is the response provided during primary credential authentication, nil if not reached due to OIDC init errors
 	PrimaryCredentialResponse *resty.Response
 
-	// TotpResponse is the response provided after the TOTP code was provided, nil if TOTP is not needed
+	// TotpEnrollResponse is the response from starting TOTP enrollment, nil if enrollment was not performed
+	TotpEnrollResponse *resty.Response
+
+	// TotpEnrollVerifyResponse is the response from verifying the enrollment TOTP code, nil if enrollment was not performed
+	TotpEnrollVerifyResponse *resty.Response
+
+	// TotpResponse is the response provided after the TOTP code was provided for an already-enrolled identity, nil if not reached
 	TotpResponse *resty.Response
 
 	// RedirectResponse is the response provided after authentication, nil if never reached
 	RedirectResponse *resty.Response
 }
 
+// OidcAuthorizeResult holds the result of starting an OIDC PKCE authorization flow.
+// It provides the auth request ID, a pre-configured resty client for making raw HTTP calls
+// to the OP login endpoints, and an exchange function that completes the flow by trading
+// an authorization code for OIDC tokens.
+type OidcAuthorizeResult struct {
+
+	// AuthRequestId is the auth request identifier returned by the /oidc/authorize endpoint.
+	AuthRequestId string
+
+	// Client is a resty client pre-configured with the correct TLS settings and redirect policy
+	// for the OIDC flow. Use it to make raw HTTP calls to OP login endpoints.
+	Client *resty.Client
+
+	// Exchange completes the OIDC flow by exchanging an authorization code (extracted from the
+	// callback redirect Location header) for OIDC tokens.
+	Exchange func(code string) (*oidc.Tokens[*oidc.IDTokenClaims], error)
+}
+
 // EdgeOidcAuthConfig represents the options necessary to complete an OAuth 2.0 PKCE authentication flow against an
 // OpenZiti controller.
 type EdgeOidcAuthConfig struct {
-	ClientTransportPool ClientTransportPool
-	Credentials         Credentials
-	ConfigTypeOverrides []string
-	HttpClient          *http.Client
-	TotpCodeProvider    TotpCodeProvider
-	RedirectUri         string
-	ApiHost             string
+	ClientTransportPool    ClientTransportPool
+	Credentials            Credentials
+	ConfigTypeOverrides    []string
+	HttpClient             *http.Client
+	TotpCodeProvider       TotpCodeProvider
+	TotpEnrollmentProvider TotpEnrollmentProvider
+	RedirectUri            string
+	ApiHost                string
 }
 
 // ApiClientConfig contains configuration options for creating API clients.
 type ApiClientConfig struct {
-	ApiUrls          []*url.URL
-	CaPool           *x509.CertPool
-	TotpCodeProvider TotpCodeProvider
-	Components       *Components
-	Proxy            func(r *http.Request) (*url.URL, error)
+	ApiUrls                []*url.URL
+	CaPool                 *x509.CertPool
+	TotpCodeProvider       TotpCodeProvider
+	TotpEnrollmentProvider TotpEnrollmentProvider
+	Components             *Components
+	Proxy                  func(r *http.Request) (*url.URL, error)
 }
 
 // exchangeTokens exchanges OIDC tokens for refreshed tokens. It uses refresh tokens preferentially,
@@ -236,6 +262,10 @@ type totpCodePayload struct {
 	AuthRequestId string `json:"id"`
 }
 
+type authRequestIdPayload struct {
+	AuthRequestId string `json:"id"`
+}
+
 func (a *authPayload) toValues() url.Values {
 	result := url.Values{
 		"id":            []string{a.AuthRequestId},
@@ -369,8 +399,10 @@ func (e *EdgeOidcAuthenticator) AuthenticateWithResponses() (*oidc.Tokens[*oidc.
 	authResponses, err := e.handlePrimaryAndSecondaryAuth(verificationParams)
 
 	if authResponses != nil {
-		oidcAuthResponses.TotpResponse = authResponses.TotpResp
 		oidcAuthResponses.PrimaryCredentialResponse = authResponses.PrimaryResp
+		oidcAuthResponses.TotpEnrollResponse = authResponses.TotpEnrollResp
+		oidcAuthResponses.TotpEnrollVerifyResponse = authResponses.TotpEnrollVerifyResp
+		oidcAuthResponses.TotpResponse = authResponses.TotpResp
 		oidcAuthResponses.RedirectResponse = authResponses.RedirectResp
 	}
 	if err != nil {
@@ -385,6 +417,29 @@ func (e *EdgeOidcAuthenticator) AuthenticateWithResponses() (*oidc.Tokens[*oidc.
 	return tokens, oidcAuthResponses, nil
 }
 
+// Authorize starts the OIDC PKCE authorization flow without completing it, returning an
+// OidcAuthorizeResult for use when tests need to make raw HTTP calls to the OP login endpoints.
+// The Exchange function in the result completes the flow by trading an authorization code for tokens.
+func (e *EdgeOidcAuthenticator) Authorize() (*OidcAuthorizeResult, error) {
+	pkceParams, err := newPkceParameters()
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate PKCE parameters: %w", err)
+	}
+
+	verificationParams, _, err := e.initOAuthFlow(pkceParams)
+	if err != nil {
+		return nil, fmt.Errorf("failed to initiate authorization flow: %w", err)
+	}
+
+	return &OidcAuthorizeResult{
+		AuthRequestId: verificationParams.AuthRequestId,
+		Client:        e.restyClient,
+		Exchange: func(code string) (*oidc.Tokens[*oidc.IDTokenClaims], error) {
+			return e.exchangeAuthorizationCodeForTokens(code, pkceParams)
+		},
+	}, nil
+}
+
 // finishOAuthFlow extracts the authorization code from the callback redirect and exchanges it for tokens.
 // The authorization server returns the code as a query parameter in the Location header of the redirect response.
 // The code is then used with the PKCE verifier to obtain OIDC tokens via the token endpoint.
@@ -430,9 +485,11 @@ func (e *EdgeOidcAuthenticator) finishOAuthFlow(redirectResp *resty.Response, ve
 // the primary credential submission and optional secondary authentication
 // steps of the OIDC flow. Fields are nil if their corresponding step was not reached.
 type PrimaryAndSecondaryAuthResponses struct {
-	RedirectResp *resty.Response
-	PrimaryResp  *resty.Response
-	TotpResp     *resty.Response
+	RedirectResp         *resty.Response
+	PrimaryResp          *resty.Response
+	TotpEnrollResp       *resty.Response
+	TotpEnrollVerifyResp *resty.Response
+	TotpResp             *resty.Response
 }
 
 // handlePrimaryAndSecondaryAuth submits credentials to the authorization endpoint and handles optional TOTP.
@@ -482,6 +539,24 @@ func (e *EdgeOidcAuthenticator) handlePrimaryAndSecondaryAuth(verificationParams
 		return authResponses, errors.New("response was not a redirect and TOTP is not required, unknown additional authentication steps are required but unsupported")
 	}
 
+	// Parse the auth queries response body to determine if TOTP is already enrolled.
+	var authQueriesResp struct {
+		AuthQueries []*rest_model.AuthQueryDetail `json:"authQueries"`
+	}
+	isTotpEnrolled := true // default to enrolled; enrollment check requires a parseable body
+	if err = json.Unmarshal(resp.Body(), &authQueriesResp); err == nil {
+		for _, q := range authQueriesResp.AuthQueries {
+			if q.TypeID == rest_model.AuthQueryTypeTOTP {
+				isTotpEnrolled = q.IsTotpEnrolled
+				break
+			}
+		}
+	}
+
+	if !isTotpEnrolled {
+		return e.handleTotpEnrollment(authResponses, payload.AuthRequestId)
+	}
+
 	if e.TotpCodeProvider == nil {
 		return authResponses, errors.New("totp is required but no totp callback was defined")
 	}
@@ -525,6 +600,81 @@ func (e *EdgeOidcAuthenticator) handlePrimaryAndSecondaryAuth(verificationParams
 	}
 }
 
+// handleTotpEnrollment performs TOTP enrollment during an OIDC authentication flow.
+// It starts enrollment to obtain a provisioning URL, delegates QR code display and code
+// collection to the configured TotpEnrollmentProvider, then verifies the resulting code
+// to complete enrollment.
+func (e *EdgeOidcAuthenticator) handleTotpEnrollment(authResponses *PrimaryAndSecondaryAuthResponses, authRequestId string) (*PrimaryAndSecondaryAuthResponses, error) {
+	if e.TotpEnrollmentProvider == nil {
+		return authResponses, errors.New("totp enrollment is required but no totp enrollment provider was configured")
+	}
+
+	enrollUri := "https://" + e.ApiHost + "/oidc/login/totp/enroll"
+	enrollVerifyUri := "https://" + e.ApiHost + "/oidc/login/totp/enroll/verify"
+
+	enrollResp, err := e.restyClient.R().SetBody(&authRequestIdPayload{
+		AuthRequestId: authRequestId,
+	}).Post(enrollUri)
+
+	authResponses.TotpEnrollResp = enrollResp
+
+	if err != nil {
+		return authResponses, fmt.Errorf("totp enrollment start failed: %w", err)
+	}
+
+	if enrollResp.StatusCode() != http.StatusCreated {
+		return authResponses, fmt.Errorf("totp enrollment start failed with status %d", enrollResp.StatusCode())
+	}
+
+	var enrollDetail rest_model.DetailMfa
+	if err = json.Unmarshal(enrollResp.Body(), &enrollDetail); err != nil {
+		return authResponses, fmt.Errorf("failed to parse totp enrollment response: %w", err)
+	}
+
+	provisioningUrl := enrollDetail.ProvisioningURL
+	if provisioningUrl == "" {
+		return authResponses, errors.New("totp enrollment response did not contain a provisioning URL")
+	}
+
+	enrollResultCh := e.TotpEnrollmentProvider.GetTotpEnrollmentCode(provisioningUrl)
+	var enrollCode string
+
+	select {
+	case enrollResult := <-enrollResultCh:
+		if enrollResult.Err != nil {
+			return authResponses, fmt.Errorf("totp enrollment cancelled: %w", enrollResult.Err)
+		}
+		enrollCode = enrollResult.Code
+	case <-time.After(30 * time.Minute):
+		return authResponses, errors.New("timeout waiting for totp enrollment code")
+	}
+
+	enrollVerifyResp, err := e.restyClient.R().SetBody(&totpCodePayload{
+		MfaCode: rest_model.MfaCode{
+			Code: &enrollCode,
+		},
+		AuthRequestId: authRequestId,
+	}).Post(enrollVerifyUri)
+
+	authResponses.TotpEnrollVerifyResp = enrollVerifyResp
+
+	if err != nil {
+		return authResponses, fmt.Errorf("totp enrollment verification failed: %w", err)
+	}
+
+	switch enrollVerifyResp.StatusCode() {
+	case http.StatusFound:
+		authResponses.RedirectResp = enrollVerifyResp
+		return authResponses, nil
+	case http.StatusOK:
+		return authResponses, errors.New("totp enrollment verified, but additional authentication is required that is not supported or not configured, cannot authenticate")
+	case http.StatusBadRequest:
+		return authResponses, errors.New("totp enrollment code did not verify")
+	default:
+		return authResponses, fmt.Errorf("unexpected response code %d from totp enrollment verification", enrollVerifyResp.StatusCode())
+	}
+}
+
 // initOAuthFlow initiates the OAuth authorization request with PKCE parameters and returns the authorization request ID.
 func (e *EdgeOidcAuthenticator) initOAuthFlow(pkceParams *pkceParameters) (*verificationParameters, *resty.Response, error) {
 	verificationParams := &verificationParameters{

edge-apis/totp.go

@@ -15,6 +15,32 @@ type TotpCodeResult struct {
 	Err  error
 }
 
+// TotpEnrollmentResult represents the outcome of a TOTP enrollment request. A non-empty Code
+// means the user scanned the QR code and entered the resulting code to complete enrollment.
+// A non-nil Err means the user cancelled or denied enrollment.
+type TotpEnrollmentResult struct {
+	Code string
+	Err  error
+}
+
+// TotpEnrollmentProvider is called during OIDC authentication when an identity has not yet
+// enrolled in TOTP but their auth policy requires it. Implementations show the provisioning
+// URL (as a QR code or plain text) to the user, collect their initial TOTP code, and return
+// it via the channel. Send a non-nil Err to cancel enrollment.
+type TotpEnrollmentProvider interface {
+	// GetTotpEnrollmentCode receives the provisioning URL for QR code display and returns
+	// a channel that delivers the TOTP code entered by the user, or an error to cancel.
+	GetTotpEnrollmentCode(provisioningUrl string) <-chan TotpEnrollmentResult
+}
+
+// TotpEnrollmentProviderFunc is a function adapter that implements TotpEnrollmentProvider.
+type TotpEnrollmentProviderFunc func(provisioningUrl string) <-chan TotpEnrollmentResult
+
+// GetTotpEnrollmentCode implements TotpEnrollmentProvider.
+func (f TotpEnrollmentProviderFunc) GetTotpEnrollmentCode(provisioningUrl string) <-chan TotpEnrollmentResult {
+	return f(provisioningUrl)
+}
+
 // TotpTokenResult represents the outcome of exchanging a TOTP code for a session token,
 // including the token value, issuance timestamp, and any errors encountered.
 type TotpTokenResult struct {