test: add tests for gcs storage client and dns package

snowflakedb · Nov 8, 2023 · 89ceef3 · 89ceef3
1 parent f7b5956
commit 89ceef3
Show file tree

Hide file tree

Showing 2 changed files with 101 additions and 0 deletions.
diff --git a/dsn_test.go b/dsn_test.go
@@ -3,6 +3,8 @@
 package gosnowflake
 
 import (
+	"crypto/ecdsa"
+	"crypto/elliptic"
 	cr "crypto/rand"
 	"crypto/rsa"
 	"crypto/x509"
@@ -15,6 +17,8 @@ import (
 	"strings"
 	"testing"
 	"time"
+
+	"github.com/aws/smithy-go/rand"
 )
 
 type tcParseDSN struct {
@@ -1358,6 +1362,33 @@ func TestParsePrivateKeyFromFileIncorrectData(t *testing.T) {
 	}
 }
 
+func TestParsePrivateKeyFromFileNotRSAPrivateKey(t *testing.T) {
+	// Generate an ECDSA private key for testing
+	ecdsaPrivateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	if err != nil {
+		t.Fatalf("failed to generate ECDSA private key: %v", err)
+	}
+
+	ecdsaPrivateKeyBytes, err := x509.MarshalECPrivateKey(ecdsaPrivateKey)
+	if err != nil {
+		t.Fatalf("failed to marshal ECDSA private key: %v", err)
+	}
+	pemBlock := &pem.Block{
+		Type:  "EC PRIVATE KEY",
+		Bytes: ecdsaPrivateKeyBytes,
+	}
+	pemData := pem.EncodeToMemory(pemBlock)
+
+	// Write the PEM data to a temporary file
+	pemFile := createTmpFile("ecdsaKey.pem", pemData)
+
+	// Attempt to parse the private key
+	_, err = parsePrivateKeyFromFile(pemFile)
+	if err == nil {
+		t.Error("expected an error when trying to parse an ECDSA private key as RSA")
+	}
+}
+
 func TestParsePrivateKeyFromFile(t *testing.T) {
 	generatedKey, _ := rsa.GenerateKey(cr.Reader, 1024)
 	pemKey, _ := x509.MarshalPKCS8PrivateKey(generatedKey)

diff --git a/gcs_storage_client_test.go b/gcs_storage_client_test.go
@@ -810,6 +810,40 @@ func TestGetFileHeaderEncryptionData(t *testing.T) {
 	}
 }
 
+func TestGetFileHeaderEncryptionDataInterfaceConversionError(t *testing.T) {
+	mockEncDataResp := "{\"EncryptionMode\":\"FullBlob\",\"WrappedContentKey\": {\"KeyId\":\"symmKey1\",\"EncryptedKey\":\"testencryptedkey12345678910==\",\"Algorithm\":\"AES_CBC_256\"},\"EncryptionAgent\": {\"Protocol\":\"1.0\",\"EncryptionAlgorithm\":\"AES_CBC_256\"},\"ContentEncryptionIV\":\"testIVkey12345678910==\",\"KeyWrappingMetadata\":{\"EncryptionLibrary\":\"Java 5.3.0\"}}"
+	mockMatDesc := "{\"queryid\":\"01abc874-0406-1bf0-0000-53b10668e056\",\"smkid\":\"92019681909886\",\"key\":\"128\"}"
+	info := execResponseStageInfo{
+		Location:     "gcs/teststage/users/34/",
+		LocationType: "GCS",
+		Creds: execResponseCredentials{
+			GcsAccessToken: "test-token-124456577",
+		},
+	}
+	meta := fileMetadata{
+		client:    1,
+		stageInfo: &info,
+		mockGcsClient: &clientMock{
+			DoFunc: func(req *http.Request) (*http.Response, error) {
+				return &http.Response{
+					Status:     "200 OK",
+					StatusCode: 200,
+					Header: http.Header{
+						"X-Goog-Meta-Encryptiondata": []string{mockEncDataResp},
+						"Content-Length":             []string{"4256"},
+						"X-Goog-Meta-Sfc-Digest":     []string{"123456789abcdef"},
+						"X-Goog-Meta-Matdesc":        []string{mockMatDesc},
+					},
+				}, nil
+			},
+		},
+	}
+	_, err := new(snowflakeGcsClient).getFileHeader(&meta, "file.txt")
+	if err == nil {
+		t.Error("should have raised an error")
+	}
+}
+
 func TestUploadFileToGcsNoStatus(t *testing.T) {
 	info := execResponseStageInfo{
 		Location:     "gcs-blob/storage/users/456/",
@@ -961,3 +995,39 @@ func TestDownloadFileWithBadRequest(t *testing.T) {
 			renewPresignedURL, downloadMeta.resStatus)
 	}
 }
+
+func Test_snowflakeGcsClient_uploadFile(t *testing.T) {
+	info := execResponseStageInfo{
+		Location:     "gcs/teststage/users/34/",
+		LocationType: "GCS",
+		Creds: execResponseCredentials{
+			GcsAccessToken: "test-token-124456577",
+		},
+	}
+	meta := fileMetadata{
+		client:    1,
+		stageInfo: &info,
+	}
+	err := new(snowflakeGcsClient).uploadFile("somedata", &meta, nil, 1, 1)
+	if err == nil {
+		t.Error("should have raised an error")
+	}
+}
+
+func Test_snowflakeGcsClient_nativeDownloadFile(t *testing.T) {
+	info := execResponseStageInfo{
+		Location:     "gcs/teststage/users/34/",
+		LocationType: "GCS",
+		Creds: execResponseCredentials{
+			GcsAccessToken: "test-token-124456577",
+		},
+	}
+	meta := fileMetadata{
+		client:    1,
+		stageInfo: &info,
+	}
+	err := new(snowflakeGcsClient).nativeDownloadFile(&meta, "dummy data", 1)
+	if err == nil {
+		t.Error("should have raised an error")
+	}
+}