Initial DDC release

internal/fetch/downloader.go

@@ -0,0 +1,152 @@
+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
+		}
+	}
+}

internal/fetch/filename.go

@@ -0,0 +1,65 @@
+package fetch
+
+import (
+	"mime"
+	"net/http"
+	"path"
+	"path/filepath"
+	"strings"
+	"unicode/utf8"
+)
+
+const (
+	maxFileNameBytes  = 200
+	maxExtensionBytes = 16
+	fallbackFileName  = "download"
+)
+
+func FileNameFromResponse(response *http.Response) string {
+	if disposition := response.Header.Get("Content-Disposition"); disposition != "" {
+		if _, parameters, err := mime.ParseMediaType(disposition); err == nil {
+			if name := SanitizeFileName(parameters["filename"]); name != "" {
+				return name
+			}
+		}
+	}
+	if response.Request != nil && response.Request.URL != nil {
+		if name := SanitizeFileName(path.Base(response.Request.URL.Path)); name != "" {
+			return name
+		}
+	}
+	return fallbackFileName
+}
+
+func SanitizeFileName(raw string) string {
+	name := path.Base(strings.ReplaceAll(raw, "\\", "/"))
+	name = strings.Map(dropUnsafeRune, name)
+	name = strings.TrimSpace(strings.TrimLeft(name, ". "))
+	if name == "" {
+		return ""
+	}
+	return truncate(name)
+}
+
+func dropUnsafeRune(character rune) rune {
+	if character < 0x20 || character == 0x7f || character == '/' || character == utf8.RuneError {
+		return -1
+	}
+	return character
+}
+
+func truncate(name string) string {
+	if len(name) <= maxFileNameBytes {
+		return name
+	}
+	extension := filepath.Ext(name)
+	if len(extension) > maxExtensionBytes {
+		extension = ""
+	}
+	base := strings.TrimSuffix(name, extension)
+	for len(base)+len(extension) > maxFileNameBytes {
+		_, size := utf8.DecodeLastRuneInString(base)
+		base = base[:len(base)-size]
+	}
+	return base + extension
+}