pocket-id/backend/internal/geolite/download.go

package geolite

import (
	"archive/tar"
	"bufio"
	"compress/gzip"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/oschwald/maxminddb-golang/v2"
)

const (
	// databaseFileName is the name of the database inside the archive published by MaxMind
	databaseFileName = "GeoLite2-City.mmdb"
	// maxDatabaseSize is the largest (decompressed) database we accept
	maxDatabaseSize = 300 << 20 // 300 MB
)

// downloadDatabase downloads the GeoLite2 City database and puts it at targetPath
// The database is streamed to a temporary file next to the target and moved into place only once it has been verified, atomically
func downloadDatabase(ctx context.Context, httpClient *http.Client, downloadURL string, licenseKey string, targetPath string) error {
	// When downloadURL contains a "%s" placeholder, it is replaced with the license key
	if strings.Contains(downloadURL, "%s") {
		downloadURL = fmt.Sprintf(downloadURL, licenseKey)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
	if err != nil {
		return fmt.Errorf("failed to create request: %w", err)
	}

	res, err := httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("failed to download database: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("failed to download database, received HTTP %d", res.StatusCode)
	}

	err = writeDatabase(res.Body, targetPath)
	if err != nil {
		return err
	}

	return nil
}

// writeDatabase extracts the database from a downloaded body and puts it at targetPath
func writeDatabase(body io.Reader, targetPath string) error {
	baseDir := filepath.Dir(targetPath)
	err := os.MkdirAll(baseDir, 0700)
	if err != nil {
		return fmt.Errorf("failed to create the database directory: %w", err)
	}

	tmpFile, err := os.CreateTemp(baseDir, "geolite.*.mmdb.tmp")
	if err != nil {
		return fmt.Errorf("failed to create temporary database file: %w", err)
	}
	tmpName := tmpFile.Name()

	// Remove the temporary file unless it has been moved into place
	moved := false
	defer func() {
		tmpFile.Close()
		if !moved {
			os.Remove(tmpName)
		}
	}()

	err = extractDatabase(body, tmpFile)
	if err != nil {
		return fmt.Errorf("failed to extract database: %w", err)
	}

	err = tmpFile.Close()
	if err != nil {
		return fmt.Errorf("failed to write database file: %w", err)
	}

	// Make sure the database isn't corrupted before it replaces the one currently in place
	db, err := maxminddb.Open(tmpName)
	if err != nil {
		return fmt.Errorf("failed to open downloaded database: %w", err)
	}
	_ = db.Close()

	err = os.Rename(tmpName, targetPath)
	if err != nil {
		return fmt.Errorf("failed to replace database file: %w", err)
	}
	moved = true

	return nil
}

// extractDatabase copies the raw MaxMind DB file out of a downloaded body and into dst
// The body is either the gzipped tarball published by MaxMind, or the database file itself
func extractDatabase(body io.Reader, dst io.Writer) error {
	// A buffered reader lets the gzip magic number be checked without consuming it
	reader := bufio.NewReader(body)

	magic, err := reader.Peek(2)
	if err != nil {
		return fmt.Errorf("failed to read magic number: %w", err)
	}

	// If the body doesn't start with the gzip magic number, assume it's a plain database file
	// Gosec returns false positive for "G602: slice index out of range"
	//nolint:gosec
	if magic[0] != 0x1f || magic[1] != 0x8b {
		return copyDatabase(dst, reader)
	}

	gzr, err := gzip.NewReader(reader)
	if err != nil {
		return fmt.Errorf("failed to create gzip reader: %w", err)
	}
	defer gzr.Close()

	tarReader := tar.NewReader(gzr)
	for {
		header, err := tarReader.Next()
		if errors.Is(err, io.EOF) {
			break
		} else if err != nil {
			return fmt.Errorf("failed to read tar archive: %w", err)
		}

		// The archive contains the database in a versioned folder, alongside other files such as the license
		if header.Typeflag != tar.TypeReg || path.Base(header.Name) != databaseFileName {
			continue
		}

		err = copyDatabase(dst, tarReader)
		if err != nil {
			return err
		}
		return nil
	}

	return errors.New(databaseFileName + " not found in archive")
}

// copyDatabase streams the database into dst, refusing anything larger than maxDatabaseSize
func copyDatabase(dst io.Writer, src io.Reader) error {
	return copyDatabaseWithLimit(dst, src, maxDatabaseSize)
}

// copyDatabaseWithLimit is copyDatabase with an explicit limit, which lets tests exercise the limit without streaming hundreds of megabytes
func copyDatabaseWithLimit(dst io.Writer, src io.Reader, limit int64) error {
	// Copy one byte more than the limit, so content that is exactly at the limit can be told apart from content that exceeds it
	written, err := io.Copy(dst, io.LimitReader(src, limit+1))
	if err != nil {
		return fmt.Errorf("failed to read database: %w", err)
	}

	if written > limit {
		return errors.New("database size exceeds maximum allowed limit")
	}

	return nil
}