Examples with KMSConfig com.github.ambry.config.KMSConfig used on opensource projects

Example 1 with KMSConfig

use of com.github.ambry.config.KMSConfig in project ambry by linkedin.

the class ChunkFillTest method fillChunksAndAssertSuccess.

/**
 * Create a {@link PutOperation} and pass in a channel with the blobSize set by the caller; and test the chunk
 * filling flow for puts.
 * Note that this test is for the chunk filling flow, not for the ChunkFiller thread (which never gets exercised,
 * as we do not even instantiate the {@link PutManager})
 */
private void fillChunksAndAssertSuccess() throws Exception {
    VerifiableProperties vProps = getNonBlockingRouterProperties();
    MockClusterMap mockClusterMap = new MockClusterMap();
    RouterConfig routerConfig = new RouterConfig(vProps);
    routerMetrics = new NonBlockingRouterMetrics(mockClusterMap, routerConfig);
    ResponseHandler responseHandler = new ResponseHandler(mockClusterMap);
    short accountId = Utils.getRandomShort(random);
    short containerId = Utils.getRandomShort(random);
    BlobProperties putBlobProperties = new BlobProperties(blobSize, "serviceId", "memberId", "contentType", false, Utils.Infinite_Time, accountId, containerId, testEncryption, null, null, null);
    Random random = new Random();
    byte[] putUserMetadata = new byte[10];
    random.nextBytes(putUserMetadata);
    putContent = new byte[blobSize];
    random.nextBytes(putContent);
    final ReadableStreamChannel putChannel = new ByteBufferReadableStreamChannel(ByteBuffer.wrap(putContent));
    FutureResult<String> futureResult = new FutureResult<String>();
    MockTime time = new MockTime();
    MockNetworkClientFactory networkClientFactory = new MockNetworkClientFactory(vProps, null, 0, 0, 0, null, time);
    if (testEncryption) {
        kms = new MockKeyManagementService(new KMSConfig(vProps), TestUtils.getRandomKey(SingleKeyManagementServiceTest.DEFAULT_KEY_SIZE_CHARS));
        cryptoService = new MockCryptoService(new CryptoServiceConfig(vProps));
        cryptoJobHandler = new CryptoJobHandler(CryptoJobHandlerTest.DEFAULT_THREAD_COUNT);
    }
    MockRouterCallback routerCallback = new MockRouterCallback(networkClientFactory.getNetworkClient(), Collections.EMPTY_LIST);
    PutOperation op = PutOperation.forUpload(routerConfig, routerMetrics, mockClusterMap, new LoggingNotificationSystem(), new InMemAccountService(true, false), putUserMetadata, putChannel, PutBlobOptions.DEFAULT, futureResult, null, routerCallback, null, kms, cryptoService, cryptoJobHandler, time, putBlobProperties, MockClusterMap.DEFAULT_PARTITION_CLASS, quotaChargeCallback);
    op.startOperation();
    numChunks = RouterUtils.getNumChunksForBlobAndChunkSize(blobSize, chunkSize);
    compositeBuffers = new ByteBuf[numChunks];
    compositeEncryptionKeys = new ByteBuffer[numChunks];
    compositeBlobIds = new BlobId[numChunks];
    final AtomicReference<Exception> operationException = new AtomicReference<Exception>(null);
    int chunksLeftToBeFilled = numChunks;
    do {
        if (testEncryption) {
            int chunksPerBatch = Math.min(routerConfig.routerMaxInMemPutChunks, chunksLeftToBeFilled);
            CountDownLatch onPollLatch = new CountDownLatch(chunksPerBatch);
            routerCallback.setOnPollLatch(onPollLatch);
            op.fillChunks();
            Assert.assertTrue("Latch should have been zeroed out", onPollLatch.await(1000, TimeUnit.MILLISECONDS));
            chunksLeftToBeFilled -= chunksPerBatch;
        } else {
            op.fillChunks();
        }
        // since the channel is ByteBuffer based.
        for (PutOperation.PutChunk putChunk : op.putChunks) {
            if (putChunk.isFree()) {
                continue;
            }
            Assert.assertEquals("Chunk should be ready.", PutOperation.ChunkState.Ready, putChunk.getState());
            ByteBuf buf = putChunk.buf.retainedDuplicate();
            totalSizeWritten += buf.readableBytes();
            compositeBuffers[putChunk.getChunkIndex()] = buf;
            if (testEncryption) {
                compositeEncryptionKeys[putChunk.getChunkIndex()] = putChunk.encryptedPerBlobKey.duplicate();
                compositeBlobIds[putChunk.getChunkIndex()] = putChunk.chunkBlobId;
            }
            putChunk.clear();
        }
    } while (!op.isChunkFillingDone());
    if (!testEncryption) {
        Assert.assertEquals("total size written out should match the blob size", blobSize, totalSizeWritten);
    }
    // for encrypted path, size will be implicitly tested via assertDataIdentity
    Exception exception = operationException.get();
    if (exception != null) {
        throw exception;
    }
    assertDataIdentity(mockClusterMap);
}

