Make sure we use the correct sized buffer

The buffer may be smaller, lean on content length to ensure we're sizing
the buffer correctly in the deliver and resume calls.

pkg/download/buffer.go

@@ -85,7 +85,7 @@ func (m *BufferMode) Fetch(ctx context.Context, url string) (io.Reader, int64, e
 			logger.Warn().
 				Int("connection_interrupted_at_byte", n).
 				Msg("Resuming Chunk Download")
-			_, err = resumeDownload(firstChunkResp.Request, buf[n:], m.Client, int64(n))
+			n, err = resumeDownload(firstChunkResp.Request, buf[n:contentLength], m.Client, int64(n))
 		}
 		firstChunk.Deliver(buf[0:n], err)
 	})
@@ -154,7 +154,7 @@ func (m *BufferMode) Fetch(ctx context.Context, url string) (io.Reader, int64, e
 					logger.Warn().
 						Int("connection_interrupted_at_byte", n).
 						Msg("Resuming Chunk Download")
-					_, err = resumeDownload(resp.Request, buf[n:], m.Client, int64(n))
+					n, err = resumeDownload(resp.Request, buf[n:contentLength], m.Client, int64(n))
 				}
 				chunk.Deliver(buf[0:n], err)
 			})

pkg/download/common.go

@@ -25,28 +25,30 @@ var (
 	errInvalidContentRange  = errors.New("invalid content range")
 )
 
-func resumeDownload(req *http.Request, buffer []byte, client client.HTTPClient, bytesReceived int64) (*http.Response, error) {
+func resumeDownload(req *http.Request, buffer []byte, client client.HTTPClient, bytesReceived int64) (int, error) {
 	var startByte int
 	logger := logging.GetLogger()
 
 	var resumeCount = 1
 	var initialBytesReceived = bytesReceived
+	var totalBytesReceived = bytesReceived
 
 	for {
 		var n int
 		if err := updateRangeRequestHeader(req, bytesReceived); err != nil {
-			return nil, err
+			return int(totalBytesReceived), err
 		}
 
 		resp, err := client.Do(req)
 		if err != nil {
-			return nil, err
+			return int(totalBytesReceived), err
 		}
 		defer resp.Body.Close()
 		if resp.StatusCode != http.StatusPartialContent {
-			return nil, fmt.Errorf("expected status code %d, got %d", http.StatusPartialContent, resp.StatusCode)
+			return int(totalBytesReceived), fmt.Errorf("expected status code %d, got %d", http.StatusPartialContent, resp.StatusCode)
 		}
 		n, err = io.ReadFull(resp.Body, buffer[startByte:])
+		totalBytesReceived += int64(n)
 		if err == io.ErrUnexpectedEOF {
 			bytesReceived = int64(n)
 			startByte += n
@@ -58,7 +60,7 @@ func resumeDownload(req *http.Request, buffer []byte, client client.HTTPClient,
 				Msg("Resuming Chunk Download")
 			continue
 		}
-		return nil, err
+		return int(totalBytesReceived), err
 
 	}
 }

pkg/download/common_test.go

@@ -112,12 +112,13 @@ func TestResumeDownload(t *testing.T) {
 				},
 			}
 
-			_, err = resumeDownload(req, buffer[tt.bytesReceived:], mockClient, tt.bytesReceived)
+			totalBytesReceived, err := resumeDownload(req, buffer[tt.bytesReceived:], mockClient, tt.bytesReceived)
 			if tt.expectedError != nil {
 				assert.Error(t, err)
 				assert.Equal(t, tt.expectedError.Error(), err.Error())
 			} else {
 				assert.NoError(t, err)
+				assert.Equal(t, len(tt.expectedOutput), totalBytesReceived)
 				assert.Equal(t, tt.expectedOutput, buffer[:len(tt.expectedOutput)])
 			}
 			assert.Equal(t, tt.expectedCalls, mockClient.callCount.Load(), "Unexpected number of HTTP client calls")

pkg/download/consistent_hashing.go

@@ -120,7 +120,7 @@ func (m *ConsistentHashingMode) Fetch(ctx context.Context, urlString string) (io
 			logger.Warn().
 				Int("connection_interrupted_at_byte", n).
 				Msg("Resuming Chunk Download")
-			_, err = resumeDownload(firstChunkResp.Request, buf[n:], m.Client, int64(n))
+			_, err = resumeDownload(firstChunkResp.Request, buf[n:contentLength], m.Client, int64(n))
 		}
 		firstChunk.Deliver(buf[0:n], err)
 	})
@@ -231,7 +231,7 @@ func (m *ConsistentHashingMode) downloadRemainingChunks(ctx context.Context, url
 					logger.Warn().
 						Int("connection_interrupted_at_byte", n).
 						Msg("Resuming Chunk Download")
-					_, err = resumeDownload(resp.Request, buf[n:], m.Client, int64(n))
+					_, err = resumeDownload(resp.Request, buf[n:contentLength], m.Client, int64(n))
 				}
 				chunk.Deliver(buf[0:n], err)
 			})