urldownload GetVolumeSize check for content length

pkg/storage/url/urldownload/storage.go

@@ -2,6 +2,7 @@ package urldownload
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"mime"
@@ -26,6 +27,10 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )
 
+var (
+	ErrNoContentLengthFound = errors.New("content-length not provided by the server")
+)
+
 // StorageProvider downloads data on request from a URL to a local
 // directory.
 
@@ -75,9 +80,39 @@ func (sp *StorageProvider) HasStorageLocally(context.Context, models.InputSource
 	return false, nil
 }
 
-func (sp *StorageProvider) GetVolumeSize(context.Context, models.InputSource) (uint64, error) {
-	// Could do a HEAD request and check Content-Length, but in some cases that's not guaranteed to be the real end file size
-	return 0, nil
+func (sp *StorageProvider) GetVolumeSize(ctx context.Context, storageSpec models.InputSource) (uint64, error) {
+	source, err := DecodeSpec(storageSpec.Source)
+	if err != nil {
+		return 0, err
+	}
+
+	u, err := IsURLSupported(source.URL)
+	if err != nil {
+		return 0, err
+	}
+
+	req, err := retryablehttp.NewRequestWithContext(ctx, http.MethodHead, u.String(), nil)
+	if err != nil {
+		return 0, err
+	}
+
+	res, err := sp.client.Do(req) //nolint:bodyclose // this is being closed - golangci-lint is wrong again
+	if err != nil {
+		return 0, err
+	}
+	defer closer.DrainAndCloseWithLogOnError(ctx, "response", res.Body)
+
+	if res.StatusCode != http.StatusOK {
+		return 0, fmt.Errorf("received non-OK response code %d while fetching size of file download", res.StatusCode)
+	}
+
+	// Ideally if the content size is not provided by server we should try and fetch the file with max size
+	// as the one provided in the storageSpec
+	if res.ContentLength < 0 {
+		return 0, ErrNoContentLengthFound
+	}
+
+	return uint64(res.ContentLength), nil
 }
 
 // PrepareStorage will download the file from the URL

pkg/storage/url/urldownload/storage_test.go

@@ -339,3 +339,86 @@ func (s *StorageSuite) TestPrepareStorageURL() {
 		})
 	}
 }
+func (s *StorageSuite) TestGetVolumeSize_WithServerReturningValidSize() {
+	path := "/initial"
+	headers := &map[string]string{
+		"Content-Length": "500",
+	}
+	code := 200
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
+		if r.URL.Path != path {
+			http.Error(w, fmt.Sprintf("invalid path: %s should be %s", r.URL.Path, path), 999)
+			return
+		}
+
+		if headers != nil {
+			// Set the headers, if any, before WriteHeader is called
+			for k, v := range *headers {
+				w.Header().Add(k, v)
+			}
+		}
+
+		w.WriteHeader(code)
+
+		_, err := w.Write([]byte(""))
+		s.NoError(err)
+	}))
+	s.T().Cleanup(ts.Close)
+
+	subject := NewStorage()
+
+	url := fmt.Sprintf("%s%s", ts.URL, path)
+	spec := models.InputSource{
+		Source: &models.SpecConfig{
+			Type: models.StorageSourceURL,
+			Params: Source{
+				URL: url,
+			}.ToMap(),
+		},
+		Target: "/inputs",
+	}
+
+	vs, err := subject.GetVolumeSize(context.Background(), spec)
+	s.Require().NoError(err)
+
+	s.Equal(uint64(500), vs, "content-length does not match")
+
+}
+
+func (s *StorageSuite) TestGetVolumeSize_WithServerReturningInvalidSize() {
+	path := "/initial"
+	code := 200
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
+		if r.URL.Path != path {
+			http.Error(w, fmt.Sprintf("invalid path: %s should be %s", r.URL.Path, path), 999)
+			return
+		}
+
+		w.WriteHeader(code)
+
+		_, err := w.Write([]byte(""))
+		s.NoError(err)
+	}))
+	s.T().Cleanup(ts.Close)
+
+	subject := NewStorage()
+
+	url := fmt.Sprintf("%s%s", ts.URL, path)
+	spec := models.InputSource{
+		Source: &models.SpecConfig{
+			Type: models.StorageSourceURL,
+			Params: Source{
+				URL: url,
+			}.ToMap(),
+		},
+		Target: "/inputs",
+	}
+
+	_, err := subject.GetVolumeSize(context.Background(), spec)
+	s.Require().ErrorIs(err, ErrNoContentLengthFound)
+
+}