feat(agent-proxy): support all machine identity auth methods

packages/cmd/agent_proxy.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net/url"
@@ -18,6 +19,7 @@ import (
 	"github.com/Infisical/infisical-merge/packages/util"
 	"github.com/fatih/color"
 	"github.com/go-resty/resty/v2"
+	infisicalSdk "github.com/infisical/go-sdk"
 	"github.com/posthog/posthog-go"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
@@ -78,12 +80,9 @@ var proxyEnvKeys = []string{
 	"OPENCLAW_PROXY_URL",
 }
 
-// Stripped so the agent never sees the long-lived MI credentials, only the scoped short-lived JWT set below.
-var credentialEnvKeys = []string{
-	util.INFISICAL_UNIVERSAL_AUTH_CLIENT_ID_NAME,
-	util.INFISICAL_UNIVERSAL_AUTH_CLIENT_SECRET_NAME,
+var credentialEnvKeys = append([]string{
 	util.INFISICAL_UNIVERSAL_AUTH_ACCESS_TOKEN_NAME,
-}
+}, util.MachineIdentityAuthEnvVars...)
 
 // Addresses of host IPC endpoints. Not secret values, but handing them to the agent points it at
 // sockets it can use: the SSH/GPG agents as a signing oracle, and the session bus, where
@@ -210,37 +209,121 @@ func telemetryAgentName(args []string) string {
 	return filepath.Base(args[0])
 }
 
-func universalAuthCredentialSource(cmd *cobra.Command) string {
-	if cmd.Flags().Changed("client-id") {
-		return "universal-auth-flag"
+func resolveAgentProxyStaticToken(cmd *cobra.Command, subject string) *models.TokenDetails {
+	token, err := util.GetInfisicalToken(cmd)
+	if err != nil {
+		util.HandleError(err, "Unable to resolve authentication")
+	}
+	if token == nil {
+		return nil
+	}
+	if token.Type == util.SERVICE_TOKEN_IDENTIFIER {
+		util.PrintErrorMessageAndExit("The agent proxy does not support service tokens. Use a machine identity access token, or authenticate with --auth-method.")
 	}
-	return "universal-auth-env"
+	failIfTokenExpired(token.Token, subject)
+	return token
 }
 
-// Returns the token and a label for the branch that produced it.
-func resolveAgentToken(cmd *cobra.Command) (*models.TokenDetails, string) {
-	clientID, _ := util.GetCmdFlagOrEnvWithDefaultValue(cmd, "client-id", []string{util.INFISICAL_UNIVERSAL_AUTH_CLIENT_ID_NAME}, "")
-	clientSecret, _ := util.GetCmdFlagOrEnvWithDefaultValue(cmd, "client-secret", []string{util.INFISICAL_UNIVERSAL_AUTH_CLIENT_SECRET_NAME}, "")
+type agentProxyCredential struct {
+	token  *models.TokenDetails
+	login  func() (infisicalSdk.MachineIdentityCredential, error)
+	source string
+}
 
-	if clientID != "" && clientSecret != "" {
-		loginResp, err := util.UniversalAuthLogin(clientID, clientSecret)
-		if err != nil {
-			util.HandleError(err, "Failed to authenticate the agent machine identity")
+// Token first, then machine identity, as in `gateway` and `pam agentic-access`. A token wins however
+// it arrived, so INFISICAL_TOKEN beats a typed --auth-method there too.
+func resolveAgentProxyCredential(cmd *cobra.Command) agentProxyCredential {
+	if token := resolveAgentProxyStaticToken(cmd, "the provided token"); token != nil {
+		return agentProxyCredential{token: token, source: "token"}
+	}
+	if login, source := resolveAgentProxyLogin(cmd); login != nil {
+		return agentProxyCredential{login: login, source: source}
+	}
+	return agentProxyCredential{}
+}
+
+// A client per login, dropped after: AutoTokenRefresh cannot be switched off (a bool tagged
+// `default:"true"`, which setDefaults restores), so one kept alive would re-authenticate alongside the
+// caller's own renewal. Its getter reads that field without the mutex, hence the returned credential.
+func resolveAgentProxyLogin(cmd *cobra.Command) (login func() (infisicalSdk.MachineIdentityCredential, error), source string) {
+	authMethod, err := util.ResolveAuthMethod(cmd)
+	if err != nil {
+		util.HandleError(err, "Unable to parse auth-method flag")
+	}
+
+	if authMethod == "" {
+		clientID, _ := util.GetCmdFlagOrEnvWithDefaultValue(cmd, "client-id", []string{util.INFISICAL_UNIVERSAL_AUTH_CLIENT_ID_NAME}, "")
+		clientSecret, _ := util.GetCmdFlagOrEnvWithDefaultValue(cmd, "client-secret", []string{util.INFISICAL_UNIVERSAL_AUTH_CLIENT_SECRET_NAME}, "")
+		if clientID == "" && clientSecret == "" {
+			return nil, ""
+		}
+		if clientID == "" {
+			util.HandleError(fmt.Errorf("client id required; pass --client-id or set %s", util.INFISICAL_UNIVERSAL_AUTH_CLIENT_ID_NAME))
+		}
+		if clientSecret == "" {
+			util.HandleError(fmt.Errorf("client secret required; pass --client-secret or set %s", util.INFISICAL_UNIVERSAL_AUTH_CLIENT_SECRET_NAME))
 		}
-		return &models.TokenDetails{
-			Type:  util.UNIVERSAL_AUTH_TOKEN_IDENTIFIER,
-			Token: loginResp.AccessToken,
-		}, universalAuthCredentialSource(cmd)
+		authMethod = string(util.AuthStrategy.UNIVERSAL_AUTH)
 	}
 
-	token, err := util.GetInfisicalToken(cmd)
+	if err := util.ValidateAuthMethod(authMethod); err != nil {
+		util.PrintErrorMessageAndExit(err.Error())
+	}
+
+	customHeaders, err := util.GetInfisicalCustomHeadersMap()
 	if err != nil {
-		util.HandleError(err, "Unable to resolve authentication")
+		util.HandleError(err, "Unable to get custom headers")
 	}
-	if token == nil {
-		util.HandleError(fmt.Errorf("authentication required; provide --client-id/--client-secret, env vars, or a token"))
+
+	login = func() (infisicalSdk.MachineIdentityCredential, error) {
+		ctx, cancel := context.WithCancel(cmd.Context())
+		defer cancel()
+
+		client := infisicalSdk.NewInfisicalClient(ctx, infisicalSdk.Config{
+			SiteUrl:       config.INFISICAL_URL,
+			UserAgent:     api.USER_AGENT,
+			CustomHeaders: customHeaders,
+		})
+		authenticate, err := util.MachineIdentityLoginFunc(cmd, client, authMethod)
+		if err != nil {
+			return infisicalSdk.MachineIdentityCredential{}, err
+		}
+		credential, err := authenticate()
+		if err != nil {
+			return infisicalSdk.MachineIdentityCredential{}, err
+		}
+		if credential.AccessToken == "" {
+			return infisicalSdk.MachineIdentityCredential{}, fmt.Errorf("authenticating with %s returned no access token", authMethod)
+		}
+		return credential, nil
 	}
-	return token, "token"
+	return login, authMethodCredentialSource(cmd, authMethod)
+}
+
+func authMethodCredentialSource(cmd *cobra.Command, authMethod string) string {
+	if cmd.Flags().Changed("auth-method") || cmd.Flags().Changed("client-id") {
+		return authMethod + "-flag"
+	}
+	return authMethod + "-env"
+}
+
+func resolveAgentToken(cmd *cobra.Command) (*models.TokenDetails, string) {
+	resolved := resolveAgentProxyCredential(cmd)
+	if resolved.token != nil {
+		return resolved.token, resolved.source
+	}
+	if resolved.login == nil {
+		util.HandleError(fmt.Errorf("authentication required; pass --auth-method [%s] with that method's credentials, --client-id/--client-secret, or a token", util.MachineIdentityAuthMethods))
+	}
+
+	credential, err := resolved.login()
+	if err != nil {
+		util.HandleError(err, "Failed to authenticate the agent machine identity")
+	}
+	return &models.TokenDetails{
+		Type:  util.UNIVERSAL_AUTH_TOKEN_IDENTIFIER,
+		Token: credential.AccessToken,
+	}, resolved.source
 }
 
 // Builds http://<projectId>:<env>/<path>:<jwt>@host:port (username=projectId, password="<env>/<path>:<jwt>", jwt last).
@@ -370,11 +453,8 @@ func fetchAgentRealSecrets(token *models.TokenDetails, projectID, environment, s
 		ExpandSecretReferences: true,
 		IncludeImport:          true,
 	}
-	if token.Type == util.SERVICE_TOKEN_IDENTIFIER {
-		params.InfisicalToken = token.Token
-	} else if token.Type == util.UNIVERSAL_AUTH_TOKEN_IDENTIFIER {
-		params.UniversalAuthAccessToken = token.Token
-	}
+	// Always an identity access token: service tokens are turned away before this.
+	params.UniversalAuthAccessToken = token.Token
 
 	secrets, err := util.GetAllEnvironmentVariables(params, "")
 	if err != nil {
@@ -474,17 +554,16 @@ func init() {
 	agentProxyConnectCmd.Flags().StringP("env", "e", "", "environment slug to fetch proxied services and secrets from (falls back to INFISICAL_ENVIRONMENT or .infisical.json)")
 	agentProxyConnectCmd.Flags().String("path", "/", "secret path (folder) scope (falls back to INFISICAL_SECRET_PATH or defaultSecretPath in .infisical.json)")
 	agentProxyConnectCmd.Flags().String("projectId", "", "project id (falls back to INFISICAL_PROJECT_ID or .infisical.json)")
-	agentProxyConnectCmd.Flags().String("client-id", "", "universal auth client id for the agent machine identity")
-	agentProxyConnectCmd.Flags().String("client-secret", "", "universal auth client secret for the agent machine identity")
-	agentProxyConnectCmd.Flags().String("token", "", "Fetch secrets using service token or machine identity access token")
+	util.RegisterMachineIdentityAuthFlags(agentProxyConnectCmd, "agent")
+	agentProxyConnectCmd.Flags().String("token", "", "machine identity access token to use instead of authenticating; takes precedence over --auth-method")
 	agentProxyConnectCmd.Flags().String("no-proxy", "", "additional comma-separated hosts to bypass the proxy (always merged with localhost,127.0.0.1)")
 	agentProxyConnectCmd.Flags().Bool("allow-readable-brokered-secrets", false, "start even if the agent can read secrets that proxied services broker to it (bypasses a misconfiguration guardrail; falls back to INFISICAL_AGENT_PROXY_ALLOW_READABLE_BROKERED_SECRETS)")
 
 	agentProxyStartCmd.Flags().Int("port", 17322, "port for the agent proxy to listen on")
 	agentProxyStartCmd.Flags().String("unmatched-host", "allow", "policy for hosts with no proxied service: allow | block")
 	agentProxyStartCmd.Flags().Int("poll-interval", 60, "seconds between permission/credential refreshes for active agents")
-	agentProxyStartCmd.Flags().String("client-id", "", "universal auth client id for the agent proxy machine identity")
-	agentProxyStartCmd.Flags().String("client-secret", "", "universal auth client secret for the agent proxy machine identity")
+	util.RegisterMachineIdentityAuthFlags(agentProxyStartCmd, "agent proxy")
+	agentProxyStartCmd.Flags().String("token", "", "machine identity access token to use instead of authenticating; takes precedence over --auth-method, and the proxy cannot renew it")
 	agentProxyStartCmd.Flags().String("log-format", "console", "log output format: console | json")
 	agentProxyStartCmd.Flags().String("log-file", "", "also write json logs to this file (in addition to the console/json stream)")
 

packages/cmd/agent_proxy_start.go

@@ -9,11 +9,15 @@ import (
 	"github.com/Infisical/infisical-merge/packages/telemetry"
 	"github.com/Infisical/infisical-merge/packages/util"
 	"github.com/fatih/color"
+	infisicalSdk "github.com/infisical/go-sdk"
 	"github.com/posthog/posthog-go"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 )
 
+// The floor on refreshProxyToken's wait, and so the shortest TTL it can stay ahead of.
+const minRefreshableTTL = 30 * time.Second
+
 func runAgentProxyStart(cmd *cobra.Command, args []string) {
 	port, _ := cmd.Flags().GetInt("port")
 	unmatchedHost, _ := cmd.Flags().GetString("unmatched-host")
@@ -33,32 +37,44 @@ func runAgentProxyStart(cmd *cobra.Command, args []string) {
 	}
 	log.Logger = log.Output(logWriter)
 
-	clientID, err := util.GetCmdFlagOrEnvWithDefaultValue(cmd, "client-id", []string{util.INFISICAL_UNIVERSAL_AUTH_CLIENT_ID_NAME}, "")
-	if err != nil || clientID == "" {
-		util.HandleError(fmt.Errorf("agent proxy credentials required; set INFISICAL_UNIVERSAL_AUTH_CLIENT_ID / _SECRET or pass --client-id / --client-secret"))
-	}
-	clientSecret, err := util.GetCmdFlagOrEnvWithDefaultValue(cmd, "client-secret", []string{util.INFISICAL_UNIVERSAL_AUTH_CLIENT_SECRET_NAME}, "")
-	if err != nil || clientSecret == "" {
-		util.HandleError(fmt.Errorf("agent proxy client secret required"))
-	}
+	resolved := resolveAgentProxyCredential(cmd)
 
-	loginResp, err := util.UniversalAuthLogin(clientID, clientSecret)
-	if err != nil {
-		util.HandleError(err, "Failed to authenticate the agent proxy machine identity")
+	var accessToken string
+	var accessTokenTTL int
+	switch {
+	case resolved.token != nil:
+		accessToken = resolved.token.Token
+		log.Warn().Msg("The agent proxy is running on a fixed token, which it cannot renew. It will stop working when that token expires; use --auth-method or client credentials to have it re-authenticate on its own.")
+	case resolved.login != nil:
+		credential, err := resolved.login()
+		if err != nil {
+			util.HandleError(err, "Failed to authenticate the agent proxy machine identity")
+		}
+		accessToken = credential.AccessToken
+		accessTokenTTL = int(credential.ExpiresIn)
+		if accessTokenTTL > 0 && accessTokenTTL <= int(minRefreshableTTL.Seconds()) {
+			util.HandleError(fmt.Errorf("the agent proxy cannot refresh an access token with a TTL of %s or less; raise the TTL on this identity's auth method", minRefreshableTTL))
+		}
+	default:
+		util.HandleError(fmt.Errorf("agent proxy credentials required; pass --auth-method [%s] with that method's credentials, --client-id/--client-secret, or a token", util.MachineIdentityAuthMethods))
 	}
+	credentialSource := resolved.source
+	login := resolved.login
 
-	Telemetry.SetActor(telemetry.IdentityClaimsFromToken(loginResp.AccessToken))
+	Telemetry.SetActor(telemetry.IdentityClaimsFromToken(accessToken))
 	Telemetry.CaptureEvent("cli-command:agent-proxy start", posthog.NewProperties().
 		Set("version", util.CLI_VERSION).
 		Set("unmatchedHost", unmatchedHost).
 		Set("pollInterval", pollInterval).
-		Set("credentialSource", universalAuthCredentialSource(cmd)))
+		Set("credentialSource", credentialSource))
 
 	log.Info().Msg(color.GreenString("Agent proxy authenticated; starting MITM proxy"))
 
 	var proxyToken atomic.Value
-	proxyToken.Store(loginResp.AccessToken)
-	go refreshProxyToken(&proxyToken, clientID, clientSecret, loginResp.AccessTokenTTL)
+	proxyToken.Store(accessToken)
+	if login != nil {
+		go refreshProxyToken(&proxyToken, login, accessTokenTTL)
+	}
 
 	err = agentproxy.Start(agentproxy.Options{
 		Port:          port,
@@ -71,8 +87,8 @@ func runAgentProxyStart(cmd *cobra.Command, args []string) {
 	}
 }
 
-func refreshProxyToken(token *atomic.Value, clientID, clientSecret string, ttlSeconds int) {
-	const retryInterval = 30 * time.Second
+func refreshProxyToken(token *atomic.Value, login func() (infisicalSdk.MachineIdentityCredential, error), ttlSeconds int) {
+	const retryInterval = minRefreshableTTL
 
 	halfTTL := func() time.Duration {
 		wait := time.Duration(ttlSeconds) * time.Second / 2
@@ -86,15 +102,15 @@ func refreshProxyToken(token *atomic.Value, clientID, clientSecret string, ttlSe
 	for {
 		time.Sleep(wait)
 
-		loginResp, err := util.UniversalAuthLogin(clientID, clientSecret)
+		credential, err := login()
 		if err != nil {
 			log.Warn().Err(err).Msgf("Failed to refresh agent proxy token, retrying in %s", retryInterval)
 			wait = retryInterval
 			continue
 		}
-		token.Store(loginResp.AccessToken)
-		if loginResp.AccessTokenTTL > 0 {
-			ttlSeconds = loginResp.AccessTokenTTL
+		token.Store(credential.AccessToken)
+		if credential.ExpiresIn > 0 {
+			ttlSeconds = int(credential.ExpiresIn)
 		}
 		wait = halfTTL()
 	}