Example 2 with KMSConfig

use of com.github.ambry.config.KMSConfig in project ambry by linkedin.

the class NonBlockingRouterTest method testResponseHandling.

/**
 * Response handling related tests for all operation managers.
 */
@Test
public void testResponseHandling() throws Exception {
    try {
        Properties props = getNonBlockingRouterProperties("DC1");
        VerifiableProperties verifiableProperties = new VerifiableProperties((props));
        setOperationParams();
        final List<ReplicaId> failedReplicaIds = new ArrayList<>();
        final AtomicInteger successfulResponseCount = new AtomicInteger(0);
        final AtomicBoolean invalidResponse = new AtomicBoolean(false);
        ResponseHandler mockResponseHandler = new ResponseHandler(mockClusterMap) {

            @Override
            public void onEvent(ReplicaId replicaId, Object e) {
                if (e instanceof ServerErrorCode) {
                    if (e == ServerErrorCode.No_Error) {
                        successfulResponseCount.incrementAndGet();
                    } else {
                        invalidResponse.set(true);
                    }
                } else {
                    failedReplicaIds.add(replicaId);
                }
            }
        };
        // Instantiate a router just to put a blob successfully.
        MockServerLayout mockServerLayout = new MockServerLayout(mockClusterMap);
        setRouter(props, mockServerLayout, new LoggingNotificationSystem());
        setOperationParams();
        // More extensive test for puts present elsewhere - these statements are here just to exercise the flow within the
        // NonBlockingRouter class, and to ensure that operations submitted to a router eventually completes.
        String blobIdStr = router.putBlob(putBlobProperties, putUserMetadata, putChannel, new PutBlobOptionsBuilder().build()).get();
        BlobId blobId = RouterUtils.getBlobIdFromString(blobIdStr, mockClusterMap);
        router.close();
        for (MockServer mockServer : mockServerLayout.getMockServers()) {
            mockServer.setServerErrorForAllRequests(ServerErrorCode.No_Error);
        }
        SocketNetworkClient networkClient = new MockNetworkClientFactory(verifiableProperties, mockSelectorState, MAX_PORTS_PLAIN_TEXT, MAX_PORTS_SSL, CHECKOUT_TIMEOUT_MS, mockServerLayout, mockTime).getNetworkClient();
        cryptoJobHandler = new CryptoJobHandler(CryptoJobHandlerTest.DEFAULT_THREAD_COUNT);
        KeyManagementService localKMS = new MockKeyManagementService(new KMSConfig(verifiableProperties), singleKeyForKMS);
        putManager = new PutManager(mockClusterMap, mockResponseHandler, new LoggingNotificationSystem(), new RouterConfig(verifiableProperties), new NonBlockingRouterMetrics(mockClusterMap, null), new RouterCallback(networkClient, new ArrayList<>()), "0", localKMS, cryptoService, cryptoJobHandler, accountService, mockTime, MockClusterMap.DEFAULT_PARTITION_CLASS);
        OperationHelper opHelper = new OperationHelper(OperationType.PUT);
        testFailureDetectorNotification(opHelper, networkClient, failedReplicaIds, null, successfulResponseCount, invalidResponse, -1);
        // Test that if a failed response comes before the operation is completed, failure detector is notified.
        testFailureDetectorNotification(opHelper, networkClient, failedReplicaIds, null, successfulResponseCount, invalidResponse, 0);
        // Test that if a failed response comes after the operation is completed, failure detector is notified.
        testFailureDetectorNotification(opHelper, networkClient, failedReplicaIds, null, successfulResponseCount, invalidResponse, PUT_REQUEST_PARALLELISM - 1);
        testNoResponseNoNotification(opHelper, failedReplicaIds, null, successfulResponseCount, invalidResponse);
        testResponseDeserializationError(opHelper, networkClient, null);
        opHelper = new OperationHelper(OperationType.GET);
        getManager = new GetManager(mockClusterMap, mockResponseHandler, new RouterConfig(verifiableProperties), new NonBlockingRouterMetrics(mockClusterMap, null), new RouterCallback(networkClient, new ArrayList<BackgroundDeleteRequest>()), localKMS, cryptoService, cryptoJobHandler, mockTime);
        testFailureDetectorNotification(opHelper, networkClient, failedReplicaIds, blobId, successfulResponseCount, invalidResponse, -1);
        // Test that if a failed response comes before the operation is completed, failure detector is notified.
        testFailureDetectorNotification(opHelper, networkClient, failedReplicaIds, blobId, successfulResponseCount, invalidResponse, 0);
        // Test that if a failed response comes after the operation is completed, failure detector is notified.
        testFailureDetectorNotification(opHelper, networkClient, failedReplicaIds, blobId, successfulResponseCount, invalidResponse, GET_REQUEST_PARALLELISM - 1);
        testNoResponseNoNotification(opHelper, failedReplicaIds, blobId, successfulResponseCount, invalidResponse);
        testResponseDeserializationError(opHelper, networkClient, blobId);
        opHelper = new OperationHelper(OperationType.DELETE);
        deleteManager = new DeleteManager(mockClusterMap, mockResponseHandler, accountService, new LoggingNotificationSystem(), new RouterConfig(verifiableProperties), new NonBlockingRouterMetrics(mockClusterMap, null), new RouterCallback(null, new ArrayList<BackgroundDeleteRequest>()), mockTime);
        testFailureDetectorNotification(opHelper, networkClient, failedReplicaIds, blobId, successfulResponseCount, invalidResponse, -1);
        // Test that if a failed response comes before the operation is completed, failure detector is notified.
        testFailureDetectorNotification(opHelper, networkClient, failedReplicaIds, blobId, successfulResponseCount, invalidResponse, 0);
        // Test that if a failed response comes after the operation is completed, failure detector is notified.
        testFailureDetectorNotification(opHelper, networkClient, failedReplicaIds, blobId, successfulResponseCount, invalidResponse, DELETE_REQUEST_PARALLELISM - 1);
        testNoResponseNoNotification(opHelper, failedReplicaIds, blobId, successfulResponseCount, invalidResponse);
        testResponseDeserializationError(opHelper, networkClient, blobId);
    } finally {
        if (putManager != null) {
            putManager.close();
        }
        if (getManager != null) {
            getManager.close();
        }
        if (deleteManager != null) {
            deleteManager.close();
        }
    }
}

