distributor/internal/app/upload_http.go

package app

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"

	"gitea.maximumdirect.net/eric/distributor/internal/config"
	"gitea.maximumdirect.net/eric/distributor/internal/ingest"
)

type uploadCoordinator interface {
	CanAccept() bool
	Submit(context.Context, UploadRequest) (UploadRunRecord, error)
	Status(UploadRunID) (UploadRunRecord, bool)
}

type uploadHTTPHandler struct {
	coordinator     uploadCoordinator
	tokens          map[string]resolvedUploadToken
	uploadPipelines map[string]struct{}
}

type resolvedUploadToken struct {
	ID               string
	Value            string
	AllowedPipelines map[string]struct{}
}

type uploadAcceptedResponse struct {
	RunID  UploadRunID  `json:"run_id"`
	Status UploadStatus `json:"status"`
}

type httpErrorResponse struct {
	Error     string `json:"error"`
	Retryable bool   `json:"retryable,omitempty"`
}

const idempotencyKeyHeader = "Idempotency-Key"

func newUploadHTTPHandler(ctx context.Context, cfg config.Config, environment config.Environment) (http.Handler, error) {
	config.ApplyDefaults(&cfg)
	tokens, err := resolveUploadTokens(cfg, environment)
	if err != nil {
		return nil, err
	}
	return uploadHTTPHandler{
		coordinator:     NewUploadCoordinator(ctx, cfg),
		tokens:          tokens,
		uploadPipelines: uploadPipelineSet(cfg),
	}, nil
}

func resolveUploadTokens(cfg config.Config, environment config.Environment) (map[string]resolvedUploadToken, error) {
	tokens := make(map[string]resolvedUploadToken)
	for _, uploadToken := range cfg.UploadTokens {
		token, ok := environment.Lookup(uploadToken.TokenEnv)
		if !ok {
			return nil, fmt.Errorf("upload token %s environment variable %s is not set", uploadToken.ID, uploadToken.TokenEnv)
		}
		if token == "" {
			return nil, fmt.Errorf("upload token %s environment variable %s is empty", uploadToken.ID, uploadToken.TokenEnv)
		}
		if existing, exists := tokens[token]; exists {
			return nil, fmt.Errorf("upload token environment variables for tokens %s and %s resolve to the same value", existing.ID, uploadToken.ID)
		}
		tokens[token] = resolvedUploadToken{
			ID:               uploadToken.ID,
			Value:            token,
			AllowedPipelines: pipelineIDSet(uploadToken.AllowPipelines),
		}
	}
	return tokens, nil
}

func uploadPipelineSet(cfg config.Config) map[string]struct{} {
	pipelines := make(map[string]struct{})
	for _, pipeline := range cfg.Pipelines {
		if pipeline.Source.Backend == config.BackendHTTPUpload {
			pipelines[pipeline.ID] = struct{}{}
		}
	}
	return pipelines
}

func pipelineIDSet(ids []string) map[string]struct{} {
	set := make(map[string]struct{}, len(ids))
	for _, id := range ids {
		set[id] = struct{}{}
	}
	return set
}

func (handler uploadHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == http.MethodGet && r.URL.Path == "/healthz":
		handler.handleHealth(w)
	case r.Method == http.MethodPost && strings.HasPrefix(r.URL.Path, "/v1/pipelines/"):
		handler.handleUpload(w, r)
	case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/runs/"):
		handler.handleRunStatus(w, r)
	default:
		writeHTTPError(w, http.StatusNotFound, "not found")
	}
}

func (handler uploadHTTPHandler) handleHealth(w http.ResponseWriter) {
	writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}

func (handler uploadHTTPHandler) handleUpload(w http.ResponseWriter, r *http.Request) {
	if r.URL.Query().Has("pipeline") || r.URL.Query().Has("pipeline_id") {
		writeHTTPError(w, http.StatusBadRequest, "pipeline id is not accepted")
		return
	}
	pipelineID, ok := uploadPipelineIDFromPath(r.URL.Path)
	if !ok {
		writeHTTPError(w, http.StatusNotFound, "not found")
		return
	}
	if !config.IsSlugLikeID(pipelineID) {
		writeHTTPError(w, http.StatusBadRequest, "invalid pipeline id")
		return
	}
	token, ok := handler.authenticate(r.Header.Get("Authorization"))
	if !ok {
		writeHTTPError(w, http.StatusUnauthorized, "unauthorized")
		return
	}
	if _, ok := handler.uploadPipelines[pipelineID]; !ok {
		writeHTTPError(w, http.StatusNotFound, "upload pipeline not found")
		return
	}
	if _, ok := token.AllowedPipelines[pipelineID]; !ok {
		writeHTTPError(w, http.StatusForbidden, "forbidden")
		return
	}
	contentType := r.Header.Get("Content-Type")
	if err := ingest.ValidateContentType(contentType); err != nil {
		writeHTTPError(w, http.StatusUnsupportedMediaType, "unsupported content type")
		return
	}
	idempotencyKey, err := uploadIdempotencyKey(r.Header)
	if err != nil {
		writeHTTPError(w, http.StatusBadRequest, "invalid idempotency key")
		return
	}
	record, err := handler.coordinator.Submit(r.Context(), UploadRequest{
		TokenID:        token.ID,
		PipelineID:     pipelineID,
		ContentType:    contentType,
		Body:           r.Body,
		IdempotencyKey: idempotencyKey,
	})
	if err != nil {
		writeUploadSubmitError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, uploadAcceptedResponse{
		RunID:  record.ID,
		Status: UploadStatusAccepted,
	})
}

