EvoBGP/internal/pipeline/cdn_snapshot.go

package pipeline

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"evobgp/internal/store"
)

func cdnSourceKey(sourceID string) string {
	return "cdn:" + strings.TrimSpace(sourceID)
}

func cachedCDNPrefixRows(st store.Backend, tenantID, moduleID string, priorSnapshot []store.PrefixRow, sourceKey string) []store.PrefixRow {
	if cached := prefixRowsForSource(priorSnapshot, sourceKey); len(cached) > 0 {
		return cached
	}
	if st != nil {
		if snap, ok, _ := st.GetModulePrefixSnapshot(tenantID, moduleID); ok && snap != nil {
			if cached := prefixRowsForSource(snap.Prefixes, sourceKey); len(cached) > 0 {
				return cached
			}
		}
	}
	return nil
}

func mergeSnapshotDropSource(rows []store.PrefixRow, sourceKey string) []store.PrefixRow {
	if len(rows) == 0 {
		return nil
	}
	out := make([]store.PrefixRow, 0, len(rows))
	for _, row := range rows {
		if row.Source != sourceKey {
			out = append(out, row)
		}
	}
	return out
}

// mergeSnapshotKeepSkippedCDN keeps non-CDN rows and cdn:* rows whose source is still skipped
// (not in fetchedSourceIDs). Deleted sources (absent from allSourceIDs) are dropped.
func mergeSnapshotKeepSkippedCDN(rows []store.PrefixRow, fetchedSourceIDs, allSourceIDs []string) []store.PrefixRow {
	fetched := make(map[string]struct{}, len(fetchedSourceIDs))
	for _, id := range fetchedSourceIDs {
		fetched[cdnSourceKey(id)] = struct{}{}
	}
	keepCDN := make(map[string]struct{})
	for _, id := range allSourceIDs {
		k := cdnSourceKey(id)
		if _, ok := fetched[k]; !ok {
			keepCDN[k] = struct{}{}
		}
	}
	out := make([]store.PrefixRow, 0, len(rows))
	for _, row := range rows {
		src := strings.TrimSpace(row.Source)
		if strings.HasPrefix(src, "cdn:") {
			if _, ok := keepCDN[src]; !ok {
				continue
			}
		}
		out = append(out, row)
	}
	return out
}

// mergeAllCDNSourcesIntoModuleSnapshot replaces fetched CDN source rows in one write.
// skipped sources (errors with EVOBGP_CDN_PARTIAL_OK) keep their prior rows.
func mergeAllCDNSourcesIntoModuleSnapshot(st store.Backend, tenantID string, mod *store.Module, priorSnapshot []store.PrefixRow, cdnRows []store.PrefixRow, fetchedSourceIDs, allSourceIDs []string) error {
	if st == nil || mod == nil {
		return nil
	}
	unlock := st.LockModuleSnapshot(tenantID, mod.ID)
	defer unlock()
	var base []store.PrefixRow
	if len(priorSnapshot) > 0 {
		base = priorSnapshot
	} else if snap, ok, _ := st.GetModulePrefixSnapshot(tenantID, mod.ID); ok && snap != nil {
		base = snap.Prefixes
	}
	kept := mergeSnapshotKeepSkippedCDN(base, fetchedSourceIDs, allSourceIDs)
	merged := append(kept, cdnRows...)
	return persistModuleSnapshot(st, tenantID, mod, merged)
}

func cdnRowsFromParsed(mod *store.Module, src *store.CDNSource, pfxStrings []string) []store.PrefixRow {
	var rows []store.PrefixRow
	for _, p := range pfxStrings {
		p = strings.TrimSpace(p)
		if p == "" {
			continue
		}
		comm := src.CommunityID
		if comm == nil && mod != nil && mod.DefaultCommunityID != nil {
			c := *mod.DefaultCommunityID
			comm = &c
		}
		rows = append(rows, store.PrefixRow{Prefix: p, CommunityID: comm, Source: cdnSourceKey(src.ID)})
	}
	return rows
}

