Repository: snowflakedb/gosnowflake
Pull request #1191 (2024-08-01 12:34:08)

SNOW-1524257 Implement GCM encryption
### Description

SNOW-1524257 Implemented GCM encryption to be ready for CSP integrations.

### Checklist
- [ ] Created tests which fail without the change (if possible)
- [ ] Extended the README / documentation, if necessary



Patch:
diff --git a/encrypt_util.go b/encrypt_util.go
index 481ebc7..3c6b186 100644
--- a/encrypt_util.go
+++ b/encrypt_util.go
@@ -15,19 +15,10 @@ import (
 	"strconv"
 )
 
-type snowflakeFileEncryption struct {
-	QueryStageMasterKey string `json:"queryStageMasterKey,omitempty"`
-	QueryID             string `json:"queryId,omitempty"`
-	SMKID               int64  `json:"smkId,omitempty"`
-}
-
-// PUT requests return a single encryptionMaterial object whereas GET requests
-// return a slice (array) of encryptionMaterial objects, both under the field
-// 'encryptionMaterial'
-type encryptionWrapper struct {
-	snowflakeFileEncryption
-	EncryptionMaterials []snowflakeFileEncryption
-}
+var (
+	defaultKeyAad  = make([]byte, 0)
+	defaultDataAad = make([]byte, 0)
+)
 
 // override default behavior for wrapper
 func (ew *encryptionWrapper) UnmarshalJSON(data []byte) error {
@@ -39,15 +30,9 @@ func (ew *encryptionWrapper) UnmarshalJSON(data []byte) error {
 	return json.Unmarshal(data, &ew.snowflakeFileEncryption)
 }
 
-type encryptMetadata struct {
-	key     string
-	iv      string
-	matdesc string
-}
-
-// encryptStream encrypts a stream buffer using AES128 block cipher in CBC mode
+// encryptStreamCBC encrypts a stream buffer using AES128 block cipher in CBC mode
 // with PKCS5 padding
-func encryptStream(
+func encryptStreamCBC(
 	sfe *snowflakeFileEncryption,
 	src io.Reader,
 	out io.Writer,
@@ -55,20 +40,20 @@ func encryptStream(
 	if chunkSize == 0 {
 		chunkSize = aes.BlockSize * 4 * 1024
 	}
-	decodedKey, err := base64.StdEncoding.DecodeString(sfe.QueryStageMasterKey)
+	kek, err := base64.StdEncoding.DecodeString(sfe.QueryStageMasterKey)
 	if err != nil {
 		return nil, err
 	}
-	keySize := len(decodedKey)
+	keySize := len(kek)
 
 	fileKey := getSecureRandom(keySize)
 	block, err := aes.NewCipher(fileKey)
 	if err != nil {
 		return nil, err
 	}
-	ivData := getSecureRandom(block.BlockSize())
+	dataIv := getSecureRandom(block.BlockSize())
 
-	mode := cipher.NewCBCEncrypter(block, ivData)
+	mode := cipher.NewCBCEncrypter(block, dataIv)
 	cipherText := make([]byte, chunkSize)
 	chunk := make([]byte, chunkSize)
 
@@ -88,23 +73,24 @@ func encryptStream(
 		}
 		// make sure only n bytes of chunk is used
 		mode.CryptBlocks(cipherText, chunk[:n])
-		out.Write(cipherText[:n])
-	}
-	if err != nil {
-		return nil, err
+		if _, err := out.Write(cipherText[:n]); err != nil {
+			return nil, err
+		}
 	}
 
 	// add padding if not yet added
 	if !padded {
 		padding := bytes.Repeat([]byte(string(rune(aes.BlockSize))), aes.BlockSize)
 		mode.CryptBlocks(cipherText, padding)
-		out.Write(cipherText[:len(padding)])
+		if _, err := out.Write(cipherText[:len(padding)]); err != nil {
+			return nil, err
+		}
 	}
 
 	// encrypt key with ECB
 	fileKey = padBytesLength(fileKey, block.BlockSize())
 	encryptedFileKey := make([]byte, len(fileKey))
-	if err = encryptECB(encryptedFileKey, fileKey, decodedKey); err != nil {
+	if err = encryptECB(encryptedFileKey, fileKey, kek); err != nil {
 		return nil, err
 	}
 
@@ -120,7 +106,7 @@ func encryptStream(
 	}
 	return &encryptMetadata{
 		base64.StdEncoding.EncodeToString(encryptedFileKey),
-		base64.StdEncoding.EncodeToString(ivData),
+		base64.StdEncoding.EncodeToString(dataIv),
 		matDescUnicode,
 	}, nil
 }
@@ -163,7 +149,7 @@ func decryptECB(decrypted []byte, keyBytes []byte, decodedKey []byte) error {
 	return nil
 }
 
-func encryptFile(
+func encryptFileCBC(
 	sfe *snowflakeFileEncryption,
 	filename string,
 	chunkSize int,
@@ -183,14 +169,14 @@ func encryptFile(
 	}
 	defer infile.Close()
 
-	meta, err := encryptStream(sfe, infile, tmpOutputFile, chunkSize)
+	meta, err := encryptStreamCBC(sfe, infile, tmpOutputFile, chunkSize)
 	if err != nil {
 		return nil, "", err
 	}
 	return meta, tmpOutputFile.Name(), nil
 }
 
-func decryptFileKey(
+func decryptFileKeyECB(
 	metadata *encryptMetadata,
 	sfe *snowflakeFileEncryption) ([]byte, []byte, error) {
 	decodedKey, err := base64.StdEncoding.DecodeString(sfe.QueryStageMasterKey)
@@ -229,7 +215,7 @@ func initCBC(decryptedKey []byte, ivBytes []byte) (cipher.BlockMode, error) {
 	return mode, err
 }
 
-func decryptFile(
+func decryptFileCBC(
 	metadata *encryptMetadata,
 	sfe *snowflakeFileEncryption,
 	filename string,
@@ -245,7 +231,7 @@ func decryptFile(
 		return "", err
 	}
 	defer infile.Close()
-	totalFileSize, err := decryptStream(metadata, sfe, chunkSize, infile, tmpOutputFile)
+	totalFileSize, err := decryptStreamCBC(metadata, sfe, chunkSize, infile, tmpOutputFile)
 	if err != nil {
 		return "", err
 	}
@@ -253,7 +239,7 @@ func decryptFile(
 	return tmpOutputFile.Name(), nil
 }
 
-func decryptStream(
+func decryptStreamCBC(
 	metadata *encryptMetadata,
 	sfe *snowflakeFileEncryption,
 	chunkSize int,
@@ -262,7 +248,7 @@ func decryptStream(
 	if chunkSize == 0 {
 		chunkSize = aes.BlockSize * 4 * 1024
 	}
-	decryptedKey, ivBytes, err := decryptFileKey(metadata, sfe)
+	decryptedKey, ivBytes, err := decryptFileKeyECB(metadata, sfe)
 	if err != nil {
 		return 0, err
 	}
@@ -298,6 +284,149 @@ func decryptStream(
 	return totalFileSize, err
 }
 
+func encryptGCM(iv []byte, plaintext []byte, encryptionKey []byte, aad []byte) ([]byte, error) {
+	aead, err := initGcm(encryptionKey)
+	if err != nil {
+		return nil, err
+	}
+	return aead.Seal(nil, iv, plaintext, aad), nil
+}
+
+func decryptGCM(iv []byte, ciphertext []byte, encryptionKey []byte, aad []byte) ([]byte, error) {
+	aead, err := initGcm(encryptionKey)
+	if err != nil {
+		return nil, err
+	}
+	return aead.Open(nil, iv, ciphertext, aad)
+}
+
+func initGcm(encryptionKey []byte) (cipher.AEAD, error) {
+	block, err := aes.NewCipher(encryptionKey)
+	if err != nil {
+		return nil, err
+	}
+	return cipher.NewGCMWithNonceSize(block, 16)
+}
+
+func encryptFileGCM(
+	sfe *snowflakeFileEncryption,
+	filename string,
+	tmpDir string) (
+	*gcmEncryptMetadata, string, error) {
+	tmpOutputFile, err := os.CreateTemp(tmpDir, baseName(filename)+"#")
+	if err != nil {
+		return nil, "", err
+	}
+	defer tmpOutputFile.Close()
+	infile, err := os.OpenFile(filename, os.O_CREATE|os.O_RDONLY, readWriteFileMode)
+	if err != nil {
+		return nil, "", err
+	}
+	defer infile.Close()
+	plaintext, err := os.ReadFile(filename)
+	if err != nil {
+		return nil, "", err
+	}
+
+	kek, err := base64.StdEncoding.DecodeString(sfe.QueryStageMasterKey)
+	if err != nil {
+		return nil, "", err
+	}
+	keySize := len(kek)
+	fileKey := getSecureRandom(keySize)
+	keyIv := getSecureRandom(keySize)
+	encryptedFileKey, err := encryptGCM(keyIv, fileKey, kek, defaultKeyAad)
+	if err != nil {
+		return nil, "", err
+	}
+
+	dataIv := getSecureRandom(keySize)
+	encryptedData, err := encryptGCM(dataIv, plaintext, fileKey, defaultDataAad)
+	if err != nil {
+		return nil, "", err
+	}
+	_, err = tmpOutputFile.Write(encryptedData)
+	if err != nil {
+		return nil, "", err
+	}
+
+	matDesc := materialDescriptor{
+		strconv.Itoa(int(sfe.SMKID)),
+		sfe.QueryID,
+		strconv.Itoa(keySize * 8),
+	}
+
+	matDescUnicode, err := matdescToUnicode(matDesc)
+	if err != nil {
+		return nil, "", err
+	}
+	meta := &gcmEncryptMetadata{
+		key:     base64.StdEncoding.EncodeToString(encryptedFileKey),
+		keyIv:   base64.StdEncoding.EncodeToString(keyIv),
+		dataIv:  base64.StdEncoding.EncodeToString(dataIv),
+		keyAad:  base64.StdEncoding.EncodeToString(defaultKeyAad),
+		dataAad: base64.StdEncoding.EncodeToString(defaultDataAad),
+		matdesc: matDescUnicode,
+	}
+	return meta, tmpOutputFile.Name(), nil
+}
+
+func decryptFileGCM(
+	metadata *gcmEncryptMetadata,
+	sfe *snowflakeFileEncryption,
+	filename string,
+	tmpDir string) (
+	string, error) {
+	kek, err := base64.StdEncoding.DecodeString(sfe.QueryStageMasterKey)
+	if err != nil {
+		return "", err
+	}
+	encryptedFileKey, err := base64.StdEncoding.DecodeString(metadata.key)
+	if err != nil {
+		return "", err
+	}
+	keyIv, err := base64.StdEncoding.DecodeString(metadata.keyIv)
+	if err != nil {
+		return "", err
+	}
+	keyAad, err := base64.StdEncoding.DecodeString(metadata.keyAad)
+	if err != nil {
+		return "", err
+	}
+	dataIv, err := base64.StdEncoding.DecodeString(metadata.dataIv)
+	if err != nil {
+		return "", err
+	}
+	dataAad, err := base64.StdEncoding.DecodeString(metadata.dataAad)
+	if err != nil {
+		return "", err
+	}
+
+	fileKey, err := decryptGCM(keyIv, encryptedFileKey, kek, keyAad)
+	if err != nil {
+		return "", err
+	}
+
+	ciphertext, err := os.ReadFile(filename)
+	if err != nil {
+		return "", err
+	}
+	plaintext, err := decryptGCM(dataIv, ciphertext, fileKey, dataAad)
+	if err != nil {
+		return "", err
+	}
+
+	tmpOutputFile, err := os.CreateTemp(tmpDir, baseName(filename)+"#")
+	if err != nil {
+		return "", err
+	}
+	tmpOutputFile.Write(plaintext)
+	if err != nil {
+		return "", err
+	}
+	return tmpOutputFile.Name(), nil
+}
+
 type materialDescriptor struct {
 	SmkID   string `json:"smkId"`
 	QueryID string `json:"queryId"`
@@ -363,3 +492,32 @@ type encryptionData struct {
 	ContentEncryptionIV string          `json:"ContentEncryptionIV,omitempty"`
 	KeyWrappingMetadata keyMetadata     `json:"KeyWrappingMetadata,omitempty"`
 }
+
+type snowflakeFileEncryption struct {
+	QueryStageMasterKey string `json:"queryStageMasterKey,omitempty"`
+	QueryID             string `json:"queryId,omitempty"`
+	SMKID               int64  `json:"smkId,omitempty"`
+}
+
+// PUT requests return a single encryptionMaterial object whereas GET requests
+// return a slice (array) of encryptionMaterial objects, both under the field
+// 'encryptionMaterial'
+type encryptionWrapper struct {
+	snowflakeFileEncryption
+	EncryptionMaterials []snowflakeFileEncryption
+}
+
+type encryptMetadata struct {
+	key     string
+	iv      string
+	matdesc string
+}
+
+type gcmEncryptMetadata struct {
+	key     string
+	keyIv   string
+	dataIv  string
+	keyAad  string
+	dataAad string
+	matdesc string
+}
diff --git a/storage_client.go b/storage_client.go
index 2079fa8..f1be2b4 100644
--- a/storage_client.go
+++ b/storage_client.go
@@ -65,14 +65,14 @@ func (rsu *remoteStorageUtil) uploadOneFile(meta *fileMetadata) error {
 			if meta.realSrcStream != nil {
 				srcStream = meta.realSrcStream
 			}
-			encryptMeta, err = encryptStream(meta.encryptionMaterial, srcStream, &encryptedStream, 0)
+			encryptMeta, err = encryptStreamCBC(meta.encryptionMaterial, srcStream, &encryptedStream, 0)
 			if err != nil {
 				return err
 			}
 			meta.realSrcStream = &encryptedStream
 			dataFile = meta.realSrcFileName
 		} else {
-			encryptMeta, dataFile, err = encryptFile(meta.encryptionMaterial, meta.realSrcFileName, 0, meta.tmpDir)
+			encryptMeta, dataFile, err = encryptFileCBC(meta.encryptionMaterial, meta.realSrcFileName, 0, meta.tmpDir)
 			if err != nil {
 				return err
 			}
@@ -219,7 +219,7 @@ func (rsu *remoteStorageUtil) downloadOneFile(meta *fileMetadata) error {
 					}
 				}
 				if meta.options.GetFileToStream {
-					totalFileSize, err := decryptStream(header.encryptionMetadata,
+					totalFileSize, err := decryptStreamCBC(header.encryptionMetadata,
 						meta.encryptionMaterial, 0, meta.dstStream, meta.sfa.streamBuffer)
 					if err != nil {
 						return err
@@ -227,7 +227,7 @@ func (rsu *remoteStorageUtil) downloadOneFile(meta *fileMetadata) error {
 					meta.sfa.streamBuffer.Truncate(totalFileSize)
 					meta.dstFileSize = int64(totalFileSize)
 				} else {
-					tmpDstFileName, err := decryptFile(header.encryptionMetadata,
+					tmpDstFileName, err := decryptFileCBC(header.encryptionMetadata,
 						meta.encryptionMaterial, fullDstFileName, 0, meta.tmpDir)
 					if err != nil {
 						return err


Test patch:
diff --git a/encrypt_util_test.go b/encrypt_util_test.go
index 03f075d..1cbae85 100644
--- a/encrypt_util_test.go
+++ b/encrypt_util_test.go
@@ -5,6 +5,7 @@ package gosnowflake
 import (
 	"bufio"
 	"compress/gzip"
+	"encoding/base64"
 	"fmt"
 	"io"
 	"math/rand"
@@ -24,7 +25,7 @@ type encryptDecryptTestFile struct {
 	numberOfLines          int
 }
 
-func TestEncryptDecryptFile(t *testing.T) {
+func TestEncryptDecryptFileCBC(t *testing.T) {
 	encMat := snowflakeFileEncryption{
 		"ztke8tIdVt1zmlQIZm0BMA==",
 		"123873c7-3a66-40c4-ab89-e3722fbccce1",
@@ -43,12 +44,12 @@ func TestEncryptDecryptFile(t *testing.T) {
 		t.Error(err)
 	}
 
-	metadata, encryptedFile, err := encryptFile(&encMat, inputFile, 0, "")
+	metadata, encryptedFile, err := encryptFileCBC(&encMat, inputFile, 0, "")
 	if err != nil {
 		t.Error(err)
 	}
 	defer os.Remove(encryptedFile)
-	decryptedFile, err := decryptFile(metadata, &encMat, encryptedFile, 0, "")
+	decryptedFile, err := decryptFileCBC(metadata, &encMat, encryptedFile, 0, "")
 	if err != nil {
 		t.Error(err)
 	}
@@ -101,7 +102,7 @@ func TestEncryptDecryptFilePadding(t *testing.T) {
 	}
 }
 
-func TestEncryptDecryptLargeFile(t *testing.T) {
+func TestEncryptDecryptLargeFileCBC(t *testing.T) {
 	encMat := snowflakeFileEncryption{
 		"ztke8tIdVt1zmlQIZm0BMA==",
 		"123873c7-3a66-40c4-ab89-e3722fbccce1",
@@ -125,12 +126,12 @@ func encryptDecryptFile(t *testing.T, encMat snowflakeFileEncryption, expected i
 	}
 	inputFile := files[0]
 
-	metadata, encryptedFile, err := encryptFile(&encMat, inputFile, 0, tmpDir)
+	metadata, encryptedFile, err := encryptFileCBC(&encMat, inputFile, 0, tmpDir)
 	if err != nil {
 		t.Error(err)
 	}
 	defer os.Remove(encryptedFile)
-	decryptedFile, err := decryptFile(metadata, &encMat, encryptedFile, 0, tmpDir)
+	decryptedFile, err := decryptFileCBC(metadata, &encMat, encryptedFile, 0, tmpDir)
 	if err != nil {
 		t.Error(err)
 	}
@@ -236,3 +237,40 @@ func generateKLinesOfNFiles(k int, n int, compress bool, tmpDir string) (string,
 	}
 	return tmpDir, nil
 }
+
+func TestEncryptDecryptGCM(t *testing.T) {
+	input := []byte("abc")
+	iv := []byte("abcdef1234567890")  // pragma: allowlist secret
+	key := []byte("1234567890abcdef") // pragma: allowlist secret
+	encrypted, err := encryptGCM(iv, input, key, nil)
+	assertNilF(t, err)
+	assertEqualE(t, base64.StdEncoding.EncodeToString(encrypted), "pgs/wjNH2TYekmN7mbhFjeHH0A==")
+
+	decrypted, err := decryptGCM(iv, encrypted, key, nil)
+	assertNilF(t, err)
+	assertDeepEqualE(t, decrypted, input)
+}
+
+func TestEncryptDecryptFileGCM(t *testing.T) {
+	tmpDir := os.TempDir()
+	tempFile, err := os.CreateTemp(tmpDir, "gcm")
+	assertNilF(t, err)
+	_, err = tempFile.Write([]byte("abc"))
+	assertNilF(t, err)
+
+	sfe := &snowflakeFileEncryption{
+		QueryStageMasterKey: "YWJjZGVmMTIzNDU2Nzg5MA==",
+		QueryID:             "unused",
+		SMKID:               123,
+	}
+	meta, encryptedFileName, err := encryptFileGCM(sfe, tempFile.Name(), tmpDir)
+	assertNilF(t, err)
+
+	decryptedFileName, err := decryptFileGCM(meta, sfe, encryptedFileName, tmpDir)
+	assertNilF(t, err)
+
+	fileContent, err := os.ReadFile(decryptedFileName)
+	assertNilF(t, err)
+
+	assertEqualE(t, string(fileContent), "abc")
+}