source.go

package keepcurrent

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

	"github.com/mholt/archives"
)

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

// 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() {
		req.Header.Add("If-Modified-Since", ifNewerThan.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 {
		return nil, ErrUnmodified
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("unexpected HTTP status %v", resp.StatusCode)
	}
	etag := resp.Header.Get("ETag")
	if etag != "" {
		s.setETag(etag)
	}
	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()
			buf, err = io.ReadAll(f)
			if err != nil {
				return err
			}
			return errFound
		}
		return nil
	})

	if errors.Is(err, errFound) {
		return io.NopCloser(bytes.NewReader(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)
}

// 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 {
		return nil, err
	}
	if !ifNewerThan.IsZero() && ifNewerThan.Before(fi.ModTime()) {
		return nil, ErrUnmodified
	}
	var result io.ReadCloser
	result = f
	if s.preprocessor != nil {
		result, err = s.preprocessor(f)
		if err != nil {
			return nil, err
		}
	}
	return result, nil
}