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) {
	if input.Identity == "" {
		return nil, fmt.Errorf("identity is required")
	}
	
	if input.Scope == "" {
		return nil, fmt.Errorf("scope is required (e.g., 'llm:read' or 'llm:read llm:write')")
	}

	// Get auth server URL from input or environment
	authServerURL := input.AuthServerURL
	if authServerURL == "" {
		authServerURL = os.Getenv("AUTH_SERVER_URL")
		if authServerURL == "" {
			return nil, fmt.Errorf("AUTH_SERVER_URL not set in input or environment")
		}
	}

	// Get credentials from input or environment
	clientID := input.ClientID
	if clientID == "" {
		clientID = os.Getenv("OAUTH_CLIENT_ID")
		if clientID == "" {
			return nil, fmt.Errorf("clientId not provided and OAUTH_CLIENT_ID not set")
		}
	}

	clientSecret := input.ClientSecret
	if clientSecret == "" {
		clientSecret = os.Getenv("OAUTH_CLIENT_SECRET")
		if clientSecret == "" {
			return nil, fmt.Errorf("clientSecret not provided and OAUTH_CLIENT_SECRET not set")
		}
	}

	// Set default duration
	durationSeconds := input.DurationSeconds
	if durationSeconds == 0 {
		durationSeconds = 3600 // 1 hour default
	}
	if durationSeconds > 86400 {
		durationSeconds = 86400 // Max 24 hours
	}

	// Build token request
	tokenReq := oauthTokenRequest{
		GrantType:    "client_credentials",
		ClientID:     clientID,
		ClientSecret: clientSecret,
		Scope:        input.Scope,
		Subject:      input.Identity, // Assume this identity
	}

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

	// Call auth server
	httpReq, err := http.NewRequestWithContext(ctx, "POST",
		fmt.Sprintf("%s/oauth/token", 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 at %s: %w", authServerURL, 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")
	}

	// Calculate expiration
	expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Unix()

	return &AssumeRoleOutput{
		Token:     tokenResp.AccessToken,
		ExpiresAt: expiresAt,
		ExpiresIn: tokenResp.ExpiresIn,
		TokenType: tokenResp.TokenType,
	}, nil
}