Proposed "Resilient" CryptoMaterialsManager

CurrentlyAvailabeMasterKeySet.java

class CurrentlyAvailableMasterKeySet {

    private final Map<String, String> context;
    private final CryptoAlgorithm algo;
    private final List<MasterKey> remainingMasterKeys;
    private final int quorum;

    CurrentlyAvailableMasterKeySet(Map<String, String> context, CryptoAlgorithm algo, List<MasterKey> masterKeys, int quorum) {
        this.context = context;
        this.algo = algo;
        this.remainingMasterKeys = new LinkedList<>(masterKeys);
        this.quorum = quorum;
    }

    DataKey generateDataKey() {
        AWSKMSException firstFailureReason;
        try {
            return tryToGenerateDataKey();
        } catch (AWSKMSException exception) {
            firstFailureReason = exception;
        }
        while (remainingMasterKeys.size() >= quorum) {
            try {
                return tryToGenerateDataKey();
            } catch (AWSKMSException ignored) {
            }
        }
        throw firstFailureReason;
    }

    private DataKey tryToGenerateDataKey() {
        return remainingMasterKeys.remove(0).generateDataKey(algo, context);
    }

    List<KeyBlob> encryptDataKeys(DataKey dataKey) {
        List<KeyBlob> keyBlobs = new LinkedList<>();
        keyBlobs.add(new KeyBlob(dataKey));

        AWSKMSException firstFailureReason = null;
        while (keyBlobs.size() < quorum && remainingMasterKeys.size() > 0) {
            try {
                keyBlobs.add(tryToEncryptDataKey(dataKey));
            } catch (AWSKMSException e) {
                firstFailureReason = null;
            }
        }
        if (keyBlobs.size() == quorum) {
            return keyBlobs;
        } else if (firstFailureReason != null) {
            throw firstFailureReason;
        } else {
            throw new IllegalStateException("Insufficient keys to reach quorum");
        }
    }

    private KeyBlob tryToEncryptDataKey(DataKey dataKey) {
        return new KeyBlob(remainingMasterKeys.remove(0).encryptDataKey(algo, context, dataKey));
    }
}

ResilientCryptoMaterialsManager.java

public class ResilientCryptoMaterialsManager implements CryptoMaterialsManager {

    private final MasterKeyProvider<?> mkp;
    private final int quorum;

    public ResilientCryptoMaterialsManager(MasterKeyProvider<?> mkp, int quorum) {
        this.mkp = mkp;
        this.quorum = quorum;
    }

    @Override
    public EncryptionMaterials getMaterialsForEncrypt(EncryptionMaterialsRequest request) {
        Map<String, String> context = request.getContext();

        CryptoAlgorithm algo = request.getRequestedAlgorithm();
        if (algo == null) {
            algo = AwsCrypto.getDefaultCryptoAlgorithm();
        }

        KeyPair trailingKeys = null;
        if (algo.getTrailingSignatureLength() > 0) {
            try {
                trailingKeys = generateTrailingSigKeyPair(algo);
                if (context.containsKey(Constants.EC_PUBLIC_KEY_FIELD)) {
                    throw new IllegalArgumentException("EncryptionContext contains reserved field "
                            + Constants.EC_PUBLIC_KEY_FIELD);
                }
                // make mutable
                context = new HashMap<>(context);
                context.put(Constants.EC_PUBLIC_KEY_FIELD, serializeTrailingKeyForEc(algo, trailingKeys));
            } catch (final GeneralSecurityException ex) {
                throw new AwsCryptoException(ex);
            }
        }

        final MasterKeyRequest.Builder mkRequestBuilder = MasterKeyRequest.newBuilder();
        mkRequestBuilder.setEncryptionContext(context);

        mkRequestBuilder.setStreaming(request.getPlaintextSize() == -1);
        if (request.getPlaintext() != null) {
            mkRequestBuilder.setPlaintext(request.getPlaintext());
        } else {
            mkRequestBuilder.setSize(request.getPlaintextSize());
        }

        @SuppressWarnings("unchecked")
        final List<MasterKey> mks
                = (List<MasterKey>)assertNonNull(mkp, "provider")
                .getMasterKeysForEncryption(mkRequestBuilder.build());

        if (mks.isEmpty()) {
            throw new IllegalArgumentException("No master keys provided");
        }

        CurrentlyAvailableMasterKeySet keySet = new CurrentlyAvailableMasterKeySet(context, algo, mks, quorum);

        DataKey<?> dataKey = keySet.generateDataKey();
        List<KeyBlob> keyBlobs = keySet.encryptDataKeys(dataKey);

        //noinspection unchecked
        return EncryptionMaterials.newBuilder()
                .setAlgorithm(algo)
                .setCleartextDataKey(dataKey.getKey())
                .setEncryptedDataKeys(keyBlobs)
                .setEncryptionContext(context)
                .setTrailingSignatureKey(trailingKeys == null ? null : trailingKeys.getPrivate())
                .setMasterKeys(mks)
                .build();
    }

    @Override
    public DecryptionMaterials decryptMaterials(DecryptionMaterialsRequest request) {
        throw new UnsupportedOperationException("ResilientCryptoMaterialsManager only relevant for encryption");
    }

    private static String serializeTrailingKeyForEc(CryptoAlgorithm algo, KeyPair trailingKeys) {
        return TrailingSignatureAlgorithm.forCryptoAlgorithm(algo).serializePublicKey(trailingKeys.getPublic());
    }

    private static KeyPair generateTrailingSigKeyPair(CryptoAlgorithm algo) throws GeneralSecurityException {
        return TrailingSignatureAlgorithm.forCryptoAlgorithm(algo).generateKey();
    }
}