ddc/internal/fetch/downloader.go

package fetch

import (
	"context"
	"errors"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"

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

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)
	SetDone(done int64)
	Add(delta int64)
	SetFileName(name string)
	SetNotice(text string)
}

type Request struct {
	URL           string
	Directory     string
	FileName      string
	ReservedBytes int64
	Connections   int
	State         *State
	Progress      Progress
	Checkpoint    func(State)
	Logger        *slog.Logger
}

func (request Request) logger() *slog.Logger {
	if request.Logger != nil {
		return request.Logger
	}
	return slog.Default()
}

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, request Request) (string, error) {
	target, err := downloader.policy.ValidateURL(request.URL)
	if err != nil {
		return "", err
	}

	source := &remoteFile{url: target.String()}
	partPath := filepath.Join(request.Directory, files.PartFileName)
	state, probe, err := downloader.prepare(ctx, source, request, partPath)
	if err != nil {
		return "", err
	}

	progress := request.Progress
	progress.SetFileName(state.FileName)
	progress.SetTotal(state.TotalBytes)
	progress.SetDone(state.DoneBytes())
	remainingBytes := max(state.TotalBytes-state.DoneBytes(), 0)
	if err := diskspace.EnsureFree(request.Directory, remainingBytes, request.ReservedBytes); err != nil {
		closeBody(probe)
		return "", err
	}

	file, err := openPartFile(partPath, probe != nil)
	if err != nil {
		closeBody(probe)
		return "", err
	}
	runErr := newSession(downloader, source, file, request, state).run(ctx, probe)
	closeErr := file.Close()
	if runErr != nil {
		return "", runErr
	}
	if closeErr != nil {
		return "", closeErr
	}

	finalPath := files.UniquePath(filepath.Join(request.Directory, state.FileName))
	if err := os.Rename(partPath, finalPath); err != nil {
		return "", err
	}
	return finalPath, nil
}

func (downloader *Downloader) prepare(ctx context.Context, source *remoteFile, request Request, partPath string) (State, *http.Response, error) {
	if request.State != nil {
		if request.State.valid() && fileExists(partPath) {
			state := *request.State
			source.validator = state.Validator
			source.totalBytes = state.TotalBytes
			request.logger().Info("resuming from saved progress",
				"bytes_done", state.DoneBytes(),
				"bytes_total", state.TotalBytes,
				"parts", len(state.Segments),
			)
			return state, nil, nil
		}
		request.logger().Warn("saved progress is unusable, starting over", "part_file_exists", fileExists(partPath))
	}

	probe, err := downloader.open(ctx, source, 0, 0)
	if err != nil {
		return State{}, nil, err
	}
	source.describe(probe)

	connections := 1
	if probe.Header.Get("Accept-Ranges") == "bytes" {
		connections = min(max(request.Connections, 1), MaxConnections)
	}
	state := State{
		FileName:   chooseFileName(request.FileName, probe),
		Validator:  source.validator,
		TotalBytes: source.totalBytes,
		Segments:   planSegments(source.totalBytes, connections),
	}
	request.logger().Info("connected to server",
		"file", state.FileName,
		"bytes_total", state.TotalBytes,
		"supports_ranges", probe.Header.Get("Accept-Ranges") == "bytes",
		"has_validator", state.Validator != "",
		"parts", len(state.Segments),
	)
	return state, probe, nil
}

func chooseFileName(requested string, response *http.Response) string {
	if name := SanitizeFileName(requested); name != "" {
		return name
	}
	return FileNameFromResponse(response)
}

func openPartFile(partPath string, fresh bool) (*os.File, error) {
	flags := os.O_WRONLY | os.O_CREATE
	if fresh {
		flags |= os.O_TRUNC
	}
	return os.OpenFile(partPath, flags, 0o644)
}

func fileExists(path string) bool {
	info, err := os.Stat(path)
	return err == nil && info.Mode().IsRegular()
}

func closeBody(response *http.Response) {
	if response != nil {
		response.Body.Close()
	}
}