sow-tools/internal/depot/remote.go

package depot

import (
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"time"
)

// NewBackend returns the backend for name ("local"|"cdn"|"bunny").
// localRoot is used only for "local". Fails closed: "bunny" with empty
// StorageHost or empty ReadKey returns an error (never prompts);
// unknown name returns an error.
func NewBackend(name, localRoot string, cfg Config) (Backend, error) {
	switch name {
	case "local":
		if localRoot == "" {
			return nil, errors.New("local backend requires a non-empty root")
		}
		return &LocalBackend{Root: localRoot}, nil
	case "cdn":
		return &httpBackend{
			name:   "cdn",
			client: newHTTPClient(cfg),
			cfg:    cfg,
		}, nil
	case "bunny":
		if cfg.StorageHost == "" {
			return nil, errors.New("bunny backend requires BUNNY_STORAGE_HOST")
		}
		if cfg.ReadKey == "" {
			return nil, errors.New("bunny backend requires BUNNY_STORAGE_READ_PASSWORD or BUNNY_STORAGE_PASSWORD")
		}
		return &httpBackend{
			name:   "bunny",
			client: newHTTPClient(cfg),
			cfg:    cfg,
		}, nil
	default:
		return nil, fmt.Errorf("unknown backend %q", name)
	}
}

// newHTTPClient builds an IPv4-only client per spec (HEAD probes are banned;
// IPv6 dial hazards are out of scope for this depot).
func newHTTPClient(cfg Config) *http.Client {
	transport := &http.Transport{
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			d := net.Dialer{Timeout: cfg.ConnectTimeout}
			return d.DialContext(ctx, "tcp4", addr)
		},
		MaxIdleConnsPerHost: cfg.ProbeJobs,
	}
	return &http.Client{Transport: transport, Timeout: cfg.ProbeMaxTime}
}

// httpBackend implements Backend for both cdn (read-only) and bunny
// (read/write) over HTTP, sharing probe/get/put logic.
type httpBackend struct {
	name   string
	client *http.Client
	cfg    Config
}

func (b *httpBackend) Name() string { return b.name }

func (b *httpBackend) storageURL(sha string) string { return b.keyURL(BlobKey(sha)) }

func (b *httpBackend) cdnURL(sha string) string {
	return fmt.Sprintf("%s/%s", b.cfg.CDNBase, BlobKey(sha))
}

// probeRetrySleep is the backoff sleeper for transient probe retries;
// tests stub it (same pattern as sweep.go's confirmSleep).
var probeRetrySleep = time.Sleep

// rangeProbe issues GET <url> with Range: bytes=0-0 and classifies the
// response. Never HEAD (banned by spec). Transient outcomes (transport
// error, or a status that is neither 2xx nor 404/410) retry up to 2 more
// times with linear backoff (1s, 2s) — the CDN edge throttles sustained
// sweeps; the bash this ports absorbed that with curl --retry 2.
func (b *httpBackend) rangeProbe(ctx context.Context, url string, headers map[string]string) (ProbeState, bool, error) {
	for attempt := 0; ; attempt++ {
		state, transient, err := b.rangeProbeOnce(ctx, url, headers)
		if !transient || attempt >= 2 {
			return state, transient, err
		}
		probeRetrySleep(time.Duration(attempt+1) * time.Second)
	}
}

func (b *httpBackend) rangeProbeOnce(ctx context.Context, url string, headers map[string]string) (ProbeState, bool, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return Unconfirmed, true, err
	}
	req.Header.Set("Range", "bytes=0-0")
	for k, v := range headers {
		req.Header.Set(k, v)
	}

	resp, err := b.client.Do(req)
	if err != nil {
		return Unconfirmed, true, err
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, resp.Body)

	switch {
	case resp.StatusCode >= 200 && resp.StatusCode < 300:
		return Present, false, nil
	case resp.StatusCode == 404 || resp.StatusCode == 410:
		return Absent, false, nil
	default:
		return Unconfirmed, true, nil
	}
}

// Probe returns the existence state of sha at this backend.
func (b *httpBackend) Probe(ctx context.Context, sha string) (ProbeState, bool, error) {
	if b.name == "cdn" {
		return b.rangeProbe(ctx, b.cdnURL(sha), nil)
	}
	return b.rangeProbe(ctx, b.storageURL(sha), map[string]string{"AccessKey": b.cfg.ReadKey})
}

// Put uploads src for sha. cdn is read-only. A depot object is named after the
// sha256 of its own bytes, so the key's sha doubles as the Checksum header.
func (b *httpBackend) Put(ctx context.Context, sha, src string) error {
	if b.name == "cdn" {
		return errors.New("cdn backend is read-only")
	}
	if b.cfg.WriteKey == "" {
		return errors.New("bunny backend requires BUNNY_STORAGE_PASSWORD to write")
	}

	state, _, err := b.Probe(ctx, sha)
	if err != nil {
		return err
	}
	if state == Present {
		return nil
	}

	return b.putFile(ctx, BlobKey(sha), src, sha)
}

// Get fetches sha into dest via temp file + rename, re-hashing and deleting
// on mismatch. Mirrors bash _depot_bunny_get: try CDN first, fall back to
// storage with ReadKey.
func (b *httpBackend) Get(ctx context.Context, sha, dest string) error {
	destDir := filepath.Dir(dest)
	if err := os.MkdirAll(destDir, 0755); err != nil {
		return err
	}

	tmpFile, err := os.CreateTemp(destDir, ".tmp-")
	if err != nil {
		return err
	}
	tmpName := tmpFile.Name()
	defer os.Remove(tmpName)

	resp, err := b.fetch(ctx, sha)
	if err != nil {
		tmpFile.Close()
		return err
	}
	defer resp.Body.Close()

	hasher := sha256.New()
	tee := io.TeeReader(resp.Body, hasher)
	_, err = io.Copy(tmpFile, tee)
	tmpFile.Close()
	if err != nil {
		return err
	}

	gotHash := fmt.Sprintf("%x", hasher.Sum(nil))
	if gotHash != sha {
		os.Remove(tmpName)
		return fmt.Errorf("hash mismatch: expected %s, got %s", sha, gotHash)
	}

	return os.Rename(tmpName, dest)
}

// fetch tries the CDN URL first (no auth), falling back to the storage URL
// with ReadKey on any non-2xx response or transport error.
func (b *httpBackend) fetch(ctx context.Context, sha string) (*http.Response, error) {
	if b.name != "cdn" {
		if resp, err := b.get(ctx, b.cdnURL(sha), nil); err == nil && resp.StatusCode >= 200 && resp.StatusCode < 300 {
			return resp, nil
		} else if err == nil {
			_, _ = io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
		}
		resp, err := b.get(ctx, b.storageURL(sha), map[string]string{"AccessKey": b.cfg.ReadKey})
		if err != nil {
			return nil, err
		}
		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
			_, _ = io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
			return nil, fmt.Errorf("bunny get %s: unexpected status %d", sha, resp.StatusCode)
		}
		return resp, nil
	}

	resp, err := b.get(ctx, b.cdnURL(sha), nil)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		_, _ = io.Copy(io.Discard, resp.Body)
		resp.Body.Close()
		return nil, fmt.Errorf("cdn get %s: unexpected status %d", sha, resp.StatusCode)
	}
	return resp, nil
}

func (b *httpBackend) get(ctx context.Context, url string, headers map[string]string) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	for k, v := range headers {
		req.Header.Set(k, v)
	}
	return b.client.Do(req)
}