NianAIGC/backend/internal/publicapi/auth.go

// Package publicapi authenticates public API clients and internal workers.
// Runtime configuration is parsed once and injected into an Authenticator;
// this package deliberately does not read process environment variables.
package publicapi

import (
	"crypto/subtle"
	"net/http"
	"strings"
)

const ownerPartLimit = 96

type PublicClient struct {
	ID  string
	Key string
}

type Config struct {
	APIKeys             string
	InternalWorkerToken string
	Production          bool
}

type AuthError struct {
	Status  int
	Message string
}

func (e *AuthError) Error() string {
	return e.Message
}

type Authenticator struct {
	clients             []PublicClient
	internalWorkerToken string
	production          bool
}

func NewAuthenticator(config Config) *Authenticator {
	return &Authenticator{
		clients:             ParseClients(config.APIKeys),
		internalWorkerToken: strings.TrimSpace(config.InternalWorkerToken),
		production:          config.Production,
	}
}

func ParseClients(configured string) []PublicClient {
	entries := strings.FieldsFunc(configured, func(character rune) bool {
		return character == ',' || character == '\n'
	})
	clients := make([]PublicClient, 0, len(entries))
	for _, rawEntry := range entries {
		entry := strings.TrimSpace(rawEntry)
		if entry == "" {
			continue
		}

		id := "default"
		key := entry
		if separator := strings.IndexByte(entry, ':'); separator >= 0 {
			id = strings.TrimSpace(entry[:separator])
			key = strings.TrimSpace(entry[separator+1:])
		}
		if id == "" || key == "" {
			continue
		}
		clients = append(clients, PublicClient{ID: id, Key: key})
	}
	return clients
}

func (a *Authenticator) Authenticate(request *http.Request) (PublicClient, string, error) {
	presented := publicCredential(request)
	if presented == "" {
		return PublicClient{}, "", &AuthError{Status: http.StatusUnauthorized, Message: "Missing API key."}
	}
	for _, client := range a.clients {
		if secureEqual(client.Key, presented) {
			return client, OwnerID(client.ID), nil
		}
	}
	return PublicClient{}, "", &AuthError{Status: http.StatusUnauthorized, Message: "Invalid API key."}
}

func (a *Authenticator) AssertInternalWorker(request *http.Request) error {
	if a.internalWorkerToken == "" && !a.production {
		return nil
	}
	if a.internalWorkerToken == "" {
		return &AuthError{Status: http.StatusInternalServerError, Message: "Worker token is not configured."}
	}

	presented := request.Header.Get("x-zhinian-worker-token")
	if presented == "" {
		presented = bearerToken(request)
	}
	if presented == "" || !secureEqual(a.internalWorkerToken, presented) {
		return &AuthError{Status: http.StatusUnauthorized, Message: "Invalid worker token."}
	}
	return nil
}

func OwnerID(id string) string {
	part := sanitizeOwnerPart(id)
	if part == "" {
		part = "unknown"
	}
	return "api:" + part
}

func publicCredential(request *http.Request) string {
	if token := bearerToken(request); token != "" {
		return token
	}
	return request.Header.Get("x-zhinian-api-key")
}

func bearerToken(request *http.Request) string {
	authorization := request.Header.Get("authorization")
	separator := strings.IndexAny(authorization, " \t\r\n\v\f")
	if separator <= 0 || !strings.EqualFold(authorization[:separator], "Bearer") {
		return ""
	}
	if strings.TrimLeft(authorization[separator:], " \t\r\n\v\f") == authorization[separator:] {
		return ""
	}
	return strings.TrimSpace(authorization[separator:])
}

func secureEqual(expected, presented string) bool {
	if len(expected) != len(presented) {
		return false
	}
	return subtle.ConstantTimeCompare([]byte(expected), []byte(presented)) == 1
}

func sanitizeOwnerPart(value string) string {
	part := make([]byte, 0, min(len(value), ownerPartLimit))
	invalidRun := false
	for _, character := range value {
		if isOwnerCharacter(character) {
			invalidRun = false
			if len(part) < ownerPartLimit {
				part = append(part, byte(character))
			}
			continue
		}
		if !invalidRun && len(part) < ownerPartLimit {
			part = append(part, '_')
		}
		invalidRun = true
	}
	return string(part)
}

func isOwnerCharacter(character rune) bool {
	return character >= 'A' && character <= 'Z' ||
		character >= 'a' && character <= 'z' ||
		character >= '0' && character <= '9' ||
		strings.ContainsRune("_.:@-", character)
}