vault-plugin-secrets-litellm/path_credentials.go

package litellm

import (
	"context"
	"fmt"
	"time"

	"github.com/hashicorp/go-uuid"
	"github.com/hashicorp/vault/sdk/framework"
	"github.com/hashicorp/vault/sdk/logical"
)

func pathCredentials(b *litellmBackend) *framework.Path {
	return &framework.Path{
		Pattern: "creds/" + framework.GenericNameRegex("name"),
		DisplayAttrs: &framework.DisplayAttributes{
			OperationPrefix: "litellm",
			OperationSuffix: "credentials",
		},
		Fields: map[string]*framework.FieldSchema{
			"name": {
				Type:        framework.TypeLowerCaseString,
				Description: "Name of the role to generate a key for.",
				Required:    true,
			},
		},
		Operations: map[logical.Operation]framework.OperationHandler{
			logical.ReadOperation: &framework.PathOperation{
				Callback: b.pathCredentialsRead,
			},
			logical.UpdateOperation: &framework.PathOperation{
				Callback: b.pathCredentialsRead,
			},
		},
		HelpSynopsis:    "Generate a LiteLLM virtual key from a role.",
		HelpDescription: "Reading from this path generates a new LiteLLM virtual key scoped by the named role's models, budget, and TTL.",
	}
}

func (b *litellmBackend) pathCredentialsRead(ctx context.Context, req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	roleName := data.Get("name").(string)

	role, err := b.getRole(ctx, req.Storage, roleName)
	if err != nil {
		return nil, err
	}
	if role == nil {
		return logical.ErrorResponse("role %q does not exist", roleName), nil
	}

	return b.createKey(ctx, req, roleName, role)
}

// createKey issues a new LiteLLM virtual key for the given role and wraps it in
// a Vault lease.
func (b *litellmBackend) createKey(ctx context.Context, req *logical.Request, roleName string, role *litellmRole) (*logical.Response, error) {
	client, err := b.getClient(ctx, req.Storage)
	if err != nil {
		return nil, err
	}

	// Resolve the lease TTL against the role and mount/system limits.
	ttl, maxTTL := b.resolveTTLs(role)

	suffix, err := uuid.GenerateUUID()
	if err != nil {
		return nil, fmt.Errorf("generating key alias suffix: %w", err)
	}
	alias := fmt.Sprintf("%s-%s-%s", role.KeyAliasPrefix, roleName, suffix[:8])

	genReq := generateKeyRequest{
		Models:   role.Models,
		KeyAlias: alias,
		Metadata: role.Metadata,
	}
	if role.MaxBudget > 0 {
		budget := role.MaxBudget
		genReq.MaxBudget = &budget
	}
	if ttl > 0 {
		genReq.Duration = durationToLiteLLM(ttl)
	}

	key, err := client.GenerateKey(ctx, genReq)
	if err != nil {
		return nil, fmt.Errorf("generating litellm key: %w", err)
	}

	internal := map[string]interface{}{
		"key":       key.Key,
		"key_alias": alias,
		"role":      roleName,
	}
	external := map[string]interface{}{
		"key":        key.Key,
		"key_alias":  alias,
		"models":     role.Models,
		"max_budget": role.MaxBudget,
		"role":       roleName,
	}

	resp := b.Secret(litellmKeyType).Response(external, internal)
	resp.Secret.TTL = ttl
	resp.Secret.MaxTTL = maxTTL
	if ttl > 0 {
		resp.Secret.Renewable = true
	}

	return resp, nil
}

// resolveTTLs clamps the role's TTL/MaxTTL against the mount and system limits.
func (b *litellmBackend) resolveTTLs(role *litellmRole) (ttl, maxTTL time.Duration) {
	sysMaxTTL := b.System().MaxLeaseTTL()

	maxTTL = role.MaxTTL
	if maxTTL <= 0 || maxTTL > sysMaxTTL {
		maxTTL = sysMaxTTL
	}

	ttl = role.TTL
	if ttl <= 0 {
		ttl = b.System().DefaultLeaseTTL()
	}
	if ttl > maxTTL {
		ttl = maxTTL
	}
	return ttl, maxTTL
}

// durationToLiteLLM converts a Go duration to a LiteLLM duration string. LiteLLM
// accepts an integer with an s/m/h/d suffix; seconds is always representable.
func durationToLiteLLM(d time.Duration) string {
	return fmt.Sprintf("%ds", int(d.Seconds()))
}