poimen-workflows/action/assume_role.go

package action

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"
)

// AssumeRoleInput is the input to AssumeRoleActivity
type AssumeRoleInput struct {
	// Identity is the user/service identity requesting access
	Identity string `json:"identity"`
	
	// ClientID is the OAuth2/OIDC client ID (from vault or env)
	ClientID string `json:"clientId,omitempty"`
	
	// ClientSecret is the OAuth2/OIDC client secret (from vault or env)
	ClientSecret string `json:"clientSecret,omitempty"`
	
	// Scope defines what APIs this token can access (e.g., "llm:read llm:write")
	Scope string `json:"scope"`
	
	// DurationSeconds is how long the token is valid (default: 3600 = 1 hour)
	DurationSeconds int `json:"durationSeconds,omitempty"`
	
	// AuthServerURL is the auth server endpoint (from env if not provided)
	AuthServerURL string `json:"authServerUrl,omitempty"`
}

// AssumeRoleOutput is the output from AssumeRoleActivity
type AssumeRoleOutput struct {
	// Token is the JWT token for calling api.riotpiao.com
	Token string `json:"token"`
	
	// ExpiresAt is when the token expires (Unix timestamp)
	ExpiresAt int64 `json:"expiresAt"`
	
	// ExpiresIn is the duration in seconds until expiration
	ExpiresIn int `json:"expiresIn"`
	
	// TokenType is typically "Bearer"
	TokenType string `json:"tokenType"`
}

// oauthTokenRequest is sent to the auth server
type oauthTokenRequest struct {
	GrantType    string `json:"grant_type"`
	ClientID     string `json:"client_id"`
	ClientSecret string `json:"client_secret"`
	Scope        string `json:"scope"`
	Subject      string `json:"subject,omitempty"` // The identity being assumed
}

// oauthTokenResponse is returned from the auth server
type oauthTokenResponse struct {
	AccessToken string `json:"access_token"`
	TokenType   string `json:"token_type"`
	ExpiresIn   int    `json:"expires_in"`
	Scope       string `json:"scope"`
}

// AssumeRoleActivity requests a temporary JWT token for accessing LLM APIs
//
// This activity works like AWS AssumeRole:
// 1. User provides identity + scope of access needed
// 2. Activity exchanges credentials with auth server
// 3. Returns JWT token valid for a limited time
// 4. Caller uses token in subsequent LLM API calls
//
// Security: Credentials should come from vault/secrets, never hardcoded
func AssumeRoleActivity(ctx context.Context, input *AssumeRoleInput) (*AssumeRoleOutput, error) {
	// Validate inputs
	if err := validateAssumeRoleInput(input); err != nil {
		return nil, err
	}

	// Resolve configuration from input + environment
	config, err := resolveAssumeRoleConfig(input)
	if err != nil {
		return nil, err
	}

	// Request token from auth server
	tokenResp, err := requestAuthToken(ctx, config, input)
	if err != nil {
		return nil, err
	}

	// Build output
	return buildAssumeRoleOutput(tokenResp), nil
}

// validateAssumeRoleInput checks required fields
func validateAssumeRoleInput(input *AssumeRoleInput) error {
	if input.Identity == "" {
		return fmt.Errorf("identity is required")
	}
	if input.Scope == "" {
		return fmt.Errorf("scope is required (e.g., 'llm:read' or 'llm:read llm:write')")
	}
	return nil
}

// assumeRoleConfig holds resolved configuration
type assumeRoleConfig struct {
	authServerURL string
	clientID      string
	clientSecret  string
	duration      int
}

// resolveAssumeRoleConfig gets config from input or environment
func resolveAssumeRoleConfig(input *AssumeRoleInput) (*assumeRoleConfig, error) {
	cfg := &assumeRoleConfig{}

	// Helper function to avoid DRY violation
	getOrEnv := func(val, envKey, fieldName string) (string, error) {
		if val != "" {
			return val, nil
		}
		if val = os.Getenv(envKey); val != "" {
			return val, nil
		}
		return "", fmt.Errorf("%s not provided and %s not set", fieldName, envKey)
	}

	var err error
	if cfg.authServerURL, err = getOrEnv(input.AuthServerURL, "AUTH_SERVER_URL", "authServerUrl"); err != nil {
		return nil, err
	}
	if cfg.clientID, err = getOrEnv(input.ClientID, "OAUTH_CLIENT_ID", "clientId"); err != nil {
		return nil, err
	}
	if cfg.clientSecret, err = getOrEnv(input.ClientSecret, "OAUTH_CLIENT_SECRET", "clientSecret"); err != nil {
		return nil, err
	}

	// Validate and set duration
	cfg.duration = input.DurationSeconds
	if cfg.duration == 0 {
		cfg.duration = 3600 // 1 hour default
	}
	if cfg.duration > 86400 {
		cfg.duration = 86400 // Max 24 hours
	}

	return cfg, nil
}

// requestAuthToken calls the auth server and returns the token response
func requestAuthToken(ctx context.Context, config *assumeRoleConfig, input *AssumeRoleInput) (*oauthTokenResponse, error) {
	tokenReq := oauthTokenRequest{
		GrantType:    "client_credentials",
		ClientID:     config.clientID,
		ClientSecret: config.clientSecret,
		Scope:        input.Scope,
		Subject:      input.Identity,
	}

	reqBody, err := json.Marshal(tokenReq)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal token request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST",
		fmt.Sprintf("%s/oauth/token", config.authServerURL),
		bytes.NewReader(reqBody))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	httpReq.Header.Set("Content-Type", "application/json")

	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("failed to call auth server: %w", err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("auth server returned status %d: %s", resp.StatusCode, string(respBody))
	}

	var tokenResp oauthTokenResponse
	if err := json.Unmarshal(respBody, &tokenResp); err != nil {
		return nil, fmt.Errorf("failed to unmarshal token response: %w", err)
	}

	if tokenResp.AccessToken == "" {
		return nil, fmt.Errorf("auth server returned empty access token")
	}

	return &tokenResp, nil
}

// buildAssumeRoleOutput constructs the output from token response
func buildAssumeRoleOutput(tokenResp *oauthTokenResponse) *AssumeRoleOutput {
	expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Unix()
	return &AssumeRoleOutput{
		Token:     tokenResp.AccessToken,
		ExpiresAt: expiresAt,
		ExpiresIn: tokenResp.ExpiresIn,
		TokenType: tokenResp.TokenType,
	}
}