Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

SNOW-1524257 Implement GCM encryption #1191

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
+245 −49
Open

SNOW-1524257 Implement GCM encryption #1191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
236 changes: 197 additions & 39 deletions encrypt_util.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,19 +15,10 @@
"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 {
Expand All @@ -39,36 +30,30 @@
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,
chunkSize int) (*encryptMetadata, error) {
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)

Expand All @@ -88,23 +73,24 @@
}
// 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
}

Check warning on line 78 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L77-L78 

Added lines #L77 - L78 were not covered by tests
}

// 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
}

Check warning on line 87 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L86-L87 

Added lines #L86 - L87 were not covered by tests
}

// 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
}

Expand All @@ -120,7 +106,7 @@
}
return &encryptMetadata{
base64.StdEncoding.EncodeToString(encryptedFileKey),
base64.StdEncoding.EncodeToString(ivData),
base64.StdEncoding.EncodeToString(dataIv),
matDescUnicode,
}, nil
}
Expand Down Expand Up @@ -163,7 +149,7 @@
return nil
}

func encryptFile(
func encryptFileCBC(
sfe *snowflakeFileEncryption,
filename string,
chunkSize int,
Expand All @@ -183,14 +169,14 @@
}
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)
Expand Down Expand Up @@ -229,7 +215,7 @@
return mode, err
}

func decryptFile(
func decryptFileCBC(
metadata *encryptMetadata,
sfe *snowflakeFileEncryption,
filename string,
Expand All @@ -245,15 +231,15 @@
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
}
tmpOutputFile.Truncate(int64(totalFileSize))
return tmpOutputFile.Name(), nil
}

func decryptStream(
func decryptStreamCBC(
metadata *encryptMetadata,
sfe *snowflakeFileEncryption,
chunkSize int,
Expand All @@ -262,7 +248,7 @@
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
}
Expand Down Expand Up @@ -298,6 +284,149 @@
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
}

Check warning on line 291 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L290-L291 

Added lines #L290 - L291 were not covered by tests
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
}

Check warning on line 299 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L298-L299 

Added lines #L298 - L299 were not covered by tests
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
}

Check warning on line 307 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L306-L307 

Added lines #L306 - L307 were not covered by tests
return cipher.NewGCMWithNonceSize(block, 16)
}

func encryptFileGCM(
sfe *snowflakeFileEncryption,
filename string,
tmpDir string) (
*gcmEncryptMetadata, string, error) {
tmpOutputFile, err := os.CreateTemp(tmpDir, baseName(filename)+"#")
sfc-gh-dheyman-1 marked this conversation as resolved.
Show resolved Hide resolved
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we add "#" here?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be honest - I don't know, but we do the same in CBC encryption. The case is that such magic things might have solved some problems on some specific OSes or something in the past and I'm afraid to replace it.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After trying to figure this out for some time I guess it's just a way to prevent collisions in filenames. The tempfile uses the same filename as the infile so I guess it's to prevent overriding the infile in case the tmpDir is empty or something

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I generally agree, but whether it has to be that way or not is another question. I think that in 99% we can replace it with .bak or anything, but I'd leave it as is.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It will most likely not cause any problems, but in case some one, for god only know what reason, has files called my_file and my_file# it might cause a collision?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It feels like we should at least check if a file of this name exists

if err != nil {
return nil, "", err
}

Check warning on line 319 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L318-L319 

Added lines #L318 - L319 were not covered by tests
defer tmpOutputFile.Close()
infile, err := os.OpenFile(filename, os.O_CREATE|os.O_RDONLY, readWriteFileMode)
if err != nil {
return nil, "", err
}

Check warning on line 324 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L323-L324 

Added lines #L323 - L324 were not covered by tests
defer infile.Close()
plaintext, err := os.ReadFile(filename)
if err != nil {
return nil, "", err
}

Check warning on line 329 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L328-L329 

Added lines #L328 - L329 were not covered by tests

kek, err := base64.StdEncoding.DecodeString(sfe.QueryStageMasterKey)
if err != nil {
return nil, "", err
}

Check warning on line 334 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L333-L334 

Added lines #L333 - L334 were not covered by tests
keySize := len(kek)
fileKey := getSecureRandom(keySize)
keyIv := getSecureRandom(keySize)
encryptedFileKey, err := encryptGCM(keyIv, fileKey, kek, defaultKeyAad)
if err != nil {
return nil, "", err
}

Check warning on line 341 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L340-L341 

Added lines #L340 - L341 were not covered by tests

dataIv := getSecureRandom(keySize)
encryptedData, err := encryptGCM(dataIv, plaintext, fileKey, defaultDataAad)
if err != nil {
return nil, "", err
}

Check warning on line 347 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L346-L347 

Added lines #L346 - L347 were not covered by tests
_, err = tmpOutputFile.Write(encryptedData)
if err != nil {
return nil, "", err
}

Check warning on line 351 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L350-L351 

Added lines #L350 - L351 were not covered by tests

matDesc := materialDescriptor{
strconv.Itoa(int(sfe.SMKID)),
sfe.QueryID,
strconv.Itoa(keySize * 8),
}

matDescUnicode, err := matdescToUnicode(matDesc)
if err != nil {
return nil, "", err
}

Check warning on line 362 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L361-L362 

Added lines #L361 - L362 were not covered by tests
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
}

