Agent-System/internal/taskrun/backend.go

package taskrun

import (
	"errors"
	"strings"
)

var (
	ErrUnauthorizedBackend = errors.New("taskrun: unauthorized backend")
	ErrSensitivePayload    = errors.New("taskrun: sensitive input rejected")
)

// Allowed backends are resolved by server policy. Callers cannot inject arbitrary strings.
var allowedBackends = map[string]struct{}{
	"stub":      {},
	"pi-agent":  {},
	"pi-native": {},
}

func isProductionEnv(env string) bool {
	switch strings.ToLower(strings.TrimSpace(env)) {
	case "main", "prod", "production":
		return true
	default:
		return false
	}
}

// ResolveBackend maps a caller hint onto the server allowlist.
// An empty requested value falls back to configured. Unauthorized values are rejected.
// Prefer ResolveBackendInEnv when the deployment environment is known.
func ResolveBackend(requested, configured string) (string, error) {
	return ResolveBackendInEnv(requested, configured, "")
}

// ResolveBackendInEnv is ResolveBackend with environment policy:
// in main/prod, stub is never allowed (neither as request nor as empty fallback).
func ResolveBackendInEnv(requested, configured, env string) (string, error) {
	requested = strings.TrimSpace(requested)
	configured = strings.TrimSpace(configured)
	prod := isProductionEnv(env)

	if requested == "" {
		if configured == "" {
			if prod {
				return "", ErrUnauthorizedBackend
			}
			return "stub", nil
		}
		if _, ok := allowedBackends[configured]; !ok {
			return "", ErrUnauthorizedBackend
		}
		if prod && configured == "stub" {
			return "", ErrUnauthorizedBackend
		}
		return configured, nil
	}
	if _, ok := allowedBackends[requested]; !ok {
		return "", ErrUnauthorizedBackend
	}
	if prod && requested == "stub" {
		return "", ErrUnauthorizedBackend
	}
	return requested, nil
}

func StripSecrets(meta map[string]any) map[string]any {
	if meta == nil {
		return map[string]any{}
	}
	out := make(map[string]any, len(meta))
	for k, v := range meta {
		lk := strings.ToLower(strings.TrimSpace(k))
		if lk == "authorization" || lk == "api_key" || lk == "token" || lk == "password" ||
			strings.Contains(lk, "secret") || strings.Contains(lk, "credential") {
			continue
		}
		out[k] = v
	}
	return out
}

func rejectSensitiveText(parts ...string) error {
	for _, p := range parts {
		low := strings.ToLower(p)
		if strings.Contains(low, "api_key=") || strings.Contains(low, "authorization: bearer") {
			return ErrSensitivePayload
		}
	}
	return nil
}