ddc/internal/fetch/downloader.go

package fetch

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

	"ddc/internal/diskspace"
	"ddc/internal/files"
	"ddc/internal/guard"
)

const (
	userAgent          = "ddc/1.0"
	bufferSize         = 1 << 20
	stallTimeout       = 2 * time.Minute
	spaceCheckInterval = int64(256) << 20
)

var (
	ErrIncomplete = errors.New("download ended before the full file was received")
	ErrStalled    = errors.New("download stalled, no data received for too long")
)

type Progress interface {
	SetTotal(total int64)
	Add(delta int64)
	SetFileName(name string)
}

type Downloader struct {
	policy *guard.Policy
	client *http.Client
}

func NewDownloader(policy *guard.Policy) *Downloader {
	return &Downloader{policy: policy, client: policy.NewHTTPClient()}
}

func (downloader *Downloader) Download(ctx context.Context, rawURL, directory string, reservedBytes int64, progress Progress) (string, error) {
	target, err := downloader.policy.ValidateURL(rawURL)
	if err != nil {
		return "", err
	}

	ctx, cancel := context.WithCancelCause(ctx)
	defer cancel(nil)

	response, err := downloader.open(ctx, target.String())
	if err != nil {
		return "", err
	}
	defer response.Body.Close()

	fileName := FileNameFromResponse(response)
	progress.SetFileName(fileName)
	if err := diskspace.EnsureFree(directory, max(response.ContentLength, 0), reservedBytes); err != nil {
		return "", err
	}
	progress.SetTotal(max(response.ContentLength, 0))

	partPath := filepath.Join(directory, files.PartFileName)
	written, err := writePart(cancel, partPath, response.Body, progress, directory, reservedBytes)
	if err != nil {
		os.Remove(partPath)
		if errors.Is(context.Cause(ctx), ErrStalled) {
			return "", ErrStalled
		}
		return "", err
	}
	if response.ContentLength > 0 && written != response.ContentLength {
		os.Remove(partPath)
		return "", ErrIncomplete
	}

	finalPath := files.UniquePath(filepath.Join(directory, fileName))
	if err := os.Rename(partPath, finalPath); err != nil {
		os.Remove(partPath)
		return "", err
	}
	return finalPath, nil
}

func (downloader *Downloader) open(ctx context.Context, target string) (*http.Response, error) {
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
	if err != nil {
		return nil, err
	}
	request.Header.Set("User-Agent", userAgent)
	request.Header.Set("Accept-Encoding", "identity")

	response, err := downloader.client.Do(request)
	if err != nil {
		return nil, err
	}
	if response.StatusCode != http.StatusOK {
		response.Body.Close()
		return nil, fmt.Errorf("server responded with %s", response.Status)
	}
	return response, nil
}

func writePart(cancel context.CancelCauseFunc, partPath string, body io.Reader, progress Progress, directory string, reservedBytes int64) (int64, error) {
	file, err := os.OpenFile(partPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
	if err != nil {
		return 0, err
	}
	written, copyErr := copyBody(cancel, file, body, progress, directory, reservedBytes)
	closeErr := file.Close()
	if copyErr != nil {
		return written, copyErr
	}
	return written, closeErr
}

func copyBody(cancel context.CancelCauseFunc, destination io.Writer, body io.Reader, progress Progress, directory string, reservedBytes int64) (int64, error) {
	stallTimer := time.AfterFunc(stallTimeout, func() { cancel(ErrStalled) })
	defer stallTimer.Stop()

	buffer := make([]byte, bufferSize)
	var written, sinceSpaceCheck int64
	for {
		count, readErr := body.Read(buffer)
		if count > 0 {
			stallTimer.Reset(stallTimeout)
			if _, err := destination.Write(buffer[:count]); err != nil {
				return written, err
			}
			written += int64(count)
			sinceSpaceCheck += int64(count)
			progress.Add(int64(count))

			if sinceSpaceCheck >= spaceCheckInterval {
				sinceSpaceCheck = 0
				if err := diskspace.EnsureFree(directory, 0, reservedBytes); err != nil {
					return written, err
				}
			}
		}
		if errors.Is(readErr, io.EOF) {
			return written, nil
		}
		if readErr != nil {
			return written, readErr
		}
	}
}