RenewCMS/internal/infrastructure/update/httpAssetFetcher.go

package update

import (
	domain "RenewCMS/internal/domain/update"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"
)

const (
	maxBinaryBytes       int64 = 200 << 20
	maxChecksumsBytes    int64 = 64 << 10
	maxSignatureBytes    int64 = 1 << 10
	assetDownloadTimeout       = 10 * time.Minute
	maxAssetRedirects          = 5
	assetFilePermission        = 0o644
	signatureSuffix            = ".sig"
)

var (
	errAssetURLRefused  = errors.New("the asset url is not https on the release feed host")
	errAssetStatus      = errors.New("the asset server answered with an error status")
	errAssetTooLarge    = errors.New("the asset is larger than its size cap")
	errTooManyRedirects = errors.New("the asset url redirected too many times")
)

type FileOpener func(path string) (io.WriteCloser, error)

type HTTPAssetFetcher struct {
	allowedHost   string
	allowInsecure bool
	sizeCap       int64
	timeout       time.Duration
	transport     http.RoundTripper
	openFile      FileOpener
}

var _ domain.AssetFetcher = &HTTPAssetFetcher{}

type FetcherOption func(*HTTPAssetFetcher)

func WithInsecureScheme() FetcherOption {
	return func(f *HTTPAssetFetcher) { f.allowInsecure = true }
}

func WithSizeCap(bytes int64) FetcherOption {
	return func(f *HTTPAssetFetcher) { f.sizeCap = bytes }
}

func WithTimeout(timeout time.Duration) FetcherOption {
	return func(f *HTTPAssetFetcher) { f.timeout = timeout }
}

func WithTransport(transport http.RoundTripper) FetcherOption {
	return func(f *HTTPAssetFetcher) { f.transport = transport }
}

func WithFileOpener(opener FileOpener) FetcherOption {
	return func(f *HTTPAssetFetcher) { f.openFile = opener }
}

func NewHTTPAssetFetcher(feedURL string, options ...FetcherOption) *HTTPAssetFetcher {
	fetcher := &HTTPAssetFetcher{
		allowedHost: domain.HostOf(feedURL),
		timeout:     assetDownloadTimeout,
		openFile:    createAssetFile,
	}
	for _, option := range options {
		option(fetcher)
	}
	return fetcher
}

func (f *HTTPAssetFetcher) Fetch(asset domain.Asset, destination string, report domain.ProgressFunc) error {
	if !f.isAllowed(asset.URL) {
		return errAssetURLRefused
	}
	ctx, cancel := context.WithTimeout(context.Background(), f.timeout)
	defer cancel()
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, asset.URL, nil)
	if err != nil {
		return err
	}
	response, err := f.client().Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("%w: %d", errAssetStatus, response.StatusCode)
	}
	limit := f.capFor(asset.Name)
	if response.ContentLength > limit {
		return errAssetTooLarge
	}
	if err := f.writeBody(response.Body, destination, limit, asset.Size, report); err != nil {
		_ = os.Remove(destination)
		return err
	}
	return nil
}

func (f *HTTPAssetFetcher) writeBody(body io.Reader, destination string, limit int64, declaredSize int64, report domain.ProgressFunc) error {
	file, err := f.openFile(destination)
	if err != nil {
		return err
	}
	counter := &progressWriter{destination: file, total: declaredSize, report: report}
	written, copyErr := io.Copy(counter, io.LimitReader(body, limit+1))
	closeErr := file.Close()
	if copyErr != nil {
		return copyErr
	}
	if closeErr != nil {
		return closeErr
	}
	if written > limit {
		return errAssetTooLarge
	}
	return nil
}

func (f *HTTPAssetFetcher) isAllowed(rawURL string) bool {
	parsed, err := url.Parse(rawURL)
	if err != nil || f.allowedHost == "" || parsed.Host != f.allowedHost {
		return false
	}
	return parsed.Scheme == "https" || (f.allowInsecure && parsed.Scheme == "http")
}

func (f *HTTPAssetFetcher) client() *http.Client {
	return &http.Client{
		Transport: f.transport,
		CheckRedirect: func(next *http.Request, via []*http.Request) error {
			if len(via) >= maxAssetRedirects {
				return errTooManyRedirects
			}
			if !f.isAllowed(next.URL.String()) {
				return errAssetURLRefused
			}
			return nil
		},
	}
}

func (f *HTTPAssetFetcher) capFor(name string) int64 {
	if f.sizeCap > 0 {
		return f.sizeCap
	}
	if name == domain.ChecksumsAssetName {
		return maxChecksumsBytes
	}
	if strings.HasSuffix(name, signatureSuffix) {
		return maxSignatureBytes
	}
	return maxBinaryBytes
}

type progressWriter struct {
	destination io.Writer
	total       int64
	done        int64
	report      domain.ProgressFunc
}

func (w *progressWriter) Write(chunk []byte) (int, error) {
	written, err := w.destination.Write(chunk)
	w.done += int64(written)
	w.report(w.done, w.total)
	return written, err
}

func createAssetFile(path string) (io.WriteCloser, error) {
	if err := os.MkdirAll(filepath.Dir(path), os.ModePerm); err != nil {
		return nil, err
	}
	return os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, assetFilePermission)
}