distributor/internal/app/upload_http.go

package app

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime"
	"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]string
	limits      map[string]int64
}

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

type httpErrorResponse struct {
	Error string `json:"error"`
}

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

func resolveUploadTokens(cfg config.Config, environment config.Environment) (map[string]string, map[string]int64, error) {
	tokens := make(map[string]string)
	limits := make(map[string]int64)
	for _, pipeline := range cfg.Pipelines {
		if pipeline.Source.Backend != config.BackendHTTPUpload {
			continue
		}
		tokenName := pipeline.Source.Upload.TokenEnv
		token, ok := environment.Lookup(tokenName)
		if !ok {
			return nil, nil, fmt.Errorf("upload token environment variable %s is not set", tokenName)
		}
		if token == "" {
			return nil, nil, fmt.Errorf("upload token environment variable %s is empty", tokenName)
		}
		if existing, exists := tokens[token]; exists {
			return nil, nil, fmt.Errorf("upload token environment variables for pipelines %s and %s resolve to the same value", existing, pipeline.ID)
		}
		tokens[token] = pipeline.ID
		limits[pipeline.ID] = int64(*pipeline.Source.Upload.MaxUploadSize)
	}
	return tokens, limits, nil
}

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 && r.URL.Path == "/upload":
		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 := handler.authenticate(r.Header.Get("Authorization"))
	if !ok {
		writeHTTPError(w, http.StatusUnauthorized, "unauthorized")
		return
	}
	contentType := r.Header.Get("Content-Type")
	if !supportedUploadContentType(contentType) {
		writeHTTPError(w, http.StatusUnsupportedMediaType, "unsupported content type")
		return
	}
	if !handler.coordinator.CanAccept() {
		writeHTTPError(w, http.StatusServiceUnavailable, "upload queue is full")
		return
	}
	body, err := readUploadBody(r.Body, handler.limits[pipelineID])
	if err != nil {
		if errors.Is(err, ingest.ErrUploadTooLarge) {
			writeHTTPError(w, http.StatusRequestEntityTooLarge, "upload exceeds maximum size")
			return
		}
		writeHTTPError(w, http.StatusBadRequest, "read upload body failed")
		return
	}
	record, err := handler.coordinator.Submit(r.Context(), UploadRequest{
		PipelineID:  pipelineID,
		ContentType: contentType,
		Body:        bytes.NewReader(body),
	})
	if err != nil {
		writeUploadSubmitError(w, err)
		return
	}
	writeJSON(w, http.StatusAccepted, uploadAcceptedResponse{
		RunID:  record.ID,
		Status: UploadStatusAccepted,
	})
}

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) (string, bool) {
	const prefix = "Bearer "
	if !strings.HasPrefix(header, prefix) {
		return "", false
	}
	token := strings.TrimSpace(strings.TrimPrefix(header, prefix))
	if token == "" {
		return "", false
	}
	pipelineID, ok := handler.tokens[token]
	return pipelineID, ok
}

func supportedUploadContentType(contentType string) bool {
	mediaType, _, err := mime.ParseMediaType(contentType)
	if err != nil {
		mediaType = contentType
	}
	switch mediaType {
	case ingest.ContentTypeTar, ingest.ContentTypeGzip, ingest.ContentTypeXGzip:
		return true
	default:
		return false
	}
}

func readUploadBody(body io.Reader, maxSize int64) ([]byte, error) {
	limited := &io.LimitedReader{R: body, N: maxSize + 1}
	data, err := io.ReadAll(limited)
	if err != nil {
		return nil, err
	}
	if int64(len(data)) > maxSize {
		return nil, ingest.ErrUploadTooLarge
	}
	return data, nil
}

func writeUploadSubmitError(w http.ResponseWriter, err error) {
	switch {
	case IsUploadQueueFull(err):
		writeHTTPError(w, http.StatusServiceUnavailable, "upload queue is full")
	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) {
	writeJSON(w, status, httpErrorResponse{Error: message})
}

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