source.go

package keepcurrent

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"sync"
	"time"

	"github.com/mholt/archives"
)

type webSource struct {
	url    string
	etag   string
	mx     sync.RWMutex
	client *http.Client
}

// drainClose discards any remaining body and closes it. net/http only returns a
// connection to the keep-alive pool once its response body has been read to EOF,
// so error/not-modified responses we don't hand to the caller must be drained.
func drainClose(rc io.ReadCloser) {
	_, _ = io.Copy(io.Discard, rc)
	_ = rc.Close()
}

// FromWeb constructs a source from the given URL.
func FromWeb(url string) Source {
	return FromWebWithClient(url, http.DefaultClient)
}

// FromWebWithClient is the same as FromWeb but with a custom http.Client
func FromWebWithClient(url string, client *http.Client) Source {
	return &webSource{url: url, client: client}
}

// Fetch implements the Source interface
func (s *webSource) Fetch(ifNewerThan time.Time) (io.ReadCloser, error) {
	req, err := http.NewRequest(http.MethodGet, s.url, nil)
	if err != nil {
		return nil, err
	}
	if !ifNewerThan.IsZero() {
		// http.TimeFormat is GMT, so the time must be in UTC first; formatted in
		// local time it would claim a different instant on any host not on UTC.
		req.Header.Add("If-Modified-Since", ifNewerThan.UTC().Format(http.TimeFormat))
	}
	if s.getETag() != "" {
		req.Header.Add("If-None-Match", s.etag)
	}
	resp, err := s.client.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode == http.StatusNotModified {
		// Drain to EOF then close so net/http can return the connection to the
		// pool for keep-alive reuse (it won't reuse one whose body wasn't fully
		// read). 304 carries no body, so this is effectively just a close here.
		// We hand the body to the caller only on the success path below.
		drainClose(resp.Body)
		return nil, ErrUnmodified
	}
	if resp.StatusCode != http.StatusOK {
		drainClose(resp.Body)
		return nil, fmt.Errorf("unexpected HTTP status %v", resp.StatusCode)
	}
	etag := resp.Header.Get("ETag")
	if etag != "" {
		s.setETag(etag)
	}
	if resp.ContentLength >= 0 {
		// Surface the Content-Length so the Runner can pre-size its read buffer.
		return sizedReadCloser{ReadCloser: resp.Body, n: resp.ContentLength}, nil
	}
	return resp.Body, nil
}

func (s *webSource) getETag() string {
	s.mx.RLock()
	defer s.mx.RUnlock()
	return s.etag
}

func (s *webSource) setETag(etag string) {
	s.mx.Lock()
	s.etag = etag
	s.mx.Unlock()
}

type tarGzSource struct {
	s            Source
	expectedName string
}

// FromTarGz wraps a source to decompress one specific file from the gzipped
// tarball.
func FromTarGz(s Source, expectedName string) Source {
	return &tarGzSource{s, expectedName}
}

var errFound = errors.New("found")

func (s *tarGzSource) Fetch(ifNewerThan time.Time) (io.ReadCloser, error) {
	rc, err := s.s.Fetch(ifNewerThan)
	if err != nil {
		return nil, err
	}
	defer rc.Close()

	// archives.CompressedArchive.Extract() reads ca.Extraction (not ca.Archival),
	// and returns "no extraction format" if it's nil. Setting Archival here was a
	// bug that made every Fetch() fail silently.
	format := archives.CompressedArchive{
		Compression: archives.Gz{},
		Extraction:  archives.Tar{},
	}

	var buf []byte
	err = format.Extract(context.Background(), rc, func(ctx context.Context, info archives.FileInfo) error {
		// NameInArchive is the full stored path (e.g. "GeoLite2-City_20260116/GeoLite2-City.mmdb").
		// Callers pass the basename, so compare on basename — mirrors the behavior of the
		// archiver/v3-based implementation this replaced.
		if path.Base(info.NameInArchive) == s.expectedName {
			f, err := info.Open()
			if err != nil {
				return err
			}
			defer f.Close()
			// Wrap the entry in a size-aware reader so readAll pre-sizes the
			// buffer from the archive entry's uncompressed size — extracting a
			// large file (e.g. a ~75MB mmdb) becomes a single allocation rather
			// than the reallocation churn of io.ReadAll.
			buf, err = readAll(sizedReadCloser{ReadCloser: f, n: info.Size()})
			if err != nil {
				return err
			}
			return errFound
		}
		return nil
	})

	if errors.Is(err, errFound) {
		// Return a size-aware reader so the Runner's read can also be pre-sized.
		return bytesReadCloser(buf), nil
	}
	if err != nil {
		return nil, err
	}
	return nil, fmt.Errorf("file %q not found in archive", s.expectedName)
}

