artifactapi/internal/virtual/engine.go

package virtual

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"sync"

	"git.unkin.net/unkin/artifactapi/internal/database"
	"git.unkin.net/unkin/artifactapi/internal/provider"
	"git.unkin.net/unkin/artifactapi/internal/proxy"
	"git.unkin.net/unkin/artifactapi/pkg/models"
)

type Engine struct {
	db          *database.DB
	proxyEngine *proxy.Engine
}

func NewEngine(db *database.DB, proxyEngine *proxy.Engine) *Engine {
	return &Engine{db: db, proxyEngine: proxyEngine}
}

func (e *Engine) Fetch(ctx context.Context, virt models.Virtual, path string, proxyBaseURL string) ([]byte, string, error) {
	if virt.PackageType == models.PackageRPM {
		return e.fetchRPM(ctx, virt, path)
	}

	merger, err := GetMerger(virt.PackageType)
	if err != nil {
		return nil, "", fmt.Errorf("unsupported virtual type %q: %w", virt.PackageType, err)
	}

	members, err := e.fetchMemberIndexes(ctx, virt, path)
	if err != nil {
		return nil, "", err
	}

	if len(members) == 0 {
		return nil, "", fmt.Errorf("no members reachable for virtual %q", virt.Name)
	}

	merged, err := merger.MergeIndexes(members, proxyBaseURL)
	if err != nil {
		return nil, "", fmt.Errorf("merge indexes: %w", err)
	}

	contentType := "application/octet-stream"
	switch virt.PackageType {
	case models.PackageHelm:
		contentType = "text/yaml"
	case models.PackagePyPI:
		contentType = "text/html"
	}

	return merged, contentType, nil
}

func (e *Engine) fetchMemberIndexes(ctx context.Context, virt models.Virtual, path string) ([]MemberIndex, error) {
	type result struct {
		index MemberIndex
		err   error
	}

	results := make([]result, len(virt.Members))
	var wg sync.WaitGroup

	for i, memberName := range virt.Members {
		wg.Add(1)
		go func(idx int, name string) {
			defer wg.Done()

			remote, err := e.db.GetRemote(ctx, name)
			if err != nil {
				results[idx] = result{err: fmt.Errorf("remote %q: %w", name, err)}
				return
			}

			if remote.RepoType == models.RepoTypeLocal {
				body, err := e.fetchLocalIndex(ctx, *remote, path)
				if err != nil {
					results[idx] = result{err: fmt.Errorf("local index %q: %w", name, err)}
					return
				}
				results[idx] = result{index: MemberIndex{RemoteName: name, RepoType: remote.RepoType, BaseURL: remote.BaseURL, Body: body}}
				return
			}

			prov, err := provider.Get(remote.PackageType)
			if err != nil {
				results[idx] = result{err: fmt.Errorf("provider %q: %w", remote.PackageType, err)}
				return
			}

			fetchResult, err := e.proxyEngine.Fetch(ctx, *remote, path, prov)
			if err != nil {
				results[idx] = result{err: fmt.Errorf("fetch %q/%s: %w", name, path, err)}
				return
			}
			defer fetchResult.Reader.Close()

			body, err := io.ReadAll(fetchResult.Reader)
			if err != nil {
				results[idx] = result{err: fmt.Errorf("read %q: %w", name, err)}
				return
			}

			results[idx] = result{index: MemberIndex{RemoteName: name, RepoType: remote.RepoType, BaseURL: remote.BaseURL, Body: body}}
		}(i, memberName)
	}

	wg.Wait()

	var members []MemberIndex
	for _, r := range results {
		if r.err != nil {
			slog.Warn("virtual member fetch failed", "error", r.err)
			continue
		}
		members = append(members, r.index)
	}

	return members, nil
}

func (e *Engine) fetchLocalIndex(ctx context.Context, remote models.Remote, path string) ([]byte, error) {
	prov, err := provider.Get(remote.PackageType)
	if err != nil {
		return nil, fmt.Errorf("no provider for %q: %w", remote.PackageType, err)
	}

	indexer, ok := prov.(provider.LocalIndexer)
	if !ok {
		return nil, fmt.Errorf("provider %q does not support local index generation", remote.PackageType)
	}

	return indexer.GenerateLocalIndex(ctx, e.db, remote.Name, path)
}

var ErrNotFound = errors.New("not found")

