easyai-ai-gateway/apps/api/internal/executionpool/identity.go

package executionpool

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/url"
	"strings"
	"time"
)

func QueueName(poolID string) string {
	poolID = strings.TrimSpace(poolID)
	if poolID == "" {
		return "gateway_tasks"
	}
	digest := sha256.Sum256([]byte(poolID))
	return "gateway_pool_" + fmt.Sprintf("%x", digest[:8])
}

func RouteProfileKey(provider, protocol, endpointHost, proxyMode, configRevision string) string {
	value := strings.Join([]string{
		strings.ToLower(strings.TrimSpace(provider)),
		strings.ToLower(strings.TrimSpace(protocol)),
		strings.ToLower(strings.TrimSpace(endpointHost)),
		strings.ToLower(strings.TrimSpace(proxyMode)),
		strings.TrimSpace(configRevision),
	}, "\x00")
	digest := sha256.Sum256([]byte(value))
	return fmt.Sprintf("route_%x", digest[:16])
}

func EndpointHost(rawURL string) string {
	parsed, err := url.Parse(strings.TrimSpace(rawURL))
	if err != nil {
		return ""
	}
	return strings.ToLower(parsed.Hostname())
}

func ValidateAdvertisedEndpoint(rawURL string, allowedSuffixes []string, allowPrivate bool) error {
	parsed, err := url.Parse(strings.TrimSpace(rawURL))
	if err != nil || parsed.Host == "" || parsed.User != nil {
		return errors.New("worker endpoint must be an absolute URL without user info")
	}
	if parsed.Scheme != "https" && parsed.Scheme != "http" {
		return errors.New("worker endpoint scheme must be http or https")
	}
	host := strings.ToLower(parsed.Hostname())
	if host == "" {
		return errors.New("worker endpoint host is required")
	}
	if ip := net.ParseIP(host); ip != nil {
		if allowPrivate && (ip.IsPrivate() || ip.IsLoopback()) {
			return nil
		}
		return errors.New("worker endpoint IP is outside the configured trust boundary")
	}
	for _, suffix := range allowedSuffixes {
		suffix = strings.ToLower(strings.TrimSpace(suffix))
		if suffix != "" && (host == strings.TrimPrefix(suffix, ".") || strings.HasSuffix(host, "."+strings.TrimPrefix(suffix, "."))) {
			return nil
		}
	}
	return errors.New("worker endpoint host is outside the configured trust boundary")
}

type ExecutionClaims struct {
	Audience  string `json:"aud"`
	TaskID    string `json:"task_id"`
	PoolID    string `json:"pool_id"`
	WorkerID  string `json:"worker_id"`
	Nonce     string `json:"nonce"`
	ExpiresAt int64  `json:"exp"`
}

type TokenSigner struct {
	Secret []byte
	Now    func() time.Time
}

func (s TokenSigner) Sign(claims ExecutionClaims) (string, error) {
	if len(s.Secret) < 32 {
		return "", errors.New("execution token secret must be at least 32 bytes")
	}
	if claims.Audience == "" || claims.TaskID == "" || claims.PoolID == "" || claims.WorkerID == "" || claims.Nonce == "" || claims.ExpiresAt <= 0 {
		return "", errors.New("execution token claims are incomplete")
	}
	payload, err := json.Marshal(claims)
	if err != nil {
		return "", err
	}
	encoded := base64.RawURLEncoding.EncodeToString(payload)
	mac := hmac.New(sha256.New, s.Secret)
	_, _ = mac.Write([]byte(encoded))
	signature := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
	return encoded + "." + signature, nil
}

func (s TokenSigner) Verify(token, audience string) (ExecutionClaims, error) {
	var claims ExecutionClaims
	if len(s.Secret) < 32 {
		return claims, errors.New("execution token secret must be at least 32 bytes")
	}
	parts := strings.Split(token, ".")
	if len(parts) != 2 {
		return claims, errors.New("invalid execution token")
	}
	signature, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return claims, errors.New("invalid execution token signature")
	}
	mac := hmac.New(sha256.New, s.Secret)
	_, _ = mac.Write([]byte(parts[0]))
	if !hmac.Equal(signature, mac.Sum(nil)) {
		return claims, errors.New("invalid execution token signature")
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil || json.Unmarshal(payload, &claims) != nil {
		return ExecutionClaims{}, errors.New("invalid execution token payload")
	}
	now := time.Now()
	if s.Now != nil {
		now = s.Now()
	}
	if claims.ExpiresAt <= now.Unix() {
		return ExecutionClaims{}, errors.New("execution token expired")
	}
	if claims.Audience != audience {
		return ExecutionClaims{}, errors.New("execution token audience mismatch")
	}
	return claims, nil
}