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

Example 1 with AMRMTokenSecretManager

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

the class TestAMRMClient method testAMRMClientOnAMRMTokenRollOver.

@Test(timeout = 60000)
public void testAMRMClientOnAMRMTokenRollOver() throws YarnException, IOException {
    AMRMClient<ContainerRequest> amClient = null;
    try {
        AMRMTokenSecretManager amrmTokenSecretManager = yarnCluster.getResourceManager().getRMContext().getAMRMTokenSecretManager();
        // start am rm client
        amClient = AMRMClient.<ContainerRequest>createAMRMClient();
        amClient.init(conf);
        amClient.start();
        Long startTime = System.currentTimeMillis();
        amClient.registerApplicationMaster("Host", 10000, "");
        org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> amrmToken_1 = getAMRMToken();
        Assert.assertNotNull(amrmToken_1);
        Assert.assertEquals(amrmToken_1.decodeIdentifier().getKeyId(), amrmTokenSecretManager.getMasterKey().getMasterKey().getKeyId());
        // At mean time, the old AMRMToken should continue to work
        while (System.currentTimeMillis() - startTime < rolling_interval_sec * 1000) {
            amClient.allocate(0.1f);
            sleep(1000);
        }
        amClient.allocate(0.1f);
        org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> amrmToken_2 = getAMRMToken();
        Assert.assertNotNull(amrmToken_2);
        Assert.assertEquals(amrmToken_2.decodeIdentifier().getKeyId(), amrmTokenSecretManager.getMasterKey().getMasterKey().getKeyId());
        Assert.assertNotEquals(amrmToken_1, amrmToken_2);
        // can do the allocate call with latest AMRMToken
        AllocateResponse response = amClient.allocate(0.1f);
        // Verify latest AMRMToken can be used to send allocation request.
        UserGroupInformation testUser1 = UserGroupInformation.createRemoteUser("testUser1");
        AMRMTokenIdentifierForTest newVersionTokenIdentifier = new AMRMTokenIdentifierForTest(amrmToken_2.decodeIdentifier(), "message");
        Assert.assertEquals("Message is changed after set to newVersionTokenIdentifier", "message", newVersionTokenIdentifier.getMessage());
        org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> newVersionToken = new org.apache.hadoop.security.token.Token<AMRMTokenIdentifier>(newVersionTokenIdentifier.getBytes(), amrmTokenSecretManager.retrievePassword(newVersionTokenIdentifier), newVersionTokenIdentifier.getKind(), new Text());
        SecurityUtil.setTokenService(newVersionToken, yarnCluster.getResourceManager().getApplicationMasterService().getBindAddress());
        testUser1.addToken(newVersionToken);
        AllocateRequest request = Records.newRecord(AllocateRequest.class);
        request.setResponseId(response.getResponseId());
        testUser1.doAs(new PrivilegedAction<ApplicationMasterProtocol>() {

            @Override
            public ApplicationMasterProtocol run() {
                return (ApplicationMasterProtocol) YarnRPC.create(conf).getProxy(ApplicationMasterProtocol.class, yarnCluster.getResourceManager().getApplicationMasterService().getBindAddress(), conf);
            }
        }).allocate(request);
        // and can not use this rolled-over token to make a allocate all.
        while (true) {
            if (amrmToken_2.decodeIdentifier().getKeyId() != amrmTokenSecretManager.getCurrnetMasterKeyData().getMasterKey().getKeyId()) {
                if (amrmTokenSecretManager.getNextMasterKeyData() == null) {
                    break;
                } else if (amrmToken_2.decodeIdentifier().getKeyId() != amrmTokenSecretManager.getNextMasterKeyData().getMasterKey().getKeyId()) {
                    break;
                }
            }
            amClient.allocate(0.1f);
            sleep(1000);
        }
        try {
            UserGroupInformation testUser2 = UserGroupInformation.createRemoteUser("testUser2");
            SecurityUtil.setTokenService(amrmToken_2, yarnCluster.getResourceManager().getApplicationMasterService().getBindAddress());
            testUser2.addToken(amrmToken_2);
            testUser2.doAs(new PrivilegedAction<ApplicationMasterProtocol>() {

                @Override
                public ApplicationMasterProtocol run() {
                    return (ApplicationMasterProtocol) YarnRPC.create(conf).getProxy(ApplicationMasterProtocol.class, yarnCluster.getResourceManager().getApplicationMasterService().getBindAddress(), conf);
                }
            }).allocate(Records.newRecord(AllocateRequest.class));
            Assert.fail("The old Token should not work");
        } catch (Exception ex) {
            Assert.assertTrue(ex instanceof InvalidToken);
            Assert.assertTrue(ex.getMessage().contains("Invalid AMRMToken from " + amrmToken_2.decodeIdentifier().getApplicationAttemptId()));
        }
        amClient.unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, null, null);
    } finally {
        if (amClient != null && amClient.getServiceState() == STATE.STARTED) {
            amClient.stop();
        }
    }
}

