88 "io"
99 "net/http"
1010 "os"
11+ "strconv"
12+ "strings"
13+ "sync"
1114 "time"
1215
1316 "github.com/projecteru2/core/log"
@@ -29,29 +32,53 @@ const (
2932 progressInterval = 1 << 20
3033)
3134
32- // progressWriter wraps an io.Writer and periodically emits download progress events.
33- type progressWriter struct {
34- w io.Writer
35+ // progressCounter emits PhaseDownload events every ~1 MiB; mutex-guarded so it
36+ // serves both the serial writer and parallel range workers.
37+ type progressCounter struct {
38+ mu sync.Mutex
3539 written int64
40+ lastReport int64
3641 total int64
3742 tracker progress.Tracker
38- lastReport int64
3943}
4044
41- func (pw * progressWriter ) Write (p []byte ) (int , error ) {
42- n , err := pw .w .Write (p )
43- pw .written += int64 (n )
44- if pw .written - pw .lastReport >= progressInterval {
45- pw .lastReport = pw .written
46- pw .tracker .OnEvent (cloudimgProgress.Event {
45+ func (pc * progressCounter ) add (n int64 ) {
46+ pc .mu .Lock ()
47+ pc .written += n
48+ report := pc .written - pc .lastReport >= progressInterval
49+ if report {
50+ pc .lastReport = pc .written
51+ }
52+ done := pc .written
53+ pc .mu .Unlock ()
54+ // Emit outside the lock: the tracker callback is user-supplied and must not
55+ // serialize the range workers.
56+ if report {
57+ pc .tracker .OnEvent (cloudimgProgress.Event {
4758 Phase : cloudimgProgress .PhaseDownload ,
48- BytesTotal : pw .total ,
49- BytesDone : pw . written ,
59+ BytesTotal : pc .total ,
60+ BytesDone : done ,
5061 })
5162 }
63+ }
64+
65+ // countingWriter forwards writes to w while reporting byte counts to a shared progressCounter.
66+ type countingWriter struct {
67+ w io.Writer
68+ pc * progressCounter
69+ }
70+
71+ func (cw countingWriter ) Write (p []byte ) (int , error ) {
72+ n , err := cw .w .Write (p )
73+ cw .pc .add (int64 (n ))
5274 return n , err
5375}
5476
77+ // byteRange is an inclusive [start, end] byte span for an HTTP Range request.
78+ type byteRange struct {
79+ start , end int64
80+ }
81+
5582// pull commits url as a blob; the URL→blob mapping is idempotent (no-op when the blob already exists).
5683func pull (ctx context.Context , conf * Config , store storage.Store [imageIndex ], url string , force bool , tracker progress.Tracker ) error {
5784 logger := log .WithFunc ("cloudimg.pull" )
@@ -104,20 +131,46 @@ func withDownload(
104131 defer os .Remove (tmpPath ) //nolint:errcheck,gosec
105132 defer tmpFile .Close () //nolint:errcheck,gosec
106133
107- digestHex , err := downloadToFile (ctx , url , tmpFile , tracker )
134+ digestHex , err := downloadToFile (ctx , url , tmpFile , tracker , conf . PullConns )
108135 if err != nil {
109136 return err
110137 }
111138 return fn (tmpFile , tmpPath , digestHex )
112139}
113140
114- func downloadToFile (ctx context.Context , url string , dst * os.File , tracker progress.Tracker ) (string , error ) {
141+ // downloadToFile probes whether url supports HTTP Range requests and downloads it into dst
142+ // accordingly: pullConns concurrent range connections when supported, a single stream otherwise.
143+ func downloadToFile (ctx context.Context , url string , dst * os.File , tracker progress.Tracker , pullConns int ) (string , error ) {
144+ logger := log .WithFunc ("cloudimg.downloadToFile" )
145+ client := & http.Client {Timeout : urlDownloadTimeout }
146+
147+ rangeURL , size , ok , err := probeRangeSupport (ctx , client , url )
148+ if err != nil {
149+ return "" , err
150+ }
151+ if ! ok {
152+ logger .Debugf (ctx , "range requests not supported for %s, falling back to serial download" , url )
153+ return downloadSerial (ctx , client , url , dst , tracker )
154+ }
155+ if size > maxDownloadBytes {
156+ return "" , fmt .Errorf ("download %s: exceeded max size (%d bytes)" , url , maxDownloadBytes )
157+ }
158+
159+ logger .Debugf (ctx , "downloading %s in %d parallel range(s)" , url , pullConns )
160+ if err := downloadRangesParallel (ctx , client , rangeURL , dst , size , pullConns , tracker ); err != nil {
161+ return "" , err
162+ }
163+ return hashDigest (dst )
164+ }
165+
166+ // downloadSerial streams url into dst over a single connection; used when the server doesn't
167+ // support Range requests or doesn't report a usable size.
168+ func downloadSerial (ctx context.Context , client * http.Client , url string , dst * os.File , tracker progress.Tracker ) (string , error ) {
115169 req , err := http .NewRequestWithContext (ctx , http .MethodGet , url , nil )
116170 if err != nil {
117171 return "" , fmt .Errorf ("create http request: %w" , err )
118172 }
119173
120- client := & http.Client {Timeout : urlDownloadTimeout }
121174 resp , err := client .Do (req )
122175 if err != nil {
123176 return "" , fmt .Errorf ("http get %s: %w" , url , err )
@@ -138,7 +191,7 @@ func downloadToFile(ctx context.Context, url string, dst *os.File, tracker progr
138191 limitedBody := io .LimitReader (resp .Body , maxDownloadBytes + 1 )
139192 reader := io .TeeReader (limitedBody , h )
140193
141- pw := & progressWriter {w : dst , total : contentLength , tracker : tracker }
194+ pw := countingWriter {w : dst , pc : & progressCounter { total : contentLength , tracker : tracker } }
142195 written , err := io .Copy (pw , reader )
143196 if err != nil {
144197 return "" , fmt .Errorf ("download %s: %w" , url , err )
@@ -153,3 +206,134 @@ func downloadToFile(ctx context.Context, url string, dst *os.File, tracker progr
153206
154207 return hex .EncodeToString (h .Sum (nil )), nil
155208}
209+
210+ // probeRangeSupport issues a GET with Range: bytes=0-0 and reports whether the server serves partial
211+ // content with a known total size. It returns resp.Request.URL (post-redirect) so follow-up range
212+ // requests hit the resolved location directly, since Go's client can drop headers across a
213+ // cross-host redirect on a freshly built request.
214+ func probeRangeSupport (ctx context.Context , client * http.Client , url string ) (finalURL string , size int64 , ok bool , err error ) {
215+ req , err := http .NewRequestWithContext (ctx , http .MethodGet , url , nil )
216+ if err != nil {
217+ return "" , 0 , false , fmt .Errorf ("create range probe request: %w" , err )
218+ }
219+ req .Header .Set ("Range" , "bytes=0-0" )
220+
221+ resp , err := client .Do (req )
222+ if err != nil {
223+ return "" , 0 , false , fmt .Errorf ("probe range support for %s: %w" , url , err )
224+ }
225+ defer resp .Body .Close () //nolint:errcheck
226+
227+ if resp .StatusCode != http .StatusPartialContent {
228+ return "" , 0 , false , nil
229+ }
230+ size , ok = parseContentRangeSize (resp .Header .Get ("Content-Range" ))
231+ if ! ok {
232+ return "" , 0 , false , nil
233+ }
234+ return resp .Request .URL .String (), size , true , nil
235+ }
236+
237+ // downloadRangesParallel splits [0,size) into pullConns contiguous ranges and downloads them
238+ // concurrently into disjoint offsets of dst.
239+ func downloadRangesParallel (ctx context.Context , client * http.Client , url string , dst * os.File , size int64 , pullConns int , tracker progress.Tracker ) error {
240+ if err := dst .Truncate (size ); err != nil {
241+ return fmt .Errorf ("truncate temp file: %w" , err )
242+ }
243+
244+ tracker .OnEvent (cloudimgProgress.Event {Phase : cloudimgProgress .PhaseDownload , BytesTotal : size })
245+ pc := & progressCounter {total : size , tracker : tracker }
246+
247+ ranges := splitRanges (size , pullConns )
248+ if _ , err := utils .Map (ctx , ranges , func (ctx context.Context , _ int , r byteRange ) (struct {}, error ) {
249+ return struct {}{}, downloadRange (ctx , client , url , dst , r , pc )
250+ }, pullConns ); err != nil {
251+ return fmt .Errorf ("download %s: %w" , url , err )
252+ }
253+
254+ if err := dst .Sync (); err != nil {
255+ return fmt .Errorf ("sync temp file: %w" , err )
256+ }
257+
258+ tracker .OnEvent (cloudimgProgress.Event {Phase : cloudimgProgress .PhaseDownload , BytesTotal : size , BytesDone : size })
259+ return nil
260+ }
261+
262+ func downloadRange (ctx context.Context , client * http.Client , url string , dst * os.File , r byteRange , pc * progressCounter ) error {
263+ req , err := http .NewRequestWithContext (ctx , http .MethodGet , url , nil )
264+ if err != nil {
265+ return fmt .Errorf ("create range request: %w" , err )
266+ }
267+ req .Header .Set ("Range" , fmt .Sprintf ("bytes=%d-%d" , r .start , r .end ))
268+
269+ resp , err := client .Do (req )
270+ if err != nil {
271+ return fmt .Errorf ("http get range %d-%d: %w" , r .start , r .end , err )
272+ }
273+ defer resp .Body .Close () //nolint:errcheck
274+
275+ if resp .StatusCode != http .StatusPartialContent {
276+ return fmt .Errorf ("http get range %d-%d: status %s" , r .start , r .end , resp .Status )
277+ }
278+ // A mismatched span would land bytes at the wrong offsets; a short body would
279+ // leave a zero hole that hashDigest then blesses with a "valid" digest.
280+ if cr := resp .Header .Get ("Content-Range" ); ! strings .HasPrefix (cr , fmt .Sprintf ("bytes %d-%d/" , r .start , r .end )) {
281+ return fmt .Errorf ("http get range %d-%d: mismatched content-range %q" , r .start , r .end , cr )
282+ }
283+
284+ want := r .end - r .start + 1
285+ w := countingWriter {w : io .NewOffsetWriter (dst , r .start ), pc : pc }
286+ n , err := io .Copy (w , io .LimitReader (resp .Body , want ))
287+ if err != nil {
288+ return fmt .Errorf ("write range %d-%d: %w" , r .start , r .end , err )
289+ }
290+ if n != want {
291+ return fmt .Errorf ("http get range %d-%d: short body %d of %d bytes" , r .start , r .end , n , want )
292+ }
293+ return nil
294+ }
295+
296+ // hashDigest re-reads dst from the start and returns its sha256 hex digest. The parallel path can't
297+ // feed a streaming TeeReader since writes land at scattered offsets, so this also verifies what
298+ // actually landed on disk.
299+ func hashDigest (dst * os.File ) (string , error ) {
300+ if _ , err := dst .Seek (0 , io .SeekStart ); err != nil {
301+ return "" , fmt .Errorf ("seek temp file: %w" , err )
302+ }
303+ h := sha256 .New ()
304+ if _ , err := io .Copy (h , dst ); err != nil {
305+ return "" , fmt .Errorf ("hash temp file: %w" , err )
306+ }
307+ return hex .EncodeToString (h .Sum (nil )), nil
308+ }
309+
310+ // splitRanges divides [0,size) into up to n contiguous, inclusive-ended byte ranges.
311+ func splitRanges (size int64 , n int ) []byteRange {
312+ chunk := (size + int64 (n ) - 1 ) / int64 (n )
313+ ranges := make ([]byteRange , 0 , n )
314+ for start := int64 (0 ); start < size ; start += chunk {
315+ end := start + chunk - 1
316+ if end >= size {
317+ end = size - 1
318+ }
319+ ranges = append (ranges , byteRange {start : start , end : end })
320+ }
321+ return ranges
322+ }
323+
324+ // parseContentRangeSize extracts the total size from a "Content-Range: bytes 0-0/12345" header value.
325+ func parseContentRangeSize (v string ) (int64 , bool ) {
326+ i := strings .LastIndexByte (v , '/' )
327+ if i < 0 || i + 1 >= len (v ) {
328+ return 0 , false
329+ }
330+ total := v [i + 1 :]
331+ if total == "*" {
332+ return 0 , false
333+ }
334+ size , err := strconv .ParseInt (total , 10 , 64 )
335+ if err != nil || size <= 0 {
336+ return 0 , false
337+ }
338+ return size , true
339+ }
0 commit comments