weatherreporter/internal/adapters/distributor/client.go

// Package distributor adapts weatherreporter report artifacts to distributor uploads.
package distributor

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	distributorbundle "gitea.maximumdirect.net/eric/distributor/pkg/bundle"
	distributorupload "gitea.maximumdirect.net/eric/distributor/pkg/upload"

	"gitea.maximumdirect.net/eric/weatherreporter/internal/config"
)

type Client struct {
	Endpoint        string
	TokenEnv        string
	Timeout         time.Duration
	newUploadClient uploadClientFactory
	pollWait        func(context.Context, time.Duration) error
}

type UploadRequest struct {
	PipelineID     string
	BundleID       string
	IdempotencyKey string
	Files          []UploadFile
	CreatedAt      time.Time
}

type UploadFile struct {
	SourcePath string
	BundlePath string
}

type UploadResult struct {
	RunID        string
	Status       string
	UploadStatus string
	StatusError  string
	RunStatus    *RunStatus
}

type RunStatus struct {
	RunID      string
	PipelineID string
	Status     string
	AcceptedAt time.Time
	StartedAt  *time.Time
	FinishedAt *time.Time
	Report     json.RawMessage
	Error      string
}

type IdempotencyConflictError struct {
	Err error
}

func (e *IdempotencyConflictError) Error() string {
	if e == nil || e.Err == nil {
		return "distributor idempotency conflict"
	}
	return e.Err.Error()
}

func (e *IdempotencyConflictError) Unwrap() error {
	if e == nil {
		return nil
	}
	return e.Err
}

type uploadClientFactory func(endpoint, token string, timeout time.Duration) (uploadClient, error)

type uploadClient interface {
	UploadFiles(ctx context.Context, opts uploadFilesOptions) (uploadFilesResult, error)
	Status(ctx context.Context, runID string) (runStatus, error)
}

type uploadFilesOptions struct {
	PipelineID     string
	BundleID       string
	IdempotencyKey string
	Files          []UploadFile
	CreatedAt      time.Time
}

type uploadFilesResult struct {
	RunID  string
	Status string
}

type runStatus struct {
	RunID      string
	PipelineID string
	Status     string
	AcceptedAt time.Time
	StartedAt  *time.Time
	FinishedAt *time.Time
	Report     json.RawMessage
	Error      string
}

const statusPollInterval = 250 * time.Millisecond

const maxDistributorResponseBytes int64 = 1 << 20

var errDistributorResponseTooLarge = fmt.Errorf("distributor response exceeds the %d-byte limit", maxDistributorResponseBytes)

func New(cfg config.DistributorNotifyConfig) *Client {
	return newClient(cfg, newDistributorUploadClient)
}

func newClient(cfg config.DistributorNotifyConfig, factory uploadClientFactory) *Client {
	if factory == nil {
		factory = newDistributorUploadClient
	}
	return &Client{
		Endpoint:        cfg.Endpoint,
		TokenEnv:        cfg.TokenEnv,
		Timeout:         cfg.Timeout,
		newUploadClient: factory,
		pollWait:        waitForPoll,
	}
}

