Examples with RMSecretManagerService org.apache.hadoop.yarn.server.resourcemanager.RMSecretManagerService used on opensource projects

Example 1 with RMSecretManagerService

use of org.apache.hadoop.yarn.server.resourcemanager.RMSecretManagerService in project hadoop by apache.

the class TestAMRMTokens method testAMRMMasterKeysUpdate.

@Test(timeout = 20000)
public void testAMRMMasterKeysUpdate() throws Exception {
    final AtomicReference<AMRMTokenSecretManager> spySecretMgrRef = new AtomicReference<AMRMTokenSecretManager>();
    MockRM rm = new MockRM(conf) {

        @Override
        protected void doSecureLogin() throws IOException {
        // Skip the login.
        }

        @Override
        protected RMSecretManagerService createRMSecretManagerService() {
            return new RMSecretManagerService(conf, rmContext) {

                @Override
                protected AMRMTokenSecretManager createAMRMTokenSecretManager(Configuration conf, RMContext rmContext) {
                    AMRMTokenSecretManager spySecretMgr = spy(super.createAMRMTokenSecretManager(conf, rmContext));
                    spySecretMgrRef.set(spySecretMgr);
                    return spySecretMgr;
                }
            };
        }
    };
    rm.start();
    MockNM nm = rm.registerNode("127.0.0.1:1234", 8000);
    RMApp app = rm.submitApp(200);
    MockAM am = MockRM.launchAndRegisterAM(app, rm, nm);
    AMRMTokenSecretManager spySecretMgr = spySecretMgrRef.get();
    // Do allocate. Should not update AMRMToken
    AllocateResponse response = am.allocate(Records.newRecord(AllocateRequest.class));
    Assert.assertNull(response.getAMRMToken());
    Token<AMRMTokenIdentifier> oldToken = rm.getRMContext().getRMApps().get(app.getApplicationId()).getRMAppAttempt(am.getApplicationAttemptId()).getAMRMToken();
    // roll over the master key
    // Do allocate again. the AM should get the latest AMRMToken
    rm.getRMContext().getAMRMTokenSecretManager().rollMasterKey();
    response = am.allocate(Records.newRecord(AllocateRequest.class));
    Assert.assertNotNull(response.getAMRMToken());
    Token<AMRMTokenIdentifier> amrmToken = ConverterUtils.convertFromYarn(response.getAMRMToken(), new Text(response.getAMRMToken().getService()));
    Assert.assertEquals(amrmToken.decodeIdentifier().getKeyId(), rm.getRMContext().getAMRMTokenSecretManager().getMasterKey().getMasterKey().getKeyId());
    // Do allocate again with the same old token and verify the RM sends
    // back the last generated token instead of generating it again.
    reset(spySecretMgr);
    UserGroupInformation ugi = UserGroupInformation.createUserForTesting(am.getApplicationAttemptId().toString(), new String[0]);
    ugi.addTokenIdentifier(oldToken.decodeIdentifier());
    response = am.doAllocateAs(ugi, Records.newRecord(AllocateRequest.class));
    Assert.assertNotNull(response.getAMRMToken());
    verify(spySecretMgr, never()).createAndGetAMRMToken(isA(ApplicationAttemptId.class));
    // Do allocate again with the updated token and verify we do not
    // receive a new token to use.
    response = am.allocate(Records.newRecord(AllocateRequest.class));
    Assert.assertNull(response.getAMRMToken());
    // Activate the next master key. Since there is new master key generated
    // in AMRMTokenSecretManager. The AMRMToken will not get updated for AM
    rm.getRMContext().getAMRMTokenSecretManager().activateNextMasterKey();
    response = am.allocate(Records.newRecord(AllocateRequest.class));
    Assert.assertNull(response.getAMRMToken());
    rm.stop();
}