type fileSource struct {
	path         string
	preprocessor func(io.ReadCloser) (io.ReadCloser, error)
}

// fileBackedReadCloser ties a preprocessed stream's lifetime to the file it was
// derived from, closing both when Close is called. A double close of the file
// (when the preprocessor's stream already closes it) is tolerated.
type fileBackedReadCloser struct {
	io.ReadCloser
	f *os.File
}

func (r fileBackedReadCloser) Close() error {
	err := r.ReadCloser.Close()
	if cerr := r.f.Close(); cerr != nil && !errors.Is(cerr, os.ErrClosed) && err == nil {
		err = cerr
	}
	return err
}

// FromFile constructs a source from the given file path.
func FromFile(path string) Source {
	return &fileSource{path, nil}
}

// FromFileWithPreprocessor constructs a source from the given file path, while modifying the file data using preprocessor function
func FromFileWithPreprocessor(path string, preprocessor func(io.ReadCloser) (io.ReadCloser, error)) Source {
	return &fileSource{path, preprocessor}
}

func (s *fileSource) Fetch(ifNewerThan time.Time) (io.ReadCloser, error) {
	f, err := os.Open(s.path)
	if err != nil {
		return nil, err
	}
	fi, err := f.Stat()
	if err != nil {
		f.Close()
		return nil, err
	}
	// Fetch's contract is "fetch only if modified since ifNewerThan", so we
	// signal ErrUnmodified when the file is no newer than the cutoff — mirroring
	// byteSource.Fetch in the tests. (Earlier this comparison was inverted.)
	if !ifNewerThan.IsZero() && !fi.ModTime().After(ifNewerThan) {
		f.Close()
		return nil, ErrUnmodified
	}
	if s.preprocessor != nil {
		// A preprocessor may change the byte count, so we can't declare a size
		// up front; readAll falls back to io.ReadAll for the preprocessed stream.
		result, err := s.preprocessor(f)
		if err != nil {
			f.Close()
			return nil, err
		}
		// The preprocessor may hand back a stream that doesn't close the file it
		// wraps (e.g. io.NopCloser), so couple Close to the underlying file to
		// avoid leaking the descriptor.
		return modTimeReadCloser{ReadCloser: fileBackedReadCloser{ReadCloser: result, f: f}, mod: fi.ModTime()}, nil
	}
	// Surface the file size so the Runner's readAll pre-sizes its buffer instead
	// of growing by reallocation. The cached payloads read here are exactly the
	// large ones that churn — e.g. geo loads a ~75MB MaxMind mmdb from disk via
	// InitFrom(FromFile(...)) at startup.
	return modTimeReadCloser{ReadCloser: sizedReadCloser{ReadCloser: f, n: fi.Size()}, mod: fi.ModTime()}, nil
}

// modTimeReadCloser is a file's contents with the file's modification time, so
// a Runner records how old the data it loaded is rather than when it loaded it.
// Without this, a Runner initialised from an older cached file would ask its
// web source only for data newer than the moment it started, and never fetch
// an update published between the file's download and that moment.
type modTimeReadCloser struct {
	io.ReadCloser
	mod time.Time
}

func (r modTimeReadCloser) sourceModTime() time.Time { return r.mod }

// size passes through the wrapped reader's size, or -1 (unknown) when it has
// none, so readAll still pre-sizes a plain file read.
func (r modTimeReadCloser) size() int64 {
	if n, ok := knownSize(r.ReadCloser); ok {
		return n
	}
	return -1
}

// writesBack reports whether sink s would copy source from's file onto itself
// unchanged, which a sync skips: the rewrite changes nothing but the file's
// modification time, which moves to now, so the next Runner to start from it
// would take an old file for a fresh one. A preprocessor on either side can
// change the contents, so those writes still happen.
func writesBack(from Source, s Sink) bool {
	src, ok := from.(*fileSource)
	if !ok || src.preprocessor != nil {
		return false
	}
	dst, ok := s.(*fileSink)
	if !ok || dst.preprocessor != nil {
		return false
	}
	// Two paths can name one file (a symlink, a relative path): compare the
	// files themselves when both exist.
	if srcInfo, err := os.Stat(src.path); err == nil {
		if dstInfo, err := os.Stat(dst.path); err == nil {
			return os.SameFile(srcInfo, dstInfo)
		}
	}
	a, errA := filepath.Abs(src.path)
	b, errB := filepath.Abs(dst.path)
	return errA == nil && errB == nil && a == b
}