Skip to content

feat(dafny): use kms:GenerateDataKey for branch key creation and version hv-2 #1478

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

Merged
merged 8 commits into from
Apr 28, 2025
Merged

feat(dafny): use kms:GenerateDataKey for branch key creation and version hv-2 #1478

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0087862
feat(dafny): use kms::GenerateDataKey for hv-2 branch key creation
josecorella Apr 24, 2025
429b229
use ec to kms for generate data key
josecorella Apr 24, 2025
25dcfd3
work
josecorella Apr 26, 2025
a1efe60
move version key to use generatedatakey
josecorella Apr 26, 2025
e9eace8
add proof todo
josecorella Apr 26, 2025
751b9ad
Update AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStor…
josecorella Apr 28, 2025
079180d
Apply suggestions from code review
josecorella Apr 28, 2025
0ccebe4
format
josecorella Apr 28, 2025
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
98 changes: 59 additions & 39 deletions AwsCryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/CreateKeys.dfy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -316,6 +316,25 @@ module {:options "/functionSyntax:4" } CreateKeys {
&& forall k <- encryptionContext :: DDB.IsValid_AttributeName(Structure.ENCRYPTION_CONTEXT_PREFIX + k)
modifies keyManagerAndStorage.Modifies
ensures keyManagerAndStorage.ValidState()
ensures output.Success?
==>
&& var decryptBKC := Structure.DecryptOnlyBranchKeyEncryptionContext(
branchKeyIdentifier,
branchKeyVersion,
timestamp,
logicalKeyStoreName,
KMSKeystoreOperations.GetKeyId(kmsConfiguration),
hierarchyVersion,
encryptionContext
);
&& var activeBKC := Structure.ActiveBranchKeyEncryptionContext(decryptBKC);
&& var beaconBKC := Structure.BeaconKeyEncryptionContext(decryptBKC);
&& HvUtils.HasUniqueTransformedKeys?(activeBKC) == true
&& var ecToKms := HvUtils.SelectKmsEncryptionContextForHv2(activeBKC);
&& var kms := keyManagerAndStorage.keyManagerStrat.kmsSimple.kmsClient;
&& |kms.History.GenerateDataKey| == |old(kms.History.GenerateDataKey)| + 2 // 2 since we call to generate the branch key and the beacon key
&& |kms.History.Encrypt| == |old(kms.History.Encrypt)| + 3 // 3 since we encrypt the active, version, and the beacon key
//TODO: add generate data key and encrypt proofs here
{
// Construct Branch Key Contexts for ACTIVE, Version and Beacon items.
var decryptOnlyBranchKeyContext := Structure.DecryptOnlyBranchKeyEncryptionContext(
Expand All @@ -330,6 +349,27 @@ module {:options "/functionSyntax:4" } CreateKeys {
var activeBranchKeyContext := Structure.ActiveBranchKeyEncryptionContext(decryptOnlyBranchKeyContext);
var beaconBranchKeyContext := Structure.BeaconKeyEncryptionContext(decryptOnlyBranchKeyContext);

if !HvUtils.HasUniqueTransformedKeys?(activeBranchKeyContext) {
return Failure(Types.BranchKeyCiphertextException(
message := ErrorMessages.NOT_UNIQUE_BRANCH_KEY_CONTEXT_KEYS
));
}
var ecToKMS := HvUtils.SelectKmsEncryptionContextForHv2(activeBranchKeyContext);

// get plaintext data key by calling kms::GenerateDataKey
var activePlaintextMaterial :- KMSKeystoreOperations.GetPlaintextDataKeyViaGenerateDataKey(
branchKeyContext := ecToKMS,
kmsConfiguration := kmsConfiguration,
keyManagerAndStorage := keyManagerAndStorage
);

// get beacon key by calling kms::GenerateDataKey
var beaconPlaintextMaterial :- KMSKeystoreOperations.GetPlaintextDataKeyViaGenerateDataKey(
branchKeyContext := ecToKMS,
kmsConfiguration := kmsConfiguration,
keyManagerAndStorage := keyManagerAndStorage
);

// Get crypto client
var crypto? := HvUtils.ProvideCryptoClient();
var crypto :- crypto?.MapFailure(
Expand All @@ -338,23 +378,6 @@ module {:options "/functionSyntax:4" } CreateKeys {
)
);

// Generate Random Bytes as Plaintext for ACTIVE & Beacon Item's
// TODO-HV-2-M4 : Improve error messages for generate random failure
var activePlaintextMaterial? := crypto.GenerateRandomBytes(
AtomicPrimitives.Types.GenerateRandomBytesInput(length := 32)
);
var activePlaintextMaterial :- activePlaintextMaterial?.MapFailure(
e => Types.AwsCryptographyPrimitives(
AwsCryptographyPrimitives := e
));

var beaconPlaintextMaterial? := crypto.GenerateRandomBytes(
AtomicPrimitives.Types.GenerateRandomBytesInput(length := 32)
);
var beaconPlaintextMaterial :- beaconPlaintextMaterial?.MapFailure(
e => Types.AwsCryptographyPrimitives(
AwsCryptographyPrimitives := e));

var CryptoAndKms := KMSKeystoreOperations.CryptoAndKms(kmsConfiguration, keyManagerAndStorage.keyManagerStrat, crypto);
var decryptOnlyBKItem :- KMSKeystoreOperations.packAndCallKMS(
branchKeyContext := decryptOnlyBranchKeyContext,
Expand Down Expand Up @@ -385,15 +408,7 @@ module {:options "/functionSyntax:4" } CreateKeys {
Beacon := beaconBKItem
)
);
// By making this ghost, we avoid the Dafny transpiler emitting in Java etc.
ghost var storage := keyManagerAndStorage.storage;
assert |storage.History.WriteNewEncryptedBranchKey| == |old(storage.History.WriteNewEncryptedBranchKey)| + 1;
ghost var writeEvent := Seq.Last(storage.History.WriteNewEncryptedBranchKey);
assert
&& writeEvent.output.Success?
&& writeEvent.input.Active == activeBKItem
&& writeEvent.input.Version == decryptOnlyBKItem
&& writeEvent.input.Beacon == beaconBKItem;

output := Success(
Types.CreateKeyOutput(
branchKeyIdentifier := branchKeyIdentifier
Expand Down Expand Up @@ -707,6 +722,7 @@ module {:options "/functionSyntax:4" } CreateKeys {
requires kmsClient.ValidState() && storage.ValidState()
modifies storage.Modifies, kmsClient.Modifies
ensures storage.ValidState() && kmsClient.ValidState()
//TODO: add generate data key and encrypt proofs here
{
if !HvUtils.HasUniqueTransformedKeys?(oldActiveItem.EncryptionContext) {
return Failure(Types.BranchKeyCiphertextException(
Expand All @@ -720,6 +736,23 @@ module {:options "/functionSyntax:4" } CreateKeys {
Types.KeyStoreException(
message := ErrorMessages.VERSION_KEY_KMS_KEY_ARN_DISAGREEMENT)
);

var ecToKMS := HvUtils.SelectKmsEncryptionContextForHv2(oldActiveItem.EncryptionContext);

var keyManagerStrategy := KmsUtils.keyManagerStrat.kmsSimple(
kmsSimple := KmsUtils.KMSTuple(kmsClient := kmsClient, grantTokens := grantTokens)
);
var keyManagerAndStorage := KmsUtils.KeyManagerAndStorage(
storage := storage,
keyManagerStrat := keyManagerStrategy
);

var newActivePlaintextMaterial :- KMSKeystoreOperations.GetPlaintextDataKeyViaGenerateDataKey(
branchKeyContext := ecToKMS,
kmsConfiguration := kmsConfiguration,
keyManagerAndStorage := keyManagerAndStorage
);

// Get crypto client
var crypto? := HvUtils.ProvideCryptoClient();
var crypto :- crypto?.MapFailure(
Expand All @@ -728,16 +761,6 @@ module {:options "/functionSyntax:4" } CreateKeys {
)
);

var newActivePlaintextMaterial? := crypto.GenerateRandomBytes(
AtomicPrimitives.Types.GenerateRandomBytesInput(length := 32)
);
var newActivePlaintextMaterial :- newActivePlaintextMaterial?.MapFailure(
e => Types.AwsCryptographyPrimitives(
AwsCryptographyPrimitives := e
));

var ecToKMS := HvUtils.SelectKmsEncryptionContextForHv2(oldActiveItem.EncryptionContext);

// we decrypt the oldActiveItem to get the ciphertext and then
// we encrypt it to a versioned key
var oldActiveItemPlaintext :- GetKeys.DecryptAndValidateKeyForHV2(
Expand All @@ -754,9 +777,6 @@ module {:options "/functionSyntax:4" } CreateKeys {
);
var activeEncryptionContext := Structure.ActiveBranchKeyEncryptionContext(decryptOnlyEncryptionContext);

var keyManagerStrategy := KmsUtils.keyManagerStrat.kmsSimple(
kmsSimple := KmsUtils.KMSTuple(kmsClient := kmsClient, grantTokens := grantTokens)
);

var CryptoAndKms := KMSKeystoreOperations.CryptoAndKms(kmsConfiguration, keyManagerStrategy, crypto);
var wrappedDecryptOnlyBranchKey :- KMSKeystoreOperations.packAndCallKMS(
Expand Down
69 changes: 69 additions & 0 deletions ...ryptographicMaterialProviders/dafny/AwsCryptographyKeyStore/src/KMSKeystoreOperations.dfy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,73 @@ module {:options "/functionSyntax:4" } KMSKeystoreOperations {
case kmsMRKeyArn(arn) => arn
}

method GetPlaintextDataKeyViaGenerateDataKey(
nameonly branchKeyContext: map<string, string>,
nameonly kmsConfiguration: Types.KMSConfiguration,
nameonly keyManagerAndStorage: KmsUtils.KeyManagerAndStorage
)
returns (output: Result<seq<uint8>, Types.Error>)
requires
// TODO-HV2-DecryptEncrypt support Decrypt/Encrypt
&& keyManagerAndStorage.keyManagerStrat.kmsSimple?
&& keyManagerAndStorage.ValidState()
&& HasKeyId(kmsConfiguration)
&& KmsArn.ValidKmsArn?(GetKeyId(kmsConfiguration))
modifies keyManagerAndStorage.Modifies
ensures keyManagerAndStorage.ValidState()
ensures output.Success?
==>
&& var kms := keyManagerAndStorage.keyManagerStrat.kmsSimple.kmsClient;
&& |kms.History.GenerateDataKey| == |old(kms.History.GenerateDataKey)| + 1
&& old(kms.History.Encrypt) == kms.History.Encrypt
&& var kmsGenerateDataKeyEvent := Seq.Last(kms.History.GenerateDataKey);
&& var kmsKeyArn := GetKeyId(kmsConfiguration);
&& KMS.GenerateDataKeyRequest(
KeyId := kmsKeyArn,
NumberOfBytes := Some(32),
EncryptionContext := Some(branchKeyContext),
GrantTokens := Some(keyManagerAndStorage.keyManagerStrat.kmsSimple.grantTokens)
) == kmsGenerateDataKeyEvent.input
&& kmsGenerateDataKeyEvent.output.Success?
&& kmsGenerateDataKeyEvent.output.value.Plaintext.Some?
&& |kmsGenerateDataKeyEvent.output.value.Plaintext.value| == 32
&& kmsGenerateDataKeyEvent.output.value.KeyId.Some?
&& kmsGenerateDataKeyEvent.output.value.KeyId.value == kmsKeyArn
&& kmsGenerateDataKeyEvent.output.value.Plaintext.value == output.value
{
var kmsKeyArn := GetKeyId(kmsConfiguration);
var generateDataKeyInput := KMS.GenerateDataKeyRequest(
KeyId := kmsKeyArn,
NumberOfBytes := Some(32),
EncryptionContext := Some(branchKeyContext),
GrantTokens := Some(keyManagerAndStorage.keyManagerStrat.kmsSimple.grantTokens)
);

var generateDataKeyResponse? := keyManagerAndStorage.keyManagerStrat.kmsSimple.kmsClient.GenerateDataKey(
generateDataKeyInput
);

var generateDataKeyResponse :- generateDataKeyResponse?
.MapFailure(e => Types.ComAmazonawsKms(ComAmazonawsKms := e));

:- Need(
&& generateDataKeyResponse.Plaintext.Some?
&& |generateDataKeyResponse.Plaintext.value| == 32,
Types.KeyManagementException(
message := "Invalid response from AWS KMS GenerateDataKey: KMS response's Plaintext is invalid."
)
);
:- Need(
&& generateDataKeyResponse.KeyId.Some?
&& generateDataKeyResponse.KeyId.value == kmsKeyArn,
Types.KeyManagementException(
message := "Invalid response from AWS KMS GenerateDataKey: KMS Key ID of response did not match request."
)
);

output := Success(generateDataKeyResponse.Plaintext.value);
}

method GenerateKey(
encryptionContext: Structure.BranchKeyContext,
kmsConfiguration: Types.KMSConfiguration,
Expand Down Expand Up @@ -341,6 +408,7 @@ module {:options "/functionSyntax:4" } KMSKeystoreOperations {
res.Success?
==>
&& |kmsClient.History.Encrypt| == |old(kmsClient.History.Encrypt)| + 1
&& kmsClient.History.GenerateDataKey == old(kmsClient.History.GenerateDataKey)
Copy link
Contributor

Choose a reason for hiding this comment

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

I think I know this problem; but doing this should not be required.
I am going to see if I can fix this.

&& var kmsKeyArn := GetKeyId(kmsConfiguration);

&& var encryptInput := Seq.Last(kmsClient.History.Encrypt).input;
Expand Down Expand Up @@ -940,6 +1008,7 @@ module {:options "/functionSyntax:4" } KMSKeystoreOperations {
&& output.Success? ==>
&& var kms := cryptoAndKms.kms.kmsSimple.kmsClient;
&& |kms.History.Encrypt| == |old(kms.History.Encrypt)| + 1
&& kms.History.GenerateDataKey == old(kms.History.GenerateDataKey)
&& var kmsEvent := Seq.Last(kms.History.Encrypt);
&& kmsEvent.output.Success?
&& var kmsInput := kmsEvent.input;
Expand Down
Loading