Crie um downloader retomável em Go com chunks concorrentes

(
  mkdir -- range-downloader &&
  cd -- range-downloader &&
  GOENV=off GOWORK=off GOTOOLCHAIN=local GOFLAGS= go mod init range-downloader
) && cd -- range-downloader

package main

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

const chunkSize int64 = 4 << 20

type identity struct {
	URL, ETag, SHA256 string
	Size, ChunkSize   int64
}

func probe(ctx context.Context, client *http.Client, url, digest string) (identity, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
	if err != nil {
		return identity{}, err
	}
	req.Header.Set("Accept-Encoding", "identity")
	resp, err := client.Do(req)
	if err != nil {
		return identity{}, err
	}
	defer resp.Body.Close()
	tag := resp.Header.Get("ETag")
	if resp.StatusCode != 200 || resp.ContentLength < 0 || resp.ContentLength > 1<<40 ||
		len(tag) < 2 || !strings.HasPrefix(tag, "\"") || !strings.HasSuffix(tag, "\"") ||
		resp.Header.Get("Content-Encoding") != "" {
		return identity{}, errors.New("need a known size (at most 1 TiB), strong ETag and unencoded HEAD 200")
	}
	return identity{url, tag, digest, resp.ContentLength, chunkSize}, nil
}

func fetchChunk(ctx context.Context, client *http.Client, id identity, dir string, start int64) error {
	end := min(start+id.ChunkSize, id.Size) - 1
	name := filepath.Join(dir, fmt.Sprintf("%d.part", start))
	if info, err := os.Lstat(name); err == nil {
		if info.Mode().IsRegular() && info.Size() == end-start+1 {
			return nil
		}
		return errors.New("invalid saved chunk; use a new output path")
	} else if !errors.Is(err, os.ErrNotExist) {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, id.URL, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Accept-Encoding", "identity")
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", start, end))
	req.Header.Set("If-Range", id.ETag)
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	expected := fmt.Sprintf("bytes %d-%d/%d", start, end, id.Size)
	if resp.StatusCode != http.StatusPartialContent || resp.Header.Get("Content-Range") != expected ||
		resp.Header.Get("ETag") != id.ETag || resp.Header.Get("Content-Encoding") != "" ||
		(resp.ContentLength != -1 && resp.ContentLength != end-start+1) {
		return errors.New("server rejected range or changed representation")
	}
	tmp, err := os.CreateTemp(dir, ".chunk-")
	if err != nil {
		return err
	}
	defer os.Remove(tmp.Name())
	defer tmp.Close()
	n, err := io.CopyBuffer(tmp, io.LimitReader(resp.Body, end-start+2), make([]byte, 32<<10))
	if err != nil {
		return err
	}
	if n != end-start+1 {
		return errors.New("incorrect chunk length")
	}
	if err := tmp.Sync(); err != nil {
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	return os.Rename(tmp.Name(), name)
}

func download(ctx context.Context, client *http.Client, url, output, digest string, workers int) error {
	sum, err := hex.DecodeString(digest)
	if err != nil || len(sum) != sha256.Size || workers < 1 || workers > 16 {
		return errors.New("provide a SHA-256 hex digest and 1–16 workers")
	}
	digest = strings.ToLower(digest)
	id, err := probe(ctx, client, url, digest)
	if err != nil {
		return err
	}
	dir := output + ".parts"
	fresh := false
	if err := os.Mkdir(dir, 0700); err == nil {
		fresh = true
	} else if !errors.Is(err, os.ErrExist) {
		return err
	}
	info, err := os.Lstat(dir)
	if err != nil {
		return err
	}
	if !info.IsDir() || info.Mode().Perm()&0077 != 0 {
		return errors.New("parts directory must be private")
	}
	lock := filepath.Join(dir, ".lock")
	if err := os.Mkdir(lock, 0700); err != nil {
		return errors.New("parts directory locked; another download may be active")
	}
	defer os.Remove(lock)
	manifest := filepath.Join(dir, "identity.json")
	if fresh {
		data, err := json.Marshal(id)
		if err != nil {
			return err
		}
		if err := os.WriteFile(manifest, data, 0600); err != nil {
			return err
		}
	} else {
		data, err := os.ReadFile(manifest)
		if err != nil {
			return err
		}
		var saved identity
		if err := json.Unmarshal(data, &saved); err != nil {
			return err
		}
		if saved != id {
			return errors.New("saved download identity changed; use a new output path")
		}
	}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	var wg sync.WaitGroup
	var once sync.Once
	var firstErr error
	jobs := make(chan int64)
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for start := range jobs {
				if ctx.Err() != nil {
					return
				}
				if err := fetchChunk(ctx, client, id, dir, start); err != nil {
					once.Do(func() { firstErr = err; cancel() })
					return
				}
			}
		}()
	}
send:
	for start := int64(0); start < id.Size; start += id.ChunkSize {
		select {
		case jobs <- start:
		case <-ctx.Done():
			break send
		}
	}
	close(jobs)
	wg.Wait()
	if firstErr != nil {
		return firstErr
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	tmp, err := os.CreateTemp(filepath.Dir(output), ".download-")
	if err != nil {
		return err
	}
	defer os.Remove(tmp.Name())
	defer tmp.Close()
	hash := sha256.New()
	for start := int64(0); start < id.Size; start += id.ChunkSize {
		if err := ctx.Err(); err != nil {
			return err
		}
		part, err := os.Open(filepath.Join(dir, fmt.Sprintf("%d.part", start)))
		if err != nil {
			return err
		}
		n, copyErr := io.Copy(io.MultiWriter(tmp, hash), io.LimitReader(part, id.ChunkSize+1))
		closeErr := part.Close()
		if copyErr != nil {
			return copyErr
		}
		if closeErr != nil {
			return closeErr
		}
		if n != min(id.ChunkSize, id.Size-start) {
			return errors.New("saved chunk length changed")
		}
		fmt.Fprintf(os.Stderr, "Assembled: %d/%d bytes\n", start+n, id.Size)
	}
	if hex.EncodeToString(hash.Sum(nil)) != digest {
		return errors.New("SHA-256 mismatch; use a new output path")
	}
	if err := tmp.Sync(); err != nil {
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	return os.Rename(tmp.Name(), output)
}

func main() {
	if len(os.Args) != 4 {
		fmt.Fprintln(os.Stderr, "usage: downloader URL OUTPUT SHA256")
		os.Exit(2)
	}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.MaxIdleConnsPerHost = 4
	defer transport.CloseIdleConnections()
	client := &http.Client{Transport: transport, Timeout: 2 * time.Minute}
	if err := download(ctx, client, os.Args[1], os.Args[2], os.Args[3], 4); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

(
  GOENV=off GOWORK=off GOTOOLCHAIN=local GOFLAGS= GOOS= GOARCH= go build -o downloader main.go &&
  ./downloader "${DOWNLOAD_URL:?Set DOWNLOAD_URL}" "${OUTPUT_PATH:?Set OUTPUT_PATH}" "${EXPECTED_SHA256:?Set EXPECTED_SHA256}"
)

package main

import (
	"bytes"
	"crypto/sha256"
	"fmt"
	"net"
	"net/http"
	"os"
	"time"
)

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintln(os.Stderr, "usage: local-origin 127.0.0.1:PORT")
		os.Exit(2)
	}
	data := append(bytes.Repeat([]byte{0x00, 0x80, 0xff, 0x0a}, 2<<20), []byte("tail\n")...)
	digest := fmt.Sprintf("%x", sha256.Sum256(data))
	listener, err := net.Listen("tcp", os.Args[1])
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	defer listener.Close()
	mux := http.NewServeMux()
	mux.HandleFunc("/file", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("ETag", "\""+digest+"\"")
		http.ServeContent(w, r, "fixture.bin", time.Time{}, bytes.NewReader(data))
	})
	fmt.Printf("DOWNLOAD_URL=http://%s/file\nEXPECTED_SHA256=%s\n", listener.Addr(), digest)
	server := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second}
	if err := server.Serve(listener); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

(
  GOENV=off GOWORK=off GOTOOLCHAIN=local GOFLAGS= GOOS= GOARCH= go build -o local-origin origin.go &&
  ./local-origin 127.0.0.1:0
)

export OUTPUT_PATH='example.bin'

(
  printf '%s  %s\n' "${EXPECTED_SHA256:?Set EXPECTED_SHA256}" "${OUTPUT_PATH:?Set OUTPUT_PATH}" |
    sha256sum --check --status
)