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.

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

chore: add KMS Keyring example in an multithreaded environment #735

Merged
merged 8 commits into from
Jan 15, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion AwsEncryptionSDK/runtimes/go/examples/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ require (
github.com/aws/aws-sdk-go-v2/service/dynamodb v1.35.1
github.com/aws/aws-sdk-go-v2/service/kms v1.36.0
github.com/aws/aws-sdk-go-v2/service/sts v1.31.1
github.com/google/uuid v1.6.0
)

require (
Expand All @@ -38,6 +39,5 @@ require (
github.com/aws/smithy-go v1.21.0 // indirect
github.com/dafny-lang/DafnyRuntimeGo/v4 v4.9.1 // indirect
github.com/dafny-lang/DafnyStandardLibGo v0.0.0 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/jmespath/go-jmespath v0.4.0 // indirect
)
9 changes: 8 additions & 1 deletion AwsEncryptionSDK/runtimes/go/examples/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,16 +23,18 @@ import (
"github.com/aws/aws-encryption-sdk/examples/keyring/rawaeskeyring"
"github.com/aws/aws-encryption-sdk/examples/keyring/rawrsakeyring"
"github.com/aws/aws-encryption-sdk/examples/misc"
"github.com/aws/aws-encryption-sdk/examples/multithreading"
"github.com/aws/aws-encryption-sdk/examples/utils"
)

func main() {
const stringToEncrypt = "Text To encrypt"
const numOfString = 10000
clientsupplier.ClientSupplierExample(
stringToEncrypt,
utils.DefaultRegionMrkKeyArn(),
utils.DefaultKMSKeyAccountID(),
[]string{"eu-west-1"})
[]string{utils.AlternateRegionMrkKeyRegion()})
misc.CommitmentPolicyExample(
stringToEncrypt,
utils.DefaultKMSKeyId(),
Expand Down Expand Up @@ -158,4 +160,9 @@ func main() {
utils.DefaultKMSKeyId(),
utils.DefaultKmsKeyRegion(),
)
// Example with multithreading
multithreading.AWSKMSMultiThreadTest(
utils.GenerateUUIDTestData(numOfString),
utils.DefaultKMSKeyId(),
utils.DefaultKmsKeyRegion())
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,193 @@
// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

/*
This example sets up the AWS KMS Keyring in an multithreaded environment.
The AWS KMS keyring uses symmetric encryption KMS keys to generate, encrypt and
decrypt data keys. This example creates a KMS Keyring and then encrypts a custom input exampleText
with an encryption context. This example also includes some sanity checks for demonstration:
1. Ciphertext and plaintext data are not the same
2. Decrypted plaintext value matches exampleText
These sanity checks are for demonstration in the example only. You do not need these in your code.
AWS KMS keyrings can be used independently or in a multi-keyring with other keyrings
of the same or a different type.
For more information on how to use KMS keyrings, see
https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/use-kms-keyring.html
For more information on KMS Key identifiers, see
https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id
*/

package multithreading

import (
"context"
"fmt"
"sync"

mpl "github.com/aws/aws-cryptographic-material-providers-library/mpl/awscryptographymaterialproviderssmithygenerated"
mpltypes "github.com/aws/aws-cryptographic-material-providers-library/mpl/awscryptographymaterialproviderssmithygeneratedtypes"
client "github.com/aws/aws-encryption-sdk/awscryptographyencryptionsdksmithygenerated"
esdktypes "github.com/aws/aws-encryption-sdk/awscryptographyencryptionsdksmithygeneratedtypes"
"github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/service/kms"
)

// Function to handle encryption
func encryptData(
ctx context.Context,
encryptionClient *client.Client,
plaintext string,
encryptionContext map[string]string,
keyring mpltypes.IKeyring) (*esdktypes.EncryptOutput, error) {
res, err := encryptionClient.Encrypt(ctx, esdktypes.EncryptInput{
Plaintext: []byte(plaintext),
EncryptionContext: encryptionContext,
Keyring: keyring,
})
return res, err
}

// Function to handle decryption
func decryptData(
ctx context.Context,
encryptionClient *client.Client,
ciphertext []byte,
encryptionContext map[string]string,
keyring mpltypes.IKeyring) (*esdktypes.DecryptOutput, error) {
res, err := encryptionClient.Decrypt(ctx, esdktypes.DecryptInput{
EncryptionContext: encryptionContext,
Keyring: keyring,
Ciphertext: ciphertext,
})
return res, err
}

func processEncryptionWorker(
ctx context.Context,
wg *sync.WaitGroup,
jobs <-chan string,
encryptionClient *client.Client,
awsKmsKeyring mpltypes.IKeyring,
encryptionContext map[string]string,
) {
defer wg.Done()
for plaintext := range jobs {
// Perform encryption
encryptResult, err := encryptData(
ctx,
encryptionClient,
plaintext,
encryptionContext,
awsKmsKeyring)
if err != nil {
panic(err)
}
// Verify ciphertext is different from plaintext
if string(encryptResult.Ciphertext) == plaintext {
panic("Ciphertext and Plaintext before encryption are the same")
}
// Perform decryption
decryptResult, err := decryptData(
ctx,
encryptionClient,
encryptResult.Ciphertext,
encryptionContext,
awsKmsKeyring,
)
if err != nil {
panic(err)
}
// If you do not specify the encryption context on Decrypt, it's recommended to check if the resulting encryption context matches.
// The encryption context was specified on decrypt; we are validating the encryption context for demonstration only.
// Before your application uses plaintext data, verify that the encryption context that
// you used to encrypt the message is included in the encryption context that was used to
// decrypt the message. The AWS Encryption SDK can add pairs, so don't require an exact match.
if err := validateEncryptionContext(encryptionContext, decryptResult.EncryptionContext); err != nil {
panic(err)
}
if string(decryptResult.Plaintext) != plaintext {
panic("Plaintext after decryption and Plaintext before encryption are NOT the same")
}
}
}

func AWSKMSMultiThreadTest(texts []string, defaultKmsKeyID, defaultKmsKeyRegion string) {
// Create the AWS KMS client
cfg, err := config.LoadDefaultConfig(context.TODO())
if err != nil {
panic(err)
}
kmsClient := kms.NewFromConfig(cfg, func(o *kms.Options) {
o.Region = defaultKmsKeyRegion
})
// Initialize the mpl client
matProv, err := mpl.NewClient(mpltypes.MaterialProvidersConfig{})
if err != nil {
panic(err)
}
// Create the keyring
ctx := context.Background()
awsKmsKeyringInput := mpltypes.CreateAwsKmsKeyringInput{
KmsClient: kmsClient,
KmsKeyId: defaultKmsKeyID,
}
awsKmsKeyring, err := matProv.CreateAwsKmsKeyring(ctx, awsKmsKeyringInput)
if err != nil {
panic(err)
}
// Instantiate the encryption SDK client.
// This builds the default client with the RequireEncryptRequireDecrypt commitment policy,
// which enforces that this client only encrypts using committing algorithm suites and enforces
// that this client will only decrypt encrypted messages that were created with a committing
// algorithm suite.
encryptionClient, err := client.NewClient(esdktypes.AwsEncryptionSdkConfig{})
if err != nil {
panic(err)
}
// Create your encryption context (Optional).
// Remember that your encryption context is NOT SECRET.
// https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#encryption-context
encryptionContext := map[string]string{
"encryption": "context",
"is not": "secret",
"but adds": "useful metadata",
"that can help you": "be confident that",
"the data you are handling": "is what you think it is",
}
// Create buffered channels to handle multiple operations
// As an example, we will have 10 workers, adjust this number as needed.
numWorkers := 10

// Create a wait group to track all goroutines
var wg sync.WaitGroup

// Create a channel to send a plaintext
jobs := make(chan string, len(texts))

// Start worker pool
for range numWorkers {
wg.Add(1)
go processEncryptionWorker(ctx, &wg, jobs, encryptionClient, awsKmsKeyring, encryptionContext)
}

// Send jobs to workers
for _, text := range texts {
jobs <- text
}
close(jobs)
// Wait for all workers to complete
wg.Wait()
fmt.Println("AWS KMS Keyring example in multithreaded environment completed successfully.")
texastony marked this conversation as resolved.
Show resolved Hide resolved
}

// This function only does subset matching because AWS Encryption SDK can add pairs, so don't require an exact match.
func validateEncryptionContext(expected, actual map[string]string) error {
for expectedKey, expectedValue := range expected {
actualValue, exists := actual[expectedKey]
if !exists || actualValue != expectedValue {
return fmt.Errorf("encryption context mismatch: expected key '%s' with value '%s'",
expectedKey, expectedValue)
}
}
return nil
}
15 changes: 15 additions & 0 deletions AwsEncryptionSDK/runtimes/go/examples/utils/exampleUtils.go
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package utils

import (
Expand All @@ -13,6 +16,7 @@ import (

"github.com/aws/aws-cryptographic-material-providers-library/primitives/awscryptographyprimitivessmithygeneratedtypes"
"github.com/aws/aws-sdk-go-v2/service/kms"
"github.com/google/uuid"
)

const (
Expand Down Expand Up @@ -319,3 +323,14 @@ func GenerateKmsEccPublicKey(eccKeyArn string, kmsClient *kms.Client) ([]byte, e
}
return response.PublicKey, nil
}

// GenerateUUIDTestData creates an array of random UUID strings
func GenerateUUIDTestData(count int) []string {
testData := make([]string, count)
for i := 0; i < count; i++ {
// Generate a random UUID
uuid := uuid.New()
testData[i] = uuid.String()
}
return testData
}
Loading