func (c *Client) Upload(ctx context.Context, req UploadRequest) (UploadResult, error) {
	if c == nil {
		return UploadResult{}, fmt.Errorf("distributor client is nil")
	}
	if c.Endpoint == "" {
		return UploadResult{}, fmt.Errorf("distributor endpoint is required")
	}
	if c.TokenEnv == "" {
		return UploadResult{}, fmt.Errorf("distributor token environment variable is required")
	}
	if req.PipelineID == "" {
		return UploadResult{}, fmt.Errorf("distributor pipeline id is required")
	}
	if req.BundleID == "" {
		return UploadResult{}, fmt.Errorf("distributor bundle id is required")
	}
	if req.IdempotencyKey == "" {
		return UploadResult{}, fmt.Errorf("distributor idempotency key is required for bundle %q", req.BundleID)
	}
	if len(req.Files) == 0 {
		return UploadResult{}, fmt.Errorf("distributor upload files are required for bundle %q", req.BundleID)
	}
	for i, file := range req.Files {
		if file.SourcePath == "" {
			return UploadResult{}, fmt.Errorf("distributor source path is required for bundle %q file %d", req.BundleID, i)
		}
		if file.BundlePath == "" {
			return UploadResult{}, fmt.Errorf("distributor bundle path is required for bundle %q file %d", req.BundleID, i)
		}
	}
	if c.newUploadClient == nil {
		return UploadResult{}, fmt.Errorf("distributor upload client factory is required for endpoint %q", c.Endpoint)
	}

	token := os.Getenv(c.TokenEnv)
	if token == "" {
		return UploadResult{}, fmt.Errorf("distributor token environment variable %q is not set", c.TokenEnv)
	}

	uploadClient, err := c.newUploadClient(c.Endpoint, token, c.Timeout)
	if err != nil {
		return UploadResult{}, fmt.Errorf("create distributor upload client for endpoint %q: %w", c.Endpoint, redactToken(err, token))
	}

	runCtx := ctx
	if runCtx == nil {
		runCtx = context.Background()
	}
	cancel := func() {}
	if c.Timeout > 0 {
		runCtx, cancel = context.WithTimeout(runCtx, c.Timeout)
	}
	defer cancel()

	result, err := uploadClient.UploadFiles(runCtx, uploadFilesOptions{
		PipelineID:     req.PipelineID,
		BundleID:       req.BundleID,
		IdempotencyKey: req.IdempotencyKey,
		Files:          append([]UploadFile(nil), req.Files...),
		CreatedAt:      req.CreatedAt,
	})
	if err != nil {
		return UploadResult{}, wrapUploadError(err, uploadErrorContext{
			Endpoint:       c.Endpoint,
			PipelineID:     req.PipelineID,
			BundleID:       req.BundleID,
			IdempotencyKey: req.IdempotencyKey,
			SourcePaths:    uploadSourcePaths(req.Files),
			BundlePaths:    uploadBundlePaths(req.Files),
			Token:          token,
		})
	}

	uploadResult := UploadResult{
		RunID:        result.RunID,
		Status:       result.Status,
		UploadStatus: result.Status,
	}
	pollWait := c.pollWait
	if pollWait == nil {
		pollWait = waitForPoll
	}
	status, statusErr := waitForRunStatus(runCtx, uploadClient, result.RunID, c.Timeout > 0, pollWait)
	status = sanitizeRunStatus(status)
	if status.RunID != "" || status.Status != "" {
		uploadResult.RunStatus = &RunStatus{
			RunID:      status.RunID,
			PipelineID: status.PipelineID,
			Status:     status.Status,
			AcceptedAt: status.AcceptedAt,
			StartedAt:  status.StartedAt,
			FinishedAt: status.FinishedAt,
			Report:     append(json.RawMessage(nil), status.Report...),
			Error:      redactTokenString(status.Error, token),
		}
		if status.Status != "" {
			uploadResult.Status = status.Status
		}
	}
	if statusErr != nil {
		uploadResult.StatusError = safeDistributorDiagnostic(statusErr, token).Error()
		return uploadResult, nil
	}
	if status.Status == "failed" {
		return uploadResult, fmt.Errorf("distributor run %q failed: %s", status.RunID, uploadResult.RunStatus.Error)
	}
	return uploadResult, nil
}

func waitForRunStatus(ctx context.Context, client uploadClient, runID string, poll bool, wait func(context.Context, time.Duration) error) (runStatus, error) {
	status, err := client.Status(ctx, runID)
	if err != nil || terminalRunStatus(status.Status) || !poll {
		return status, err
	}

	for {
		if err := wait(ctx, statusPollInterval); err != nil {
			return status, fmt.Errorf("distributor run %q did not reach terminal status before timeout: %w", runID, err)
		}

		next, err := client.Status(ctx, runID)
		if err != nil {
			return status, err
		}
		status = next
		if terminalRunStatus(status.Status) {
			return status, nil
		}
	}
}

