Optimization

cmd/dowload_test.go → cmd/download_test.go

@@ -1,24 +1,16 @@
 package cmd
 
 import (
-	"crypto/sha256"
-	"encoding/base64"
 	"fmt"
 	"testing"
 
 	"github.com/cisco-open/grabit/test"
 	"github.com/stretchr/testify/assert"
 )
 
-func getSha256Integrity(content string) string {
-	hasher := sha256.New()
-	hasher.Write([]byte(content))
-	return fmt.Sprintf("sha256-%s", base64.StdEncoding.EncodeToString(hasher.Sum(nil)))
-}
-
 func TestRunDownload(t *testing.T) {
 	content := `abcdef`
-	contentIntegrity := getSha256Integrity(content)
+	contentIntegrity := test.GetSha256Integrity(content)
 	port := test.TestHttpHandler(content, t)
 	testfilepath := test.TmpFile(t, fmt.Sprintf(`
 	[[Resource]]
@@ -41,7 +33,7 @@ func TestRunDownload(t *testing.T) {
 
 func TestRunDownloadWithTags(t *testing.T) {
 	content := `abcdef`
-	contentIntegrity := getSha256Integrity(content)
+	contentIntegrity := test.GetSha256Integrity(content)
 	port := test.TestHttpHandler(content, t)
 	testfilepath := test.TmpFile(t, fmt.Sprintf(`
 	[[Resource]]
@@ -66,7 +58,7 @@ func TestRunDownloadWithTags(t *testing.T) {
 
 func TestRunDownloadWithoutTags(t *testing.T) {
 	content := `abcdef`
-	contentIntegrity := getSha256Integrity(content)
+	contentIntegrity := test.GetSha256Integrity(content)
 	port := test.TestHttpHandler(content, t)
 	testfilepath := test.TmpFile(t, fmt.Sprintf(`
 	[[Resource]]
@@ -126,7 +118,7 @@ func TestRunDownloadFailsIntegrityTest(t *testing.T) {
 
 func TestRunDownloadTriesAllUrls(t *testing.T) {
 	content := `abcdef`
-	contentIntegrity := getSha256Integrity(content)
+	contentIntegrity := test.GetSha256Integrity(content)
 	port := test.TestHttpHandler(content, t)
 	testfilepath := test.TmpFile(t, fmt.Sprintf(`
 	[[Resource]]

internal/resource.go

@@ -88,8 +88,37 @@ func (l *Resource) Download(dir string, mode os.FileMode, ctx context.Context) e
 	if err != nil {
 		return err
 	}
+
 	var downloadError error = nil
 	for _, u := range l.Urls {
+		localName := ""
+		if l.Filename != "" {
+			localName = l.Filename
+		} else {
+			localName = path.Base(u)
+		}
+		resPath := filepath.Join(dir, localName)
+
+		// Check if the destination file already exists and has the correct integrity.
+		_, err := os.Stat(resPath)
+		if err != nil {
+			if !os.IsNotExist(err) {
+				return fmt.Errorf("error checking destination file presence '%s': '%v'", resPath, err)
+			}
+		} else {
+			err = checkIntegrityFromFile(resPath, algo, l.Integrity, u)
+			if err != nil {
+				return fmt.Errorf("existing file at '%s' with incorrect integrity: '%v'", resPath, err)
+			}
+			if mode != NoFileMode {
+				err = os.Chmod(resPath, mode.Perm())
+				if err != nil {
+					return err
+				}
+			}
+			return nil
+		}
+
 		// Download file in the target directory so that the call to
 		// os.Rename is atomic.
 		lpath, err := GetUrlToDir(u, dir, ctx)
@@ -101,14 +130,6 @@ func (l *Resource) Download(dir string, mode os.FileMode, ctx context.Context) e
 		if err != nil {
 			return err
 		}
-
-		localName := ""
-		if l.Filename != "" {
-			localName = l.Filename
-		} else {
-			localName = path.Base(u)
-		}
-		resPath := filepath.Join(dir, localName)
 		err = os.Rename(lpath, resPath)
 		if err != nil {
 			return err
@@ -120,6 +141,7 @@ func (l *Resource) Download(dir string, mode os.FileMode, ctx context.Context) e
 			}
 		}
 		ok = true
+		break
 	}
 	if !ok {
 		if downloadError != nil {

internal/resource_test.go

@@ -4,8 +4,11 @@
 package internal
 
 import (
+	"context"
 	"fmt"
 	"net/http"
+	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/cisco-open/grabit/test"
@@ -50,3 +53,34 @@ func TestNewResourceFromUrl(t *testing.T) {
 		}
 	}
 }
+
+func TestResourceDownloadWithValidFileAlreadyPresent(t *testing.T) {
+	content := `abcdef`
+	contentIntegrity := test.GetSha256Integrity(content)
+	port := 33 // unused because the file is already present.
+	testFileName := "test.html"
+	resource := Resource{Urls: []string{fmt.Sprintf("http://localhost:%d/%s", port, testFileName)}, Integrity: contentIntegrity, Tags: []string{}, Filename: ""}
+	outputDir := test.TmpDir(t)
+	err := os.WriteFile(filepath.Join(outputDir, testFileName), []byte(content), 0644)
+	assert.Nil(t, err)
+	err = resource.Download(outputDir, 0644, context.Background())
+	assert.Nil(t, err)
+	for _, file := range []string{testFileName} {
+		test.AssertFileContains(t, fmt.Sprintf("%s/%s", outputDir, file), content)
+	}
+}
+
+func TestResourceDownloadWithInValidFileAlreadyPresent(t *testing.T) {
+	content := `abcdef`
+	contentIntegrity := test.GetSha256Integrity(content)
+	port := 33 // unused because the file, although invalid, is already present.
+	testFileName := "test.html"
+	resource := Resource{Urls: []string{fmt.Sprintf("http://localhost:%d/%s", port, testFileName)}, Integrity: contentIntegrity, Tags: []string{}, Filename: ""}
+	outputDir := test.TmpDir(t)
+	err := os.WriteFile(filepath.Join(outputDir, testFileName), []byte("invalid"), 0644)
+	assert.Nil(t, err)
+	err = resource.Download(outputDir, 0644, context.Background())
+	assert.NotNil(t, err)
+	assert.Contains(t, err.Error(), "integrity mismatch")
+	assert.Contains(t, err.Error(), "existing file")
+}

test/utils.go

@@ -4,6 +4,9 @@
 package test
 
 import (
+	"crypto/sha256"
+	"encoding/base64"
+	"fmt"
 	"log"
 	"net"
 	"net/http"
@@ -14,6 +17,12 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func GetSha256Integrity(content string) string {
+	hasher := sha256.New()
+	hasher.Write([]byte(content))
+	return fmt.Sprintf("sha256-%s", base64.StdEncoding.EncodeToString(hasher.Sum(nil)))
+}
+
 func TmpFile(t *testing.T, content string) string {
 	f, err := os.CreateTemp(t.TempDir(), "test")
 	if err != nil {