// mergeCDNSourceIntoModuleSnapshot replaces rows for one CDN source in the module snapshot.
func mergeCDNSourceIntoModuleSnapshot(st store.Backend, tenantID string, mod *store.Module, sourceID string, newRows []store.PrefixRow) error {
	if st == nil || mod == nil {
		return nil
	}
	unlock := st.LockModuleSnapshot(tenantID, mod.ID)
	defer unlock()
	sourceKey := cdnSourceKey(sourceID)
	var base []store.PrefixRow
	if snap, ok, _ := st.GetModulePrefixSnapshot(tenantID, mod.ID); ok && snap != nil {
		base = mergeSnapshotDropSource(snap.Prefixes, sourceKey)
	}
	merged := append(base, newRows...)
	return persistModuleSnapshot(st, tenantID, mod, merged)
}

func parseCDNBody(body string, src *store.CDNSource) ([]string, error) {
	pfxs, err := ExtractCIDRs(body, src.SourceKind, src.PrefixPath)
	if err != nil {
		return nil, err
	}
	out := make([]string, 0, len(pfxs))
	for _, pfx := range pfxs {
		out = append(out, pfx.String())
	}
	return out, nil
}

// fetchCDNSourceRows loads CDN prefixes without persisting the module snapshot (caller merges once).
func fetchCDNSourceRows(ctx context.Context, st store.Backend, hc *http.Client, tenantID, moduleID string, mod *store.Module, src *store.CDNSource, priorSnapshot []store.PrefixRow, now time.Time) ([]store.PrefixRow, error) {
	u := strings.TrimSpace(src.URL)
	if u == "" {
		return nil, nil
	}
	if _, err := ValidateCDNURL(u); err != nil {
		return nil, err
	}
	if err := ResolveCDNURLHost(ctx, u); err != nil {
		return nil, err
	}
	sourceKey := cdnSourceKey(src.ID)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
	if err != nil {
		return nil, err
	}
	if etag := strings.TrimSpace(src.Etag); etag != "" {
		req.Header.Set("If-None-Match", etag)
	}
	resp, err := upstreamHTTPDo(ctx, hc, req)
	if err != nil {
		return nil, fmt.Errorf("cdn fetch %s: %w", u, err)
	}

	if resp.StatusCode == http.StatusNotModified {
		if cached := cachedCDNPrefixRows(st, tenantID, moduleID, priorSnapshot, sourceKey); len(cached) > 0 {
			_ = resp.Body.Close()
			return cached, nil
		}
		_ = resp.Body.Close()
		req2, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
		if err != nil {
			return nil, err
		}
		resp, err = upstreamHTTPDo(ctx, hc, req2)
		if err != nil {
			return nil, fmt.Errorf("cdn fetch %s: %w", u, err)
		}
		if resp.StatusCode == http.StatusNotModified {
			_ = resp.Body.Close()
			return nil, fmt.Errorf("cdn url %s: 304 without cached prefixes", u)
		}
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, resp.Body)
		return nil, fmt.Errorf("cdn url %s: %s", u, resp.Status)
	}
	body, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	if err != nil {
		return nil, err
	}
	prefixStrs, err := parseCDNBody(string(body), src)
	if err != nil {
		return nil, fmt.Errorf("cdn parse %s: %w", u, err)
	}
	etag := strings.TrimSpace(resp.Header.Get("ETag"))
	patch := &store.CDNSourcePatch{}
	if etag != "" && etag != strings.TrimSpace(src.Etag) {
		e := etag
		patch.Etag = &e
	}
	refreshedAt := now
	patch.LastRefreshedAt = &refreshedAt
	_, _ = st.UpdateCDNSource(tenantID, moduleID, src.ID, patch)

	return cdnRowsFromParsed(mod, src, prefixStrs), nil
}