// MemberRedirect maps an rpm virtual package path (<member>/<href>, as written
// by MergeRPM) to the owning member's route. ok is false for any other path.
func (e *Engine) MemberRedirect(ctx context.Context, virt models.Virtual, path, proxyBaseURL string) (string, bool) {
	if virt.PackageType != models.PackageRPM || strings.HasPrefix(path, "repodata/") {
		return "", false
	}
	name, rest, found := strings.Cut(path, "/")
	if !found || !slices.Contains(virt.Members, name) {
		return "", false
	}
	remote, err := e.db.GetRemote(ctx, name)
	if err != nil {
		return "", false
	}
	return fmt.Sprintf("%s/api/v1/%s/%s/%s", strings.TrimRight(proxyBaseURL, "/"), remote.RepoType, name, rest), true
}

func (e *Engine) fetchRPM(ctx context.Context, virt models.Virtual, path string) ([]byte, string, error) {
	if !strings.HasPrefix(path, "repodata/") {
		return nil, "", ErrNotFound
	}

	results := make([]*RPMMember, len(virt.Members))
	var wg sync.WaitGroup
	for i, name := range virt.Members {
		wg.Add(1)
		go func() {
			defer wg.Done()
			m, err := e.fetchRPMMember(ctx, name)
			if err != nil {
				slog.Warn("virtual member fetch failed", "virtual", virt.Name, "member", name, "error", err)
				return
			}
			results[i] = m
		}()
	}
	wg.Wait()

	var members []RPMMember
	for _, m := range results {
		if m != nil {
			members = append(members, *m)
		}
	}
	if len(members) == 0 {
		return nil, "", fmt.Errorf("no members reachable for virtual %q", virt.Name)
	}

	// ponytail: merges on every repodata request; cache keyed on member repomd checksums if member count/size makes this slow.
	repo, err := MergeRPM(members)
	if err != nil {
		return nil, "", fmt.Errorf("merge rpm repodata: %w", err)
	}
	if path == "repodata/repomd.xml" {
		return repo.Repomd, "application/xml", nil
	}
	if body, ok := repo.Files[path]; ok {
		return body, "application/gzip", nil
	}
	return nil, "", ErrNotFound
}

func (e *Engine) fetchRPMMember(ctx context.Context, name string) (*RPMMember, error) {
	remote, err := e.db.GetRemote(ctx, name)
	if err != nil {
		return nil, fmt.Errorf("remote %q: %w", name, err)
	}
	repomd, err := e.fetchMemberPath(ctx, *remote, "repodata/repomd.xml")
	if err != nil {
		return nil, err
	}
	locs, ts, err := parseRepomd(repomd)
	if err != nil {
		return nil, err
	}
	m := &RPMMember{RemoteName: name, Timestamp: ts, Data: map[string][]byte{}}
	for t, href := range locs {
		raw, err := e.fetchMemberPath(ctx, *remote, href)
		if err != nil {
			return nil, err
		}
		if m.Data[t], err = decompress(href, raw); err != nil {
			return nil, fmt.Errorf("decompress %s: %w", href, err)
		}
	}
	return m, nil
}

// fetchMemberPath reads one path from a member the same way its own route
// would serve it: local index, synthesized remote (e.g. github_rpm), or proxy.
func (e *Engine) fetchMemberPath(ctx context.Context, remote models.Remote, path string) ([]byte, error) {
	prov, err := provider.Get(remote.PackageType)
	if err != nil {
		return nil, fmt.Errorf("provider %q: %w", remote.PackageType, err)
	}

	var serve func(w http.ResponseWriter, r *http.Request) bool
	if remote.RepoType == models.RepoTypeLocal {
		indexer, ok := prov.(provider.LocalIndexer)
		if !ok {
			return nil, fmt.Errorf("provider %q does not serve a local index", remote.PackageType)
		}
		serve = func(w http.ResponseWriter, r *http.Request) bool {
			return indexer.ServeLocalIndex(w, r, e.db, remote.Name, path)
		}
	} else if rs, ok := prov.(provider.RemoteServer); ok {
		serve = func(w http.ResponseWriter, r *http.Request) bool {
			return rs.ServeRemote(w, r, remote, path, "", e.db)
		}
	}

	if serve != nil {
		rec := httptest.NewRecorder()
		if serve(rec, httptest.NewRequestWithContext(ctx, http.MethodGet, "/"+path, nil)) {
			if rec.Code != http.StatusOK {
				return nil, fmt.Errorf("%s/%s: status %d", remote.Name, path, rec.Code)
			}
			return rec.Body.Bytes(), nil
		}
		if remote.RepoType == models.RepoTypeLocal {
			return nil, fmt.Errorf("%s/%s: %w", remote.Name, path, ErrNotFound)
		}
	}

	res, err := e.proxyEngine.Fetch(ctx, remote, path, prov)
	if err != nil {
		return nil, fmt.Errorf("fetch %s/%s: %w", remote.Name, path, err)
	}
	defer res.Reader.Close()
	return io.ReadAll(res.Reader)
}