add file size sending and improve tests for sigv4

Signed-off-by: Ashley Davis <ashley.davis@cyberark.com>

Author: SgtCoDFish

internal/cyberark/dataupload/dataupload.go

@@ -124,7 +124,7 @@ func (c *CyberArkClient) PutSnapshot(ctx context.Context, snapshot Snapshot) err
 	checksumHex := hex.EncodeToString(checksum)
 	checksumBase64 := base64.StdEncoding.EncodeToString(checksum)
 
-	presignedUploadURL, username, err := c.retrievePresignedUploadURL(ctx, checksumHex, snapshot.ClusterID)
+	presignedUploadURL, username, err := c.retrievePresignedUploadURL(ctx, checksumHex, snapshot.ClusterID, int64(encodedBody.Len()))
 	if err != nil {
 		return fmt.Errorf("while retrieving snapshot upload URL: %s", err)
 	}
@@ -168,20 +168,30 @@ func (c *CyberArkClient) PutSnapshot(ctx context.Context, snapshot Snapshot) err
 	return nil
 }
 
-func (c *CyberArkClient) retrievePresignedUploadURL(ctx context.Context, checksum string, clusterID string) (string, string, error) {
+// RetrievePresignedUploadURLRequest is the JSON body sent to the inventory API to request a presigned upload URL.
+type RetrievePresignedUploadURLRequest struct {
+	ClusterID string `json:"cluster_id"`
+	Checksum  string `json:"checksum_sha256"`
+
+	// AgentVersion is the v-prefixed version of the agent uploading the snapshot.
+	// Note that the backend relies on this version being v-prefixed semver.
+	AgentVersion string `json:"agent_version"`
+
+	// FileSize is the size of the data we'll upload in bytes
+	FileSize int64 `json:"file_size"`
+}
+
+func (c *CyberArkClient) retrievePresignedUploadURL(ctx context.Context, checksum string, clusterID string, fileSize int64) (string, string, error) {
 	uploadURL, err := url.JoinPath(c.baseURL, apiPathSnapshotLinks)
 	if err != nil {
 		return "", "", err
 	}
 
-	request := struct {
-		ClusterID    string `json:"cluster_id"`
-		Checksum     string `json:"checksum_sha256"`
-		AgentVersion string `json:"agent_version"`
-	}{
+	request := RetrievePresignedUploadURLRequest{
 		ClusterID:    clusterID,
 		Checksum:     checksum,
 		AgentVersion: version.PreflightVersion,
+		FileSize:     fileSize,
 	}
 
 	encodedBody := &bytes.Buffer{}

internal/cyberark/dataupload/mock.go

@@ -2,13 +2,17 @@ package dataupload
 
 import (
 	"bytes"
+	"crypto/rand"
 	"crypto/sha256"
 	"encoding/base64"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -25,9 +29,19 @@ const (
 	successClusterID = "ffffffff-ffff-ffff-ffff-ffffffffffff"
 )
 
+type uploadValues struct {
+	ClusterID string
+	FileSize  int64
+}
+
 type mockDataUploadServer struct {
 	t         testing.TB
 	serverURL string
+
+	mux *http.ServeMux
+
+	expectedUploadValues      map[string]uploadValues
+	expectedUploadValuesMutex sync.Mutex
 }
 
 // MockDataUploadServer starts a server which mocks the CyberArk
@@ -45,39 +59,48 @@ type mockDataUploadServer struct {
 // responses.
 func MockDataUploadServer(t testing.TB) (string, *http.Client) {
 	mux := http.NewServeMux()
-	server := httptest.NewTLSServer(mux)
-	t.Cleanup(server.Close)
 	mds := &mockDataUploadServer{
-		t:         t,
-		serverURL: server.URL,
+		t: t,
+
+		expectedUploadValues: make(map[string]uploadValues),
 	}
-	mux.Handle("/", mds)
+
+	mux.HandleFunc("POST "+apiPathSnapshotLinks, mds.handleSnapshotLinks)
+
+	// The path includes random data to ensure that each request is treated separately by the mock server, allowing us to track data across calls.
+	// It also ensures that the client isn't using some pre-saved path and is actually using the presigned URL returned by the mock server in the previous step, which is important for test validity.
+	mux.HandleFunc("PUT /presigned-upload/{randData}", mds.handlePresignedUpload)
+
+	server := httptest.NewTLSServer(mds)
+	t.Cleanup(server.Close)
+
+	mds.mux = mux
+	mds.serverURL = server.URL
+
 	httpClient := server.Client()
 	httpClient.Transport = transport.NewDebuggingRoundTripper(httpClient.Transport, transport.DebugByContext)
 	return server.URL, httpClient
 }
 
 func (mds *mockDataUploadServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	mds.t.Log(r.Method, r.RequestURI)
-	switch r.URL.Path {
-	case apiPathSnapshotLinks:
-		mds.handleSnapshotLinks(w, r)
-		return
-	case "/presigned-upload":
-		mds.handlePresignedUpload(w, r)
-		return
-	default:
-		w.WriteHeader(http.StatusNotFound)
-	}
+
+	mds.mux.ServeHTTP(w, r)
 }
 
-func (mds *mockDataUploadServer) handleSnapshotLinks(w http.ResponseWriter, r *http.Request) {
-	if r.Method != http.MethodPost {
-		w.WriteHeader(http.StatusMethodNotAllowed)
-		_, _ = w.Write([]byte(`{"message":"method not allowed"}`))
-		return
+// randHex reads 8 random bytes and returns them as a hex string. It is used to generate
+// unique paths per-request to ensure that file size is tracked across calls.
+func randHex() string {
+	b := make([]byte, 8)
+	_, err := rand.Read(b)
+	if err != nil {
+		panic("failed to read random bytes: " + err.Error())
 	}
 
+	return hex.EncodeToString(b)
+}
+
+func (mds *mockDataUploadServer) handleSnapshotLinks(w http.ResponseWriter, r *http.Request) {
 	if r.Header.Get("User-Agent") != version.UserAgent() {
 		http.Error(w, "should set user agent on all requests", http.StatusInternalServerError)
 		return
@@ -99,13 +122,11 @@ func (mds *mockDataUploadServer) handleSnapshotLinks(w http.ResponseWriter, r *h
 		return
 	}
 
+	var req RetrievePresignedUploadURLRequest
+
 	decoder := json.NewDecoder(r.Body)
-	var req struct {
-		ClusterID    string `json:"cluster_id"`
-		Checksum     string `json:"checksum_sha256"`
-		AgentVersion string `json:"agent_version"`
-	}
 	decoder.DisallowUnknownFields()
+
 	if err := decoder.Decode(&req); err != nil {
 		http.Error(w, `{"error": "Invalid request format"}`, http.StatusBadRequest)
 		return
@@ -135,10 +156,33 @@ func (mds *mockDataUploadServer) handleSnapshotLinks(w http.ResponseWriter, r *h
 		return
 	}
 
+	if req.FileSize <= 0 {
+		http.Error(w, "file size must be greater than 0", http.StatusInternalServerError)
+		return
+	}
+
+	randomData := randHex()
+
+	mds.expectedUploadValuesMutex.Lock()
+	defer mds.expectedUploadValuesMutex.Unlock()
+
+	uploadValues := uploadValues{
+		ClusterID: req.ClusterID,
+		FileSize:  req.FileSize,
+	}
+
+	mds.expectedUploadValues[randomData] = uploadValues
+
+	presignedURL, err := url.JoinPath(mds.serverURL, "presigned-upload", randomData)
+	if err != nil {
+		http.Error(w, "failed to generate presigned URL", http.StatusInternalServerError)
+		mds.t.Logf("failed to generate presigned URL: %v", err)
+		return
+	}
+
 	// Write response body
 	w.WriteHeader(http.StatusOK)
 	w.Header().Set("Content-Type", "application/json")
-	presignedURL := mds.serverURL + "/presigned-upload"
 	_ = json.NewEncoder(w).Encode(struct {
 		URL string `json:"url"`
 	}{presignedURL})
@@ -155,9 +199,18 @@ const amzExampleChecksumError = `<?xml version="1.0" encoding="UTF-8"?>
 </Error>`
 
 func (mds *mockDataUploadServer) handlePresignedUpload(w http.ResponseWriter, r *http.Request) {
-	if r.Method != http.MethodPut {
-		w.WriteHeader(http.StatusMethodNotAllowed)
-		_, _ = w.Write([]byte(`{"message":"method not allowed"}`))
+	randData := r.PathValue("randData")
+	if randData == "" {
+		http.Error(w, "missing randData in path; should match that returned in presigned url", http.StatusInternalServerError)
+		return
+	}
+
+	mds.expectedUploadValuesMutex.Lock()
+	uploadValues, ok := mds.expectedUploadValues[randData]
+	mds.expectedUploadValuesMutex.Unlock()
+
+	if !ok {
+		http.Error(w, "didn't find a prior call to generate presigned URL", http.StatusInternalServerError)
 		return
 	}
 
@@ -178,9 +231,65 @@ func (mds *mockDataUploadServer) handlePresignedUpload(w http.ResponseWriter, r
 		return
 	}
 
+	sseHeader := r.Header.Get("X-Amz-Server-Side-Encryption")
+	if sseHeader != "AES256" {
+		http.Error(w, "should set x-amz-server-side-encryption header to AES256 on all requests", http.StatusInternalServerError)
+		return
+	}
+
+	taggingHeader := r.Header.Get("X-Amz-Tagging")
+	if taggingHeader == "" {
+		http.Error(w, "should set x-amz-tagging header on all requests", http.StatusInternalServerError)
+		return
+	}
+
+	tags, err := url.ParseQuery(taggingHeader)
+	if err != nil {
+		http.Error(w, "x-amz-tagging header should be encoded as a valid query string", http.StatusInternalServerError)
+		return
+	}
+
+	if tags.Get("agent_version") != version.PreflightVersion {
+		http.Error(w, fmt.Sprintf("x-amz-tagging should contain an agent_version tag with value %s", version.PreflightVersion), http.StatusInternalServerError)
+		return
+	}
+
+	if tags.Get("tenant_id") == "" {
+		// TODO: if we change setup a bit, we can check the tenant_id matches the expected tenant_id from the test config, but for now, just check it's set
+		http.Error(w, "x-amz-tagging should contain a tenant_id tag", http.StatusInternalServerError)
+		return
+	}
+
+	if tags.Get("upload_type") != "k8s_snapshot" {
+		http.Error(w, "x-amz-tagging should contain an upload_type tag with value k8s_snapshot", http.StatusInternalServerError)
+		return
+	}
+
+	if tags.Get("uploader_id") != uploadValues.ClusterID {
+		http.Error(w, "x-amz-tagging should contain an uploader_id tag which matches the cluster ID sent in the RetrievePresignedUploadURL request", http.StatusInternalServerError)
+		return
+	}
+
+	if tags.Get("username") == "" {
+		// TODO: if we change setup a bit, we can check the username matches the expected username from the test config
+		// but for now, just check it's set
+		http.Error(w, "x-amz-tagging should contain a username tag", http.StatusInternalServerError)
+		return
+	}
+
+	if tags.Get("vendor") != "k8s" {
+		http.Error(w, "x-amz-tagging should contain a vendor tag with value k8s", http.StatusInternalServerError)
+		return
+	}
+
 	body, err := io.ReadAll(r.Body)
 	require.NoError(mds.t, err)
 
+	if uploadValues.FileSize != int64(len(body)) {
+		http.Error(w, fmt.Sprintf("file size in request body should match that sent in RetrievePresignedUploadURL request; expected %d, got %d", uploadValues.FileSize, len(body)), http.StatusInternalServerError)
+		return
+	}
+
 	hash := sha256.New()
 	_, err = hash.Write(body)
 	require.NoError(mds.t, err)

internal/cyberark/servicediscovery/testdata/discovery_success.json.template

@@ -3,6 +3,7 @@
   "dr_region": "us-east-2",
   "subdomain": "venafi-test",
   "platform_id": "platform-123",
+  "tenant_id": "tenant-123",
   "identity_id": "identity-456",
   "default_url": "https://venafi-test.integration-cyberark.cloud",
   "tenant_flags": {