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 } } }