Example 2 with AMRMTokenSecretManager

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

the class TestAMRMClientOnRMRestart method testAMRMClientOnAMRMTokenRollOverOnRMRestart.

// Test verify for AM issued with rolled-over AMRMToken
// is still able to communicate with restarted RM.
@Test(timeout = 30000)
public void testAMRMClientOnAMRMTokenRollOverOnRMRestart() throws Exception {
    conf.setLong(YarnConfiguration.RM_AMRM_TOKEN_MASTER_KEY_ROLLING_INTERVAL_SECS, rolling_interval_sec);
    conf.setLong(YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS, am_expire_ms);
    MemoryRMStateStore memStore = new MemoryRMStateStore();
    memStore.init(conf);
    // start first RM
    MyResourceManager2 rm1 = new MyResourceManager2(conf, memStore);
    rm1.start();
    DrainDispatcher dispatcher = (DrainDispatcher) rm1.getRMContext().getDispatcher();
    Long startTime = System.currentTimeMillis();
    // Submit the application
    RMApp app = rm1.submitApp(1024);
    dispatcher.await();
    MockNM nm1 = new MockNM("h1:1234", 15120, rm1.getResourceTrackerService());
    nm1.registerNode();
    // Node heartbeat
    nm1.nodeHeartbeat(true);
    dispatcher.await();
    ApplicationAttemptId appAttemptId = app.getCurrentAppAttempt().getAppAttemptId();
    rm1.sendAMLaunched(appAttemptId);
    dispatcher.await();
    AMRMTokenSecretManager amrmTokenSecretManagerForRM1 = rm1.getRMContext().getAMRMTokenSecretManager();
    org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> token = amrmTokenSecretManagerForRM1.createAndGetAMRMToken(appAttemptId);
    UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
    ugi.addTokenIdentifier(token.decodeIdentifier());
    AMRMClient<ContainerRequest> amClient = new MyAMRMClientImpl(rm1);
    amClient.init(conf);
    amClient.start();
    amClient.registerApplicationMaster("h1", 10000, "");
    amClient.allocate(0.1f);
    // At mean time, the old AMRMToken should continue to work
    while (System.currentTimeMillis() - startTime < rolling_interval_sec * 1000) {
        amClient.allocate(0.1f);
        try {
            Thread.sleep(1000);
        } catch (InterruptedException e) {
        // DO NOTHING
        }
    }
    Assert.assertTrue(amrmTokenSecretManagerForRM1.getMasterKey().getMasterKey().getKeyId() != token.decodeIdentifier().getKeyId());
    amClient.allocate(0.1f);
    // active the nextMasterKey, and replace the currentMasterKey
    org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> newToken = amrmTokenSecretManagerForRM1.createAndGetAMRMToken(appAttemptId);
    int waitCount = 0;
    while (waitCount++ <= 50) {
        if (amrmTokenSecretManagerForRM1.getCurrnetMasterKeyData().getMasterKey().getKeyId() != token.decodeIdentifier().getKeyId()) {
            break;
        }
        try {
            amClient.allocate(0.1f);
        } catch (Exception ex) {
            break;
        }
        Thread.sleep(500);
    }
    Assert.assertTrue(amrmTokenSecretManagerForRM1.getNextMasterKeyData() == null);
    Assert.assertTrue(amrmTokenSecretManagerForRM1.getCurrnetMasterKeyData().getMasterKey().getKeyId() == newToken.decodeIdentifier().getKeyId());
    // start 2nd RM
    conf.set(YarnConfiguration.RM_SCHEDULER_ADDRESS, "0.0.0.0:" + ServerSocketUtil.getPort(45020, 10));
    final MyResourceManager2 rm2 = new MyResourceManager2(conf, memStore);
    rm2.start();
    nm1.setResourceTrackerService(rm2.getResourceTrackerService());
    ((MyAMRMClientImpl) amClient).updateRMProxy(rm2);
    dispatcher = (DrainDispatcher) rm2.getRMContext().getDispatcher();
    AMRMTokenSecretManager amrmTokenSecretManagerForRM2 = rm2.getRMContext().getAMRMTokenSecretManager();
    Assert.assertTrue(amrmTokenSecretManagerForRM2.getCurrnetMasterKeyData().getMasterKey().getKeyId() == newToken.decodeIdentifier().getKeyId());
    Assert.assertTrue(amrmTokenSecretManagerForRM2.getNextMasterKeyData() == null);
    try {
        UserGroupInformation testUser = UserGroupInformation.createRemoteUser("testUser");
        SecurityUtil.setTokenService(token, rm2.getApplicationMasterService().getBindAddress());
        testUser.addToken(token);
        testUser.doAs(new PrivilegedAction<ApplicationMasterProtocol>() {

            @Override
            public ApplicationMasterProtocol run() {
                return (ApplicationMasterProtocol) YarnRPC.create(conf).getProxy(ApplicationMasterProtocol.class, rm2.getApplicationMasterService().getBindAddress(), conf);
            }
        }).allocate(Records.newRecord(AllocateRequest.class));
        Assert.fail("The old Token should not work");
    } catch (Exception ex) {
        Assert.assertTrue(ex instanceof InvalidToken);
        Assert.assertTrue(ex.getMessage().contains("Invalid AMRMToken from " + token.decodeIdentifier().getApplicationAttemptId()));
    }
    // make sure the recovered AMRMToken works for new RM
    amClient.allocate(0.1f);
    amClient.unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, null, null);
    amClient.stop();
    rm1.stop();
    rm2.stop();
}

Example 3 with AMRMTokenSecretManager

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

the class TestUtils method getMockRMContext.

/**
   * Get a mock {@link RMContext} for use in test cases.
   * @return a mock {@link RMContext} for use in test cases
   */
@SuppressWarnings({ "rawtypes", "unchecked" })
public static RMContext getMockRMContext() {
    // Null dispatcher
    Dispatcher nullDispatcher = new Dispatcher() {

        private final EventHandler handler = new EventHandler() {

            @Override
            public void handle(Event event) {
            }
        };

        @Override
        public void register(Class<? extends Enum> eventType, EventHandler handler) {
        }

        @Override
        public EventHandler<Event> getEventHandler() {
            return handler;
        }
    };
    // No op
    ContainerAllocationExpirer cae = new ContainerAllocationExpirer(nullDispatcher);
    Configuration conf = new Configuration();
    RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
    RMContextImpl rmContext = new RMContextImpl(nullDispatcher, cae, null, null, null, new AMRMTokenSecretManager(conf, null), new RMContainerTokenSecretManager(conf), new NMTokenSecretManagerInRM(conf), new ClientToAMTokenSecretManagerInRM());
    RMNodeLabelsManager nlm = mock(RMNodeLabelsManager.class);
    when(nlm.getQueueResource(any(String.class), any(Set.class), any(Resource.class))).thenAnswer(new Answer<Resource>() {

        @Override
        public Resource answer(InvocationOnMock invocation) throws Throwable {
            Object[] args = invocation.getArguments();
            return (Resource) args[2];
        }
    });
    when(nlm.getResourceByLabel(any(String.class), any(Resource.class))).thenAnswer(new Answer<Resource>() {

        @Override
        public Resource answer(InvocationOnMock invocation) throws Throwable {
            Object[] args = invocation.getArguments();
            return (Resource) args[1];
        }
    });
    rmContext.setNodeLabelManager(nlm);
    rmContext.setSystemMetricsPublisher(mock(SystemMetricsPublisher.class));
    rmContext.setRMApplicationHistoryWriter(mock(RMApplicationHistoryWriter.class));
    ResourceScheduler mockScheduler = mock(ResourceScheduler.class);
    when(mockScheduler.getResourceCalculator()).thenReturn(new DefaultResourceCalculator());
    rmContext.setScheduler(mockScheduler);
    return rmContext;
}

Example 4 with AMRMTokenSecretManager

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

the class TestZKRMStateStorePerf method initStore.

private void initStore(String hostPort) {
    Optional<String> optHostPort = Optional.fromNullable(hostPort);
    RMContext rmContext = mock(RMContext.class);
    conf = new YarnConfiguration();
    conf.set(YarnConfiguration.RM_ZK_ADDRESS, optHostPort.or((curatorTestingServer == null) ? "" : curatorTestingServer.getConnectString()));
    conf.set(YarnConfiguration.ZK_RM_STATE_STORE_PARENT_PATH, workingZnode);
    store = new ZKRMStateStore();
    store.setResourceManager(new ResourceManager());
    store.init(conf);
    store.start();
    when(rmContext.getStateStore()).thenReturn(store);
    appTokenMgr = new AMRMTokenSecretManager(conf, rmContext);
    appTokenMgr.start();
    clientToAMTokenMgr = new ClientToAMTokenSecretManagerInRM();
}

Example 5 with AMRMTokenSecretManager

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

the class RMStateStoreTestBase method testAMRMTokenSecretManagerStateStore.

public void testAMRMTokenSecretManagerStateStore(RMStateStoreHelper stateStoreHelper) throws Exception {
    System.out.println("Start testing");
    RMStateStore store = stateStoreHelper.getRMStateStore();
    TestDispatcher dispatcher = new TestDispatcher();
    store.setRMDispatcher(dispatcher);
    RMContext rmContext = mock(RMContext.class);
    when(rmContext.getStateStore()).thenReturn(store);
    Configuration conf = new YarnConfiguration();
    AMRMTokenSecretManager appTokenMgr = new AMRMTokenSecretManager(conf, rmContext);
    //create and save the first masterkey
    MasterKeyData firstMasterKeyData = appTokenMgr.createNewMasterKey();
    AMRMTokenSecretManagerState state1 = AMRMTokenSecretManagerState.newInstance(firstMasterKeyData.getMasterKey(), null);
    rmContext.getStateStore().storeOrUpdateAMRMTokenSecretManager(state1, false);
    // load state
    store = stateStoreHelper.getRMStateStore();
    when(rmContext.getStateStore()).thenReturn(store);
    store.setRMDispatcher(dispatcher);
    RMState state = store.loadState();
    Assert.assertNotNull(state.getAMRMTokenSecretManagerState());
    Assert.assertEquals(firstMasterKeyData.getMasterKey(), state.getAMRMTokenSecretManagerState().getCurrentMasterKey());
    Assert.assertNull(state.getAMRMTokenSecretManagerState().getNextMasterKey());
    //create and save the second masterkey
    MasterKeyData secondMasterKeyData = appTokenMgr.createNewMasterKey();
    AMRMTokenSecretManagerState state2 = AMRMTokenSecretManagerState.newInstance(firstMasterKeyData.getMasterKey(), secondMasterKeyData.getMasterKey());
    rmContext.getStateStore().storeOrUpdateAMRMTokenSecretManager(state2, true);
    // load state
    store = stateStoreHelper.getRMStateStore();
    when(rmContext.getStateStore()).thenReturn(store);
    store.setRMDispatcher(dispatcher);
    RMState state_2 = store.loadState();
    Assert.assertNotNull(state_2.getAMRMTokenSecretManagerState());
    Assert.assertEquals(firstMasterKeyData.getMasterKey(), state_2.getAMRMTokenSecretManagerState().getCurrentMasterKey());
    Assert.assertEquals(secondMasterKeyData.getMasterKey(), state_2.getAMRMTokenSecretManagerState().getNextMasterKey());
    // re-create the masterKeyData based on the recovered masterkey
    // should have the same secretKey
    appTokenMgr.recover(state_2);
    Assert.assertEquals(appTokenMgr.getCurrnetMasterKeyData().getSecretKey(), firstMasterKeyData.getSecretKey());
    Assert.assertEquals(appTokenMgr.getNextMasterKeyData().getSecretKey(), secondMasterKeyData.getSecretKey());
    store.close();
}