Merge pull request #19 from rabadin/read-files-in-chunks

Read files in chunks to avoid reading entire files in memory
cisco-open · Aug 28, 2024 · b627bf7 · b627bf7
2 parents b8b02c1 + edfd914
commit b627bf7
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 8 deletions.
diff --git a/internal/resource_test.go b/internal/resource_test.go
@@ -41,11 +41,11 @@ func TestNewResourceFromUrl(t *testing.T) {
 
 	for _, data := range tests {
 		resource, err := NewResourceFromUrl(data.urls, algo, []string{}, "")
-		assert.Equal(t, err == nil, data.valid)
+		assert.Equal(t, data.valid, err == nil)
 		if err != nil {
 			assert.Contains(t, err.Error(), data.errorContains)
 		} else {
-			assert.Equal(t, *resource, data.res)
+			assert.Equal(t, data.res, *resource)
 		}
 	}
 }
diff --git a/internal/sri.go b/internal/sri.go
@@ -4,8 +4,10 @@
 package internal
 
 import (
+	"bufio"
 	"encoding/base64"
 	"fmt"
+	"io"
 	"os"
 	"strings"
 )
@@ -15,12 +17,31 @@ func getIntegrityFromFile(path string, algo string) (string, error) {
 	if err != nil {
 		return "", err
 	}
-	f, err := os.ReadFile(path)
+	hasher := hash.hash()
+	f, err := os.Open(path)
+	defer f.Close()
 	if err != nil {
 		return "", fmt.Errorf("cannot open file '%s'", path)
 	}
-	hasher := hash.hash()
-	hasher.Write(f)
+	reader := bufio.NewReader(f)
+	// each block size is 10MB
+	const chunkSize = 10 * 1024 * 1024
+	buf := make([]byte, chunkSize)
+	for {
+		n, err := reader.Read(buf)
+		buf = buf[:n]
+
+		if n == 0 {
+			if err != nil {
+				if err == io.EOF {
+					break
+				}
+				return "", err
+			}
+			break
+		}
+		hasher.Write(buf)
+	}
 	h := hasher.Sum(nil)
 	return fmt.Sprintf("%s-%s", algo, base64.StdEncoding.EncodeToString(h)), nil
 }

diff --git a/internal/sri_test.go b/internal/sri_test.go
@@ -35,11 +35,11 @@ func TestGetAlgoFromIntegrity(t *testing.T) {
 
 	for _, data := range tests {
 		algo, err := getAlgoFromIntegrity(data.sriString)
-		assert.Equal(t, err == nil, data.valid)
+		assert.Equal(t, data.valid, err == nil)
 		if err != nil {
 			assert.Contains(t, err.Error(), data.errorContains)
 		} else {
-			assert.Equal(t, algo, data.resString)
+			assert.Equal(t, data.resString, algo)
 		}
 	}
 }
@@ -48,7 +48,7 @@ func TestGetIntegrityFromFile(t *testing.T) {
 	path := tmpFile(t, "abcdef")
 	sri, err := getIntegrityFromFile(path, "sha256")
 	assert.Nil(t, err)
-	assert.Equal(t, sri, "sha256-vvV+x/U6bUC+tkCngKY5yDvCmsipgW8fxsXG3Nk8RyE=")
+	assert.Equal(t, "sha256-vvV+x/U6bUC+tkCngKY5yDvCmsipgW8fxsXG3Nk8RyE=", sri)
 }
 
 func TestCheckIntegrityFromFile(t *testing.T) {