Example 3 with KMSConfig

use of com.github.ambry.config.KMSConfig in project ambry by linkedin.

the class SingleKeyManagementServiceTest method testSingleKMS.

/**
 * Test {@link SingleKeyManagementService#getKey(RestRequest, short, short)}
 */
@Test
public void testSingleKMS() throws Exception {
    int[] keySizes = { 16, 32, 64, 128 };
    for (int keySize : keySizes) {
        String key = TestUtils.getRandomKey(keySize);
        Properties props = getKMSProperties(key, DEFAULT_RANDOM_KEY_SIZE_BITS);
        VerifiableProperties verifiableProperties = new VerifiableProperties((props));
        KMSConfig config = new KMSConfig(verifiableProperties);
        SecretKeySpec secretKeySpec = new SecretKeySpec(Hex.decode(key), config.kmsKeyGenAlgo);
        KeyManagementService<SecretKeySpec> kms = new SingleKeyManagementServiceFactory(verifiableProperties, CLUSTER_NAME, REGISTRY).getKeyManagementService();
        SecretKeySpec keyFromKMS = kms.getKey(null, Utils.getRandomShort(TestUtils.RANDOM), Utils.getRandomShort(TestUtils.RANDOM));
        Assert.assertEquals("Secret key mismatch ", secretKeySpec, keyFromKMS);
    }
}

Example 4 with KMSConfig

use of com.github.ambry.config.KMSConfig in project ambry by linkedin.

the class GetManagerTest method resetEncryptionObjects.

/**
 * Resets objects related to encryption testing
 * @throws GeneralSecurityException
 */
private void resetEncryptionObjects() throws GeneralSecurityException {
    if (testEncryption) {
        VerifiableProperties vProps = new VerifiableProperties(new Properties());
        kms = new SingleKeyManagementService(new KMSConfig(vProps), TestUtils.getRandomKey(SingleKeyManagementServiceTest.DEFAULT_KEY_SIZE_CHARS));
        cryptoService = new GCMCryptoService(new CryptoServiceConfig(vProps));
        cryptoJobHandler = new CryptoJobHandler(CryptoJobHandlerTest.DEFAULT_THREAD_COUNT);
    }
}

Example 5 with KMSConfig

use of com.github.ambry.config.KMSConfig in project ambry by linkedin.

the class CryptoJobHandlerTest method testEncryptionFailure.

/**
 * Tests {@link CryptoJobHandler} for failures during encryption
 * @throws InterruptedException
 * @throws GeneralSecurityException
 */
@Test
public void testEncryptionFailure() throws InterruptedException, GeneralSecurityException {
    cryptoJobHandler.close();
    MockCryptoService mockCryptoService = new MockCryptoService(new CryptoServiceConfig(verifiableProperties));
    mockCryptoService.exceptionOnEncryption.set(new GeneralSecurityException("Exception to test", new IllegalStateException()));
    cryptoJobHandler = new CryptoJobHandler(DEFAULT_THREAD_COUNT);
    SecretKeySpec perBlobSecretKey = kms.getRandomKey();
    testFailureOnEncryption(perBlobSecretKey, mockCryptoService, kms);
    mockCryptoService.clearStates();
    cryptoJobHandler.close();
    MockKeyManagementService mockKms = new MockKeyManagementService(new KMSConfig(verifiableProperties), defaultKey);
    mockKms.exceptionToThrow.set(new GeneralSecurityException("Exception to test", new IllegalStateException()));
    cryptoJobHandler = new CryptoJobHandler(DEFAULT_THREAD_COUNT);
    testFailureOnEncryption(perBlobSecretKey, cryptoService, mockKms);
}