func waitForPoll(ctx context.Context, interval time.Duration) error {
	timer := time.NewTimer(interval)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func terminalRunStatus(status string) bool {
	return status == "succeeded" || status == "failed"
}

type distributorUploadClient struct {
	client *distributorupload.Client
}

type boundedResponseTransport struct {
	base  http.RoundTripper
	limit int64
}

func (t boundedResponseTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	base := t.base
	if base == nil {
		base = http.DefaultTransport
	}
	response, err := base.RoundTrip(req)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	data, err := io.ReadAll(io.LimitReader(response.Body, t.limit+1))
	if err != nil {
		return nil, err
	}
	if int64(len(data)) > t.limit {
		return nil, errDistributorResponseTooLarge
	}
	response.Body = io.NopCloser(bytes.NewReader(data))
	response.ContentLength = int64(len(data))
	return response, nil
}

type RemoteResponseError struct {
	StatusCode int
	Retryable  bool
}

func (e *RemoteResponseError) Error() string {
	if e == nil || e.StatusCode == 0 {
		return "distributor request failed"
	}
	return fmt.Sprintf("distributor request failed with HTTP status %d", e.StatusCode)
}

func newDistributorUploadClient(endpoint, token string, timeout time.Duration) (uploadClient, error) {
	httpClient := &http.Client{
		Transport: boundedResponseTransport{base: http.DefaultTransport, limit: maxDistributorResponseBytes},
	}
	if timeout > 0 {
		httpClient.Timeout = timeout
	}
	client, err := distributorupload.NewClient(distributorupload.ClientOptions{
		Endpoint:   endpoint,
		Token:      token,
		HTTPClient: httpClient,
	})
	if err != nil {
		return nil, err
	}
	return distributorUploadClient{client: client}, nil
}

func (c distributorUploadClient) UploadFiles(ctx context.Context, opts uploadFilesOptions) (uploadFilesResult, error) {
	files := make([]distributorbundle.BundleFile, 0, len(opts.Files))
	for _, file := range opts.Files {
		files = append(files, distributorbundle.BundleFile{
			SourcePath: file.SourcePath,
			Path:       file.BundlePath,
		})
	}
	result, err := c.client.UploadFiles(ctx, distributorupload.UploadFilesOptions{
		PipelineID:     opts.PipelineID,
		ID:             opts.BundleID,
		Created:        opts.CreatedAt,
		IdempotencyKey: opts.IdempotencyKey,
		Files:          files,
	})
	if err != nil {
		return uploadFilesResult{}, err
	}
	return uploadFilesResult{
		RunID:  result.RunID,
		Status: result.Status,
	}, nil
}

func (c distributorUploadClient) Status(ctx context.Context, runID string) (runStatus, error) {
	status, err := c.client.Status(ctx, runID)
	if err != nil {
		return runStatus{}, err
	}
	return sanitizeRunStatus(runStatus{
		RunID:      status.RunID,
		PipelineID: status.PipelineID,
		Status:     status.Status,
		AcceptedAt: status.AcceptedAt,
		StartedAt:  status.StartedAt,
		FinishedAt: status.FinishedAt,
		Report:     append(json.RawMessage(nil), status.Report...),
		Error:      status.Error,
	}), nil
}

type uploadErrorContext struct {
	Endpoint       string
	PipelineID     string
	BundleID       string
	IdempotencyKey string
	SourcePaths    []string
	BundlePaths    []string
	Token          string
}

func wrapUploadError(err error, ctx uploadErrorContext) error {
	var conflict *distributorupload.IdempotencyConflictError
	isConflict := errors.As(err, &conflict)
	err = safeDistributorDiagnostic(err, ctx.Token)
	if isConflict {
		return &IdempotencyConflictError{
			Err: fmt.Errorf("upload distributor bundle %q to pipeline %q at endpoint %q with idempotency key %q from sources %q as bundle paths %q: idempotency conflict: %w", ctx.BundleID, ctx.PipelineID, ctx.Endpoint, ctx.IdempotencyKey, ctx.SourcePaths, ctx.BundlePaths, err),
		}
	}
	return fmt.Errorf("upload distributor bundle %q to pipeline %q at endpoint %q with idempotency key %q from sources %q as bundle paths %q: %w", ctx.BundleID, ctx.PipelineID, ctx.Endpoint, ctx.IdempotencyKey, ctx.SourcePaths, ctx.BundlePaths, err)
}

func safeDistributorDiagnostic(err error, token string) error {
	if err == nil {
		return nil
	}
	if errors.Is(err, errDistributorResponseTooLarge) {
		return errDistributorResponseTooLarge
	}
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return redactToken(err, token)
	}
	var httpErr *distributorupload.HTTPError
	if errors.As(err, &httpErr) {
		return &RemoteResponseError{StatusCode: httpErr.StatusCode, Retryable: httpErr.Retryable}
	}
	return errors.New("distributor request failed")
}

func sanitizeRunStatus(status runStatus) runStatus {
	status.Report = nil
	if status.Error != "" {
		status.Error = "distributor reported a failed run"
	}
	return status
}

func uploadSourcePaths(files []UploadFile) []string {
	paths := make([]string, 0, len(files))
	for _, file := range files {
		paths = append(paths, file.SourcePath)
	}
	return paths
}

func uploadBundlePaths(files []UploadFile) []string {
	paths := make([]string, 0, len(files))
	for _, file := range files {
		paths = append(paths, file.BundlePath)
	}
	return paths
}

func redactToken(err error, token string) error {
	if err == nil || token == "" {
		return err
	}
	return errors.New(redactTokenString(err.Error(), token))
}

func redactTokenString(value, token string) string {
	if token == "" {
		return value
	}
	return strings.ReplaceAll(value, token, "[redacted]")
}