Add SHA256 checksum verification and robust URL validation for image

cmd/image/qcow2ova/get-image.go

@@ -15,6 +15,8 @@
 package qcow2ova
 
 import (
+	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
 	"io"
 	"net/http"
@@ -30,8 +32,34 @@ const (
 	DefaultGetTimeout = 30 * time.Minute
 )
 
+// verifyCheckSum validates SHA256 of a downloaded file
+func verifyCheckSum(filePath, expected string) error {
+	if expected == "" {
+		klog.V(1).Infof("No checksum provided for %s, skipping verification", filePath)
+		return nil
+	}
+	f, err := os.Open(filePath)
+	if err != nil {
+		return fmt.Errorf("failed to open file for checksum: %v", err)
+	}
+	defer f.Close()
+
+	h := sha256.New()
+	if _, err := io.Copy(h, f); err != nil {
+		return fmt.Errorf("failed to calculate checksum: %v", err)
+	}
+
+	actual := hex.EncodeToString(h.Sum(nil))
+	if actual != expected {
+		return fmt.Errorf("checksum mismatch for %s:\n expected: %s\n actual:	%s", filePath, expected, actual)
+	}
+	klog.V(1).Infof("Checksum verification PASSED FOR %s", filePath)
+	return nil
+}
+
 // Downloads or copy the image into the target dir mentioned
-func getImage(downloadDir string, srcUrl string, timeout time.Duration) (string, error) {
+// Added checksum verification (optional)
+func getImage(downloadDir string, srcUrl string, timeout time.Duration, expectedSha string) (string, error) {
 	if timeout == 0 {
 		timeout = DefaultGetTimeout
 	}
@@ -71,6 +99,11 @@ func getImage(downloadDir string, srcUrl string, timeout time.Duration) (string,
 		}
 		klog.V(1).Info("Download Completed!")
 	}
+	// Verify checksum if provided
+	if err := verifyCheckSum(dest, expectedSha); err != nil {
+		return "", err
+	}
+
 	return dest, nil
 }
 

cmd/image/qcow2ova/get-image_test.go

@@ -15,6 +15,7 @@
 package qcow2ova
 
 import (
+	"crypto/sha256"
 	"fmt"
 	"log"
 	"net/http"
@@ -66,10 +67,17 @@ func Test_getImage(t *testing.T) {
 		log.Fatal(err)
 	}
 
+	// Generate SHA256 for tmpfile (for checksum test)
+	h := sha256.New()
+	h.Write(content)
+	validSHA := fmt.Sprintf("%x", h.Sum(nil))
+	invalidSHA := "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
+
 	type args struct {
 		dir     string
 		src     string
 		timeout time.Duration
+		sha     string
 	}
 	tests := []struct {
 		name    string
@@ -78,45 +86,51 @@ func Test_getImage(t *testing.T) {
 		wantErr bool
 	}{
 		{
-			name:    "getImage of type file",
-			args:    args{destDir, tmpfn, 0},
+			name:    "getImage of type file with valid checksum",
+			args:    args{destDir, tmpfn, 0, validSHA},
 			want:    filepath.Join(destDir, "tmpfile"),
 			wantErr: false,
 		},
+		{
+			name:    "getImage of type file with invalid checksum",
+			args:    args{destDir, tmpfn, 0, invalidSHA},
+			want:    "",
+			wantErr: true,
+		},
 		{
 			name:    "getImage does not exist",
-			args:    args{destDir, "/file/doesnot/exist", 0},
+			args:    args{destDir, "/file/doesnot/exist", 0, ""},
 			want:    "",
 			wantErr: true,
 		},
 		{
 			name:    "getImage of type URL",
-			args:    args{destDir, ts.URL + "/file1", httpProcessingTime * 2},
+			args:    args{destDir, ts.URL + "/file1", httpProcessingTime * 2, ""},
 			want:    filepath.Join(destDir, "file1"),
 			wantErr: false,
 		},
 		{
 			name:    "getImage of type URL with default timeout",
-			args:    args{destDir, ts.URL + "/file2", 0},
+			args:    args{destDir, ts.URL + "/file2", 0, ""},
 			want:    filepath.Join(destDir, "file2"),
 			wantErr: false,
 		},
 		{
 			name:    "getImage of type URL - timeout failure",
-			args:    args{destDir, ts.URL + "/file", httpProcessingTime / 2},
+			args:    args{destDir, ts.URL + "/file", httpProcessingTime / 2, ""},
 			want:    "",
 			wantErr: true,
 		},
 		{
 			name:    "getImage of type URL - server side error",
-			args:    args{destDir, ts.URL + "/fail", 0},
+			args:    args{destDir, ts.URL + "/fail", 0, ""},
 			want:    "",
 			wantErr: true,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := getImage(tt.args.dir, tt.args.src, tt.args.timeout)
+			got, err := getImage(tt.args.dir, tt.args.src, tt.args.timeout, tt.args.sha)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("getImage() error = %v, wantErr %v", err, tt.wantErr)
 				return

cmd/image/qcow2ova/qcow2ova.go

@@ -215,7 +215,7 @@ Qcow2 images location:
 			os.Exit(1)
 		}()
 
-		image, err := getImage(tmpDir, opt.ImageURL, 0)
+		image, err := getImage(tmpDir, opt.ImageURL, 0, "")
 		if err != nil {
 			return fmt.Errorf("failed to download the %s into %s, error: %v", opt.ImageURL, tmpDir, err)
 		}