Check warning on line 383 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L382-L383 

Added lines #L382 - L383 were not covered by tests
encryptedFileKey, err := base64.StdEncoding.DecodeString(metadata.key)
if err != nil {
return "", err
}

Check warning on line 387 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L386-L387 

Added lines #L386 - L387 were not covered by tests
keyIv, err := base64.StdEncoding.DecodeString(metadata.keyIv)
if err != nil {
return "", err
}

Check warning on line 391 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L390-L391 

Added lines #L390 - L391 were not covered by tests
keyAad, err := base64.StdEncoding.DecodeString(metadata.keyAad)
if err != nil {
return "", err
}

Check warning on line 395 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L394-L395 

Added lines #L394 - L395 were not covered by tests
dataIv, err := base64.StdEncoding.DecodeString(metadata.dataIv)
if err != nil {
return "", err
}

Check warning on line 399 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L398-L399 

Added lines #L398 - L399 were not covered by tests
dataAad, err := base64.StdEncoding.DecodeString(metadata.dataAad)
if err != nil {
return "", err
}

Check warning on line 403 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L402-L403 

Added lines #L402 - L403 were not covered by tests

fileKey, err := decryptGCM(keyIv, encryptedFileKey, kek, keyAad)
if err != nil {
return "", err
}

Check warning on line 408 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L407-L408 

Added lines #L407 - L408 were not covered by tests

ciphertext, err := os.ReadFile(filename)
if err != nil {
return "", err
}

Check warning on line 413 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L412-L413 

Added lines #L412 - L413 were not covered by tests
plaintext, err := decryptGCM(dataIv, ciphertext, fileKey, dataAad)
if err != nil {
return "", err
}

Check warning on line 417 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L416-L417 

Added lines #L416 - L417 were not covered by tests

tmpOutputFile, err := os.CreateTemp(tmpDir, baseName(filename)+"#")
if err != nil {
return "", err
}

Check warning on line 422 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L421-L422 

Added lines #L421 - L422 were not covered by tests
tmpOutputFile.Write(plaintext)
if err != nil {
return "", err
}

Check warning on line 426 in encrypt_util.go

View check run for this annotation

Codecov / codecov/patch

encrypt_util.go#L425-L426 

Added lines #L425 - L426 were not covered by tests
return tmpOutputFile.Name(), nil
}

type materialDescriptor struct {
SmkID string `json:"smkId"`
QueryID string `json:"queryId"`
Expand Down Expand Up @@ -363,3 +492,32 @@
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
}
Loading
Loading