func uploadPipelineIDFromPath(path string) (string, bool) {
	const prefix = "/v1/pipelines/"
	const suffix = "/upload"
	if !strings.HasPrefix(path, prefix) || !strings.HasSuffix(path, suffix) {
		return "", false
	}
	pipelineID := strings.TrimSuffix(strings.TrimPrefix(path, prefix), suffix)
	if pipelineID == "" || strings.Contains(pipelineID, "/") {
		return "", false
	}
	return pipelineID, true
}

func (handler uploadHTTPHandler) handleRunStatus(w http.ResponseWriter, r *http.Request) {
	rawRunID := strings.TrimPrefix(r.URL.Path, "/runs/")
	if rawRunID == "" || strings.Contains(rawRunID, "/") {
		writeHTTPError(w, http.StatusNotFound, "not found")
		return
	}
	record, ok := handler.coordinator.Status(UploadRunID(rawRunID))
	if !ok {
		writeHTTPError(w, http.StatusNotFound, "run not found")
		return
	}
	writeJSON(w, http.StatusOK, record)
}

func (handler uploadHTTPHandler) authenticate(header string) (resolvedUploadToken, bool) {
	const prefix = "Bearer "
	if !strings.HasPrefix(header, prefix) {
		return resolvedUploadToken{}, false
	}
	token := strings.TrimSpace(strings.TrimPrefix(header, prefix))
	if token == "" {
		return resolvedUploadToken{}, false
	}
	resolved, ok := handler.tokens[token]
	return resolved, ok
}

func uploadIdempotencyKey(header http.Header) (string, error) {
	values := header.Values(idempotencyKeyHeader)
	if len(values) == 0 {
		return "", nil
	}
	if len(values) != 1 {
		return "", fmt.Errorf("idempotency key must appear at most once")
	}
	key := values[0]
	if key == "" {
		return "", fmt.Errorf("idempotency key is required when header is present")
	}
	if len(key) > 128 {
		return "", fmt.Errorf("idempotency key must be at most 128 bytes")
	}
	for index := 0; index < len(key); index++ {
		character := key[index]
		if character >= 'a' && character <= 'z' ||
			character >= 'A' && character <= 'Z' ||
			character >= '0' && character <= '9' ||
			character == '.' ||
			character == '_' ||
			character == '-' ||
			character == ':' {
			continue
		}
		return "", fmt.Errorf("idempotency key contains unsupported character")
	}
	return key, nil
}

func writeUploadSubmitError(w http.ResponseWriter, err error) {
	var idempotencyConflict UploadIdempotencyConflictError
	switch {
	case IsUploadQueueFull(err):
		writeHTTPError(w, http.StatusServiceUnavailable, "upload queue is full")
	case errors.As(err, &idempotencyConflict):
		writeHTTPErrorRetryable(w, http.StatusConflict, idempotencyConflict.Error(), idempotencyConflict.Retryable)
	case errors.Is(err, ingest.ErrUploadTooLarge):
		writeHTTPError(w, http.StatusRequestEntityTooLarge, "upload exceeds maximum size")
	case errors.Is(err, ingest.ErrUnsupportedContentType):
		writeHTTPError(w, http.StatusUnsupportedMediaType, "unsupported content type")
	default:
		writeHTTPError(w, http.StatusBadRequest, "upload rejected")
	}
}

func writeHTTPError(w http.ResponseWriter, status int, message string) {
	writeHTTPErrorRetryable(w, status, message, false)
}

func writeHTTPErrorRetryable(w http.ResponseWriter, status int, message string, retryable bool) {
	writeJSON(w, status, httpErrorResponse{Error: message, Retryable: retryable})
}

func writeJSON(w http.ResponseWriter, status int, value any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(value)
}