Java Code Examples for org.apache.hadoop.security.Credentials#addToken()

public static ByteBuffer getTokens(UserGroupInformation ugi, Token<StramDelegationTokenIdentifier> delegationToken)
{
  try {
    Collection<Token<? extends TokenIdentifier>> tokens = ugi.getCredentials().getAllTokens();
    Credentials credentials = new Credentials();
    for (Token<? extends TokenIdentifier> token : tokens) {
      if (!token.getKind().equals(AMRMTokenIdentifier.KIND_NAME)) {
        credentials.addToken(token.getService(), token);
        LOG.debug("Passing container token {}", token);
      }
    }
    credentials.addToken(delegationToken.getService(), delegationToken);
    DataOutputBuffer dataOutput = new DataOutputBuffer();
    credentials.writeTokenStorageToStream(dataOutput);
    byte[] tokenBytes = dataOutput.getData();
    ByteBuffer cTokenBuf = ByteBuffer.wrap(tokenBytes);
    return cTokenBuf.duplicate();
  } catch (IOException e) {
    throw new RuntimeException("Error generating delegation token", e);
  }
}
 

@SuppressWarnings("deprecation")
@Test
public void testGetTokensForNamenodes() throws IOException,
    URISyntaxException {
  Path TEST_ROOT_DIR =
      new Path(System.getProperty("test.build.data", "test/build/data"));
  // ick, but need fq path minus file:/
  String binaryTokenFile =
      FileSystem.getLocal(conf)
        .makeQualified(new Path(TEST_ROOT_DIR, "tokenFile")).toUri()
        .getPath();

  MockFileSystem fs1 = createFileSystemForServiceName("service1");
  Credentials creds = new Credentials();
  Token<?> token1 = fs1.getDelegationToken(renewer);
  creds.addToken(token1.getService(), token1);
  // wait to set, else the obtain tokens call above will fail with FNF
  conf.set(MRJobConfig.MAPREDUCE_JOB_CREDENTIALS_BINARY, binaryTokenFile);
  creds.writeTokenStorageFile(new Path(binaryTokenFile), conf);
  TokenCache.obtainTokensForNamenodesInternal(fs1, creds, conf);
  String fs_addr = fs1.getCanonicalServiceName();
  Token<?> nnt = TokenCache.getDelegationToken(creds, fs_addr);
  assertNotNull("Token for nn is null", nnt);
}
 

@Test
public void testFsWithMyOwnExistsAndChildTokens() throws Exception {
  Credentials credentials = new Credentials();
  Text service1 = new Text("singleTokenFs1");
  Text service2 = new Text("singleTokenFs2");
  Text myService = new Text("multiTokenFs");
  Token<?> token = mock(Token.class);
  credentials.addToken(myService, token);

  MockFileSystem fs1 = createFileSystemForServiceName(service1);
  MockFileSystem fs2 = createFileSystemForServiceName(service2);
  MockFileSystem multiFs = createFileSystemForServiceName(myService, fs1, fs2);
  
  multiFs.addDelegationTokens(renewer, credentials);
  verifyTokenFetch(multiFs, false);  // we had added its token to credentials
  verifyTokenFetch(fs1, true);
  verifyTokenFetch(fs2, true);
  
  assertEquals(3, credentials.numberOfTokens());
  assertSame(token, credentials.getToken(myService));
  assertNotNull(credentials.getToken(service1));
  assertNotNull(credentials.getToken(service2));
}
 

@SuppressWarnings("deprecation")
@Test
public void testGetTokensForNamenodes() throws IOException,
    URISyntaxException {
  Path TEST_ROOT_DIR =
      new Path(System.getProperty("test.build.data", "test/build/data"));
  // ick, but need fq path minus file:/
  String binaryTokenFile =
      FileSystem.getLocal(conf)
        .makeQualified(new Path(TEST_ROOT_DIR, "tokenFile")).toUri()
        .getPath();

  MockFileSystem fs1 = createFileSystemForServiceName("service1");
  Credentials creds = new Credentials();
  Token<?> token1 = fs1.getDelegationToken(renewer);
  creds.addToken(token1.getService(), token1);
  // wait to set, else the obtain tokens call above will fail with FNF
  conf.set(MRJobConfig.MAPREDUCE_JOB_CREDENTIALS_BINARY, binaryTokenFile);
  creds.writeTokenStorageFile(new Path(binaryTokenFile), conf);
  TokenCache.obtainTokensForNamenodesInternal(fs1, creds, conf);
  String fs_addr = fs1.getCanonicalServiceName();
  Token<?> nnt = TokenCache.getDelegationToken(creds, fs_addr);
  assertNotNull("Token for nn is null", nnt);
}
 

@Override
public void handle(Channel channel, Token<DelegationTokenIdentifier> token,
    String serviceUrl) throws IOException {
  Assert.assertEquals(testToken, token);

  Credentials creds = new Credentials();
  creds.addToken(new Text(serviceUrl), token);
  DataOutputBuffer out = new DataOutputBuffer();
  creds.write(out);
  int fileLength = out.getData().length;
  ChannelBuffer cbuffer = ChannelBuffers.buffer(fileLength);
  cbuffer.writeBytes(out.getData());
  HttpResponse response = new DefaultHttpResponse(HTTP_1_1, OK);
  response.setHeader(HttpHeaders.Names.CONTENT_LENGTH,
      String.valueOf(fileLength));
  response.setContent(cbuffer);
  channel.write(response).addListener(ChannelFutureListener.CLOSE);
}
 

/**
 * Recursively obtain the tokens for this FileSystem and all descended
 * FileSystems as determined by getChildFileSystems().
 * @param renewer the user allowed to renew the delegation tokens
 * @param credentials cache in which to add the new delegation tokens
 * @param tokens list in which to add acquired tokens
 * @throws IOException
 */
private void collectDelegationTokens(final String renewer,
                                     final Credentials credentials,
                                     final List<Token<?>> tokens)
                                         throws IOException {
  final String serviceName = getCanonicalServiceName();
  // Collect token of the this filesystem and then of its embedded children
  if (serviceName != null) { // fs has token, grab it
    final Text service = new Text(serviceName);
    Token<?> token = credentials.getToken(service);
    if (token == null) {
      token = getDelegationToken(renewer);
      if (token != null) {
        tokens.add(token);
        credentials.addToken(service, token);
      }
    }
  }
  // Now collect the tokens from the children
  final FileSystem[] children = getChildFileSystems();
  if (children != null) {
    for (final FileSystem fs : children) {
      fs.collectDelegationTokens(renewer, credentials, tokens);
    }
  }
}
 

private void addTimelineDelegationToken(
    ContainerLaunchContext clc) throws YarnException, IOException {
  Credentials credentials = new Credentials();
  DataInputByteBuffer dibb = new DataInputByteBuffer();
  ByteBuffer tokens = clc.getTokens();
  if (tokens != null) {
    dibb.reset(tokens);
    credentials.readTokenStorageStream(dibb);
    tokens.rewind();
  }
  // If the timeline delegation token is already in the CLC, no need to add
  // one more
  for (org.apache.hadoop.security.token.Token<? extends TokenIdentifier> token : credentials
      .getAllTokens()) {
    if (token.getKind().equals(TimelineDelegationTokenIdentifier.KIND_NAME)) {
      return;
    }
  }
  org.apache.hadoop.security.token.Token<TimelineDelegationTokenIdentifier>
      timelineDelegationToken = getTimelineDelegationToken();
  if (timelineDelegationToken == null) {
    return;
  }
  credentials.addToken(timelineService, timelineDelegationToken);
  if (LOG.isDebugEnabled()) {
    LOG.debug("Add timline delegation token into credentials: "
        + timelineDelegationToken);
  }
  DataOutputBuffer dob = new DataOutputBuffer();
  credentials.writeTokenStorageToStream(dob);
  tokens = ByteBuffer.wrap(dob.getData(), 0, dob.getLength());
  clc.setTokens(tokens);
}
 

@Test
public void mergeAll() {
  Credentials creds = new Credentials();
  creds.addToken(service[0], token[0]);
  creds.addToken(service[1], token[1]);
  creds.addSecretKey(secret[0], secret[0].getBytes());
  creds.addSecretKey(secret[1], secret[1].getBytes());
  
  Credentials credsToAdd = new Credentials();
  // one duplicate with different value, one new
  credsToAdd.addToken(service[0], token[3]);
  credsToAdd.addToken(service[2], token[2]);
  credsToAdd.addSecretKey(secret[0], secret[3].getBytes());
  credsToAdd.addSecretKey(secret[2], secret[2].getBytes());
  
  creds.mergeAll(credsToAdd);
  assertEquals(3, creds.numberOfTokens());
  assertEquals(3, creds.numberOfSecretKeys());
  // existing token & secret should not be overwritten
  assertEquals(token[0], creds.getToken(service[0]));
  assertEquals(secret[0], new Text(creds.getSecretKey(secret[0])));
  // non-duplicate token & secret should be present
  assertEquals(token[1], creds.getToken(service[1]));
  assertEquals(secret[1], new Text(creds.getSecretKey(secret[1])));
  // new token & secret should be added
  assertEquals(token[2], creds.getToken(service[2]));
  assertEquals(secret[2], new Text(creds.getSecretKey(secret[2])));
}
 

@Test
public void testFsWithNestedDuplicatesChildren() throws Exception {
  Credentials credentials = new Credentials();
  Text service1 = new Text("singleTokenFs1");
  Text service2 = new Text("singleTokenFs2");
  Text service4 = new Text("singleTokenFs4");
  Text multiService = new Text("multiTokenFs");
  Token<?> token2 = mock(Token.class);
  credentials.addToken(service2, token2);
  
  MockFileSystem fs1 = createFileSystemForServiceName(service1);
  MockFileSystem fs1B = createFileSystemForServiceName(service1);
  MockFileSystem fs2 = createFileSystemForServiceName(service2);
  MockFileSystem fs3 = createFileSystemForServiceName(null);
  MockFileSystem fs4 = createFileSystemForServiceName(service4);
  // now let's get dirty!  ensure dup tokens aren't fetched even when
  // repeated and dupped in a nested fs.  fs4 is a real test of the drill
  // down: multi-filter-multi-filter-filter-fs4.
  MockFileSystem multiFs = createFileSystemForServiceName(multiService,
      fs1, fs1B, fs2, fs2, new FilterFileSystem(fs3),
      new FilterFileSystem(new FilterFileSystem(fs4)));
  MockFileSystem superMultiFs = createFileSystemForServiceName(null,
      fs1, fs1B, fs1, new FilterFileSystem(fs3), new FilterFileSystem(multiFs));
  superMultiFs.addDelegationTokens(renewer, credentials);
  verifyTokenFetch(superMultiFs, false); // does not have its own token
  verifyTokenFetch(multiFs, true); // has its own token
  verifyTokenFetch(fs1, true);
  verifyTokenFetch(fs2, false); // we had added its token to credentials
  verifyTokenFetch(fs3, false); // has no tokens
  verifyTokenFetch(fs4, true);
  
  assertEquals(4, credentials.numberOfTokens()); //fs1+fs2+fs4+multifs (fs3=0)
  assertNotNull(credentials.getToken(service1));
  assertNotNull(credentials.getToken(service2));
  assertSame(token2, credentials.getToken(service2));
  assertNotNull(credentials.getToken(multiService));
  assertNotNull(credentials.getToken(service4));
}
 

private void setupTokens(
    ContainerLaunchContext container, ContainerId containerID)
    throws IOException {
  Map<String, String> environment = container.getEnvironment();
  environment.put(ApplicationConstants.APPLICATION_WEB_PROXY_BASE_ENV,
      application.getWebProxyBase());
  // Set AppSubmitTime and MaxAppAttempts to be consumable by the AM.
  ApplicationId applicationId =
      application.getAppAttemptId().getApplicationId();
  environment.put(
      ApplicationConstants.APP_SUBMIT_TIME_ENV,
      String.valueOf(rmContext.getRMApps()
          .get(applicationId)
          .getSubmitTime()));
  environment.put(ApplicationConstants.MAX_APP_ATTEMPTS_ENV,
      String.valueOf(rmContext.getRMApps().get(
          applicationId).getMaxAppAttempts()));

  Credentials credentials = new Credentials();
  DataInputByteBuffer dibb = new DataInputByteBuffer();
  if (container.getTokens() != null) {
    // TODO: Don't do this kind of checks everywhere.
    dibb.reset(container.getTokens());
    credentials.readTokenStorageStream(dibb);
  }

  // Add AMRMToken
  Token<AMRMTokenIdentifier> amrmToken = createAndSetAMRMToken();
  if (amrmToken != null) {
    credentials.addToken(amrmToken.getService(), amrmToken);
  }
  DataOutputBuffer dob = new DataOutputBuffer();
  credentials.writeTokenStorageToStream(dob);
  container.setTokens(ByteBuffer.wrap(dob.getData(), 0, dob.getLength()));
}
 

/**
 * Obtain Kerberos security token for HBase.
 */
private static void obtainTokenForHBase(Credentials credentials, Configuration conf) throws IOException {
	if (UserGroupInformation.isSecurityEnabled()) {
		LOG.info("Attempting to obtain Kerberos security token for HBase");
		try {
			// ----
			// Intended call: HBaseConfiguration.addHbaseResources(conf);
			Class
					.forName("org.apache.hadoop.hbase.HBaseConfiguration")
					.getMethod("addHbaseResources", Configuration.class)
					.invoke(null, conf);
			// ----

			LOG.info("HBase security setting: {}", conf.get("hbase.security.authentication"));

			if (!"kerberos".equals(conf.get("hbase.security.authentication"))) {
				LOG.info("HBase has not been configured to use Kerberos.");
				return;
			}

			LOG.info("Obtaining Kerberos security token for HBase");
			// ----
			// Intended call: Token<AuthenticationTokenIdentifier> token = TokenUtil.obtainToken(conf);
			Token<?> token = (Token<?>) Class
					.forName("org.apache.hadoop.hbase.security.token.TokenUtil")
					.getMethod("obtainToken", Configuration.class)
					.invoke(null, conf);
			// ----

			if (token == null) {
				LOG.error("No Kerberos security token for HBase available");
				return;
			}

			credentials.addToken(token.getService(), token);
			LOG.info("Added HBase Kerberos security token to credentials.");
		} catch (ClassNotFoundException
				| NoSuchMethodException
				| IllegalAccessException
				| InvocationTargetException e) {
			LOG.info("HBase is not available (not packaged with this application): {} : \"{}\".",
					e.getClass().getSimpleName(), e.getMessage());
		}
	}
}
 

@Test (timeout = 60000)
public void testDelegationTokenRestoredInDelegationTokenRenewer()
    throws Exception {
  conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS, 2);
  conf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHENTICATION,
      "kerberos");
  UserGroupInformation.setConfiguration(conf);

  MemoryRMStateStore memStore = new MemoryRMStateStore();
  memStore.init(conf);
  RMState rmState = memStore.getState();

  Map<ApplicationId, ApplicationStateData> rmAppState =
      rmState.getApplicationState();
  MockRM rm1 = new TestSecurityMockRM(conf, memStore);
  rm1.start();

  HashSet<Token<RMDelegationTokenIdentifier>> tokenSet =
      new HashSet<Token<RMDelegationTokenIdentifier>>();

  // create an empty credential
  Credentials ts = new Credentials();

  // create tokens and add into credential
  Text userText1 = new Text("user1");
  RMDelegationTokenIdentifier dtId1 =
      new RMDelegationTokenIdentifier(userText1, new Text("renewer1"),
        userText1);
  Token<RMDelegationTokenIdentifier> token1 =
      new Token<RMDelegationTokenIdentifier>(dtId1,
        rm1.getRMContext().getRMDelegationTokenSecretManager());
  SecurityUtil.setTokenService(token1, rmAddr);
  ts.addToken(userText1, token1);
  tokenSet.add(token1);

  Text userText2 = new Text("user2");
  RMDelegationTokenIdentifier dtId2 =
      new RMDelegationTokenIdentifier(userText2, new Text("renewer2"),
        userText2);
  Token<RMDelegationTokenIdentifier> token2 =
      new Token<RMDelegationTokenIdentifier>(dtId2,
        rm1.getRMContext().getRMDelegationTokenSecretManager());
  SecurityUtil.setTokenService(token2, rmAddr);
  ts.addToken(userText2, token2);
  tokenSet.add(token2);

  // submit an app with customized credential
  RMApp app = rm1.submitApp(200, "name", "user",
      new HashMap<ApplicationAccessType, String>(), false, "default", 1, ts);

  // assert app info is saved
  ApplicationStateData appState = rmAppState.get(app.getApplicationId());
  Assert.assertNotNull(appState);

  // assert delegation tokens exist in rm1 DelegationTokenRenewr
  Assert.assertEquals(tokenSet, rm1.getRMContext()
    .getDelegationTokenRenewer().getDelegationTokens());

  // assert delegation tokens are saved
  DataOutputBuffer dob = new DataOutputBuffer();
  ts.writeTokenStorageToStream(dob);
  ByteBuffer securityTokens =
      ByteBuffer.wrap(dob.getData(), 0, dob.getLength());
  securityTokens.rewind();
  Assert.assertEquals(securityTokens, appState
    .getApplicationSubmissionContext().getAMContainerSpec()
    .getTokens());

  // start new RM
  MockRM rm2 = new TestSecurityMockRM(conf, memStore);
  rm2.start();

  // Need to wait for a while as now token renewal happens on another thread
  // and is asynchronous in nature.
  waitForTokensToBeRenewed(rm2);

  // verify tokens are properly populated back to rm2 DelegationTokenRenewer
  Assert.assertEquals(tokenSet, rm2.getRMContext()
    .getDelegationTokenRenewer().getDelegationTokens());
}
 

@SuppressWarnings("deprecation")
private void testDagCredentials(boolean doMerge) throws IOException {
  TezConfiguration conf = new TezConfiguration();
  conf.setBoolean(TezConfiguration.TEZ_AM_CREDENTIALS_MERGE, doMerge);
  conf.setBoolean(TezConfiguration.TEZ_LOCAL_MODE, true);
  conf.set(TezConfiguration.TEZ_AM_STAGING_DIR, TEST_DIR.toString());
  ApplicationId appId = ApplicationId.newInstance(1, 1);
  ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(appId, 1);

  // create some sample AM credentials
  Credentials amCreds = new Credentials();
  JobTokenSecretManager jtsm = new JobTokenSecretManager();
  JobTokenIdentifier identifier = new JobTokenIdentifier(
      new Text(appId.toString()));
  Token<JobTokenIdentifier> sessionToken =
      new Token<JobTokenIdentifier>(identifier, jtsm);
  sessionToken.setService(identifier.getJobId());
  TokenCache.setSessionToken(sessionToken, amCreds);
  TestTokenSecretManager ttsm = new TestTokenSecretManager();
  Text tokenAlias1 = new Text("alias1");
  Token<TestTokenIdentifier> amToken1 = new Token<TestTokenIdentifier>(
      new TestTokenIdentifier(new Text("amtoken1")), ttsm);
  amCreds.addToken(tokenAlias1, amToken1);
  Text tokenAlias2 = new Text("alias2");
  Token<TestTokenIdentifier> amToken2 = new Token<TestTokenIdentifier>(
      new TestTokenIdentifier(new Text("amtoken2")), ttsm);
  amCreds.addToken(tokenAlias2, amToken2);

  FileSystem fs = FileSystem.getLocal(conf);
  FSDataOutputStream sessionJarsPBOutStream =
      TezCommonUtils.createFileForAM(fs, new Path(TEST_DIR.toString(),
          TezConstants.TEZ_AM_LOCAL_RESOURCES_PB_FILE_NAME));
  DAGProtos.PlanLocalResourcesProto.getDefaultInstance()
      .writeDelimitedTo(sessionJarsPBOutStream);
  sessionJarsPBOutStream.close();
  DAGAppMaster am = new DAGAppMaster(attemptId,
      ContainerId.newInstance(attemptId, 1),
      "127.0.0.1", 0, 0, new SystemClock(), 1, true,
      TEST_DIR.toString(), new String[] {TEST_DIR.toString()},
      new String[] {TEST_DIR.toString()},
      new TezApiVersionInfo().getVersion(), amCreds,
      "someuser", null);
  am.init(conf);
  am.start();

  // create some sample DAG credentials
  Credentials dagCreds = new Credentials();
  Token<TestTokenIdentifier> dagToken1 = new Token<TestTokenIdentifier>(
      new TestTokenIdentifier(new Text("dagtoken1")), ttsm);
  dagCreds.addToken(tokenAlias2, dagToken1);
  Text tokenAlias3 = new Text("alias3");
  Token<TestTokenIdentifier> dagToken2 = new Token<TestTokenIdentifier>(
      new TestTokenIdentifier(new Text("dagtoken2")), ttsm);
  dagCreds.addToken(tokenAlias3, dagToken2);

  TezDAGID dagId = TezDAGID.getInstance(appId, 1);
  DAGPlan dagPlan = DAGPlan.newBuilder()
      .setName("somedag")
      .setCredentialsBinary(
          DagTypeConverters.convertCredentialsToProto(dagCreds))
      .build();
  DAGImpl dag = am.createDAG(dagPlan, dagId);
  Credentials fetchedDagCreds = dag.getCredentials();
  am.stop();

  Token<? extends TokenIdentifier> fetchedToken1 =
      fetchedDagCreds.getToken(tokenAlias1);
  if (doMerge) {
    assertNotNull("AM creds missing from DAG creds", fetchedToken1);
    compareTestTokens(amToken1, fetchedDagCreds.getToken(tokenAlias1));
  } else {
    assertNull("AM creds leaked to DAG creds", fetchedToken1);
  }
  compareTestTokens(dagToken1, fetchedDagCreds.getToken(tokenAlias2));
  compareTestTokens(dagToken2, fetchedDagCreds.getToken(tokenAlias3));
}
 

/**
 * Basic idea of the test:
 * 0. Setup token KEEP_ALIVE
 * 1. create tokens.
 * 2. register them for renewal - to be cancelled on app complete
 * 3. Complete app.
 * 4. Verify token is alive within the KEEP_ALIVE time
 * 5. Verify token has been cancelled after the KEEP_ALIVE_TIME
 * @throws IOException
 * @throws URISyntaxException
 */
@Test(timeout=60000)
public void testDTKeepAlive1 () throws Exception {
  Configuration lconf = new Configuration(conf);
  lconf.setBoolean(YarnConfiguration.LOG_AGGREGATION_ENABLED, true);
  //Keep tokens alive for 6 seconds.
  lconf.setLong(YarnConfiguration.RM_NM_EXPIRY_INTERVAL_MS, 6000l);
  //Try removing tokens every second.
  lconf.setLong(
      YarnConfiguration.RM_DELAYED_DELEGATION_TOKEN_REMOVAL_INTERVAL_MS,
      1000l);
  DelegationTokenRenewer localDtr =
      createNewDelegationTokenRenewer(lconf, counter);
  RMContext mockContext = mock(RMContext.class);
  when(mockContext.getSystemCredentialsForApps()).thenReturn(
    new ConcurrentHashMap<ApplicationId, ByteBuffer>());
  ClientRMService mockClientRMService = mock(ClientRMService.class);
  when(mockContext.getClientRMService()).thenReturn(mockClientRMService);
  when(mockContext.getDelegationTokenRenewer()).thenReturn(
      localDtr);
  when(mockContext.getDispatcher()).thenReturn(dispatcher);
  InetSocketAddress sockAddr =
      InetSocketAddress.createUnresolved("localhost", 1234);
  when(mockClientRMService.getBindAddress()).thenReturn(sockAddr);
  localDtr.setRMContext(mockContext);
  localDtr.init(lconf);
  localDtr.start();
  
  MyFS dfs = (MyFS)FileSystem.get(lconf);
  LOG.info("dfs="+(Object)dfs.hashCode() + ";conf="+lconf.hashCode());
  
  Credentials ts = new Credentials();
  // get the delegation tokens
  MyToken token1 = dfs.getDelegationToken("user1");

  String nn1 = DelegationTokenRenewer.SCHEME + "://host1:0";
  ts.addToken(new Text(nn1), token1);

  // register the tokens for renewal
  ApplicationId applicationId_0 =  BuilderUtils.newApplicationId(0, 0);
  localDtr.addApplicationAsync(applicationId_0, ts, true, "user");
  waitForEventsToGetProcessed(localDtr);
  if (!eventQueue.isEmpty()){
    Event evt = eventQueue.take();
    if (evt instanceof RMAppEvent) {
      Assert.assertEquals(((RMAppEvent)evt).getType(), RMAppEventType.START);
    } else {
      fail("RMAppEvent.START was expected!!");
    }
  }
  
  localDtr.applicationFinished(applicationId_0);
  waitForEventsToGetProcessed(localDtr);

  //Token should still be around. Renewal should not fail.
  token1.renew(lconf);

  //Allow the keepalive time to run out
  Thread.sleep(10000l);

  //The token should have been cancelled at this point. Renewal will fail.
  try {
    token1.renew(lconf);
    fail("Renewal of cancelled token should have failed");
  } catch (InvalidToken ite) {}
}
 

@Test
public void testAppSubmissionWithoutDelegationToken() throws Exception {
  conf.setBoolean(YarnConfiguration.RM_PROXY_USER_PRIVILEGES_ENABLED, true);
  // create token2
  Text userText2 = new Text("user2");
  DelegationTokenIdentifier dtId2 =
      new DelegationTokenIdentifier(new Text("user2"), new Text("renewer2"),
        userText2);
  final Token<DelegationTokenIdentifier> token2 =
      new Token<DelegationTokenIdentifier>(dtId2.getBytes(),
        "password2".getBytes(), dtId2.getKind(), new Text("service2"));
  final MockRM rm = new TestSecurityMockRM(conf, null) {
    @Override
    protected DelegationTokenRenewer createDelegationTokenRenewer() {
      return new DelegationTokenRenewer() {
        @Override
        protected Token<?>[] obtainSystemTokensForUser(String user,
            final Credentials credentials) throws IOException {
          credentials.addToken(token2.getService(), token2);
          return new Token<?>[] { token2 };
        }
      };
    }
  };
  rm.start();

  // submit an app without delegationToken
  RMApp app = rm.submitApp(200);

  // wait for the new retrieved hdfs token.
  GenericTestUtils.waitFor(new Supplier<Boolean>() {
    public Boolean get() {
      return rm.getRMContext().getDelegationTokenRenewer()
        .getDelegationTokens().contains(token2);
    }
  }, 1000, 20000);

  // check nm can retrieve the token
  final MockNM nm1 =
      new MockNM("127.0.0.1:1234", 15120, rm.getResourceTrackerService());
  nm1.registerNode();
  NodeHeartbeatResponse response = nm1.nodeHeartbeat(true);
  ByteBuffer tokenBuffer =
      response.getSystemCredentialsForApps().get(app.getApplicationId());
  Assert.assertNotNull(tokenBuffer);
  Credentials appCredentials = new Credentials();
  DataInputByteBuffer buf = new DataInputByteBuffer();
  tokenBuffer.rewind();
  buf.reset(tokenBuffer);
  appCredentials.readTokenStorageStream(buf);
  Assert.assertTrue(appCredentials.getAllTokens().contains(token2));
}
 

/**
 * store job token
 * @param t
 */
@InterfaceAudience.Private
public static void setJobToken(Token<? extends TokenIdentifier> t, 
    Credentials credentials) {
  credentials.addToken(JOB_TOKEN, t);
}
 

/**
 * Overrides parent method to adds RM delegation token to the given context. If YARN is running with HA RM,
 * delegation tokens for each RM service will be added.
 */
protected void addRMToken(ContainerLaunchContext context, YarnClient yarnClient, ApplicationId appId) {
  if (!UserGroupInformation.isSecurityEnabled()) {
    return;
  }

  try {
    Text renewer = new Text(UserGroupInformation.getCurrentUser().getShortUserName());
    org.apache.hadoop.yarn.api.records.Token rmDelegationToken = yarnClient.getRMDelegationToken(renewer);

    // The following logic is copied from ClientRMProxy.getRMDelegationTokenService, which is not available in
    // YARN older than 2.4
    List<String> services = new ArrayList<>();
    if (HAUtil.isHAEnabled(configuration)) {
      // If HA is enabled, we need to enumerate all RM hosts
      // and add the corresponding service name to the token service
      // Copy the yarn conf since we need to modify it to get the RM addresses
      YarnConfiguration yarnConf = new YarnConfiguration(configuration);
      for (String rmId : HAUtil.getRMHAIds(configuration)) {
        yarnConf.set(YarnConfiguration.RM_HA_ID, rmId);
        InetSocketAddress address = yarnConf.getSocketAddr(YarnConfiguration.RM_ADDRESS,
                                                           YarnConfiguration.DEFAULT_RM_ADDRESS,
                                                           YarnConfiguration.DEFAULT_RM_PORT);
        services.add(SecurityUtil.buildTokenService(address).toString());
      }
    } else {
      services.add(SecurityUtil.buildTokenService(YarnUtils.getRMAddress(configuration)).toString());
    }

    Credentials credentials = YarnUtils.decodeCredentials(context.getTokens());

    // casting needed for later Hadoop version
    @SuppressWarnings("RedundantCast")
    Token<TokenIdentifier> token = ConverterUtils.convertFromYarn(rmDelegationToken, (InetSocketAddress) null);

    token.setService(new Text(Joiner.on(',').join(services)));
    credentials.addToken(new Text(token.getService()), token);

    LOG.debug("Added RM delegation token {} for application {}", token, appId);
    credentials.addToken(token.getService(), token);

    context.setTokens(YarnUtils.encodeCredentials(credentials));

  } catch (Exception e) {
    throw Throwables.propagate(e);
  }
}
 

@Test
public void testCreateTaskExecutorCredentials() throws Exception {
	File root = temporaryFolder.getRoot();
	File home = new File(root, "home");
	boolean created = home.mkdir();
	assertTrue(created);

	Configuration flinkConf = new Configuration();
	YarnConfiguration yarnConf = new YarnConfiguration();

	Map<String, String> env = new HashMap<>();
	env.put(YarnConfigKeys.ENV_APP_ID, "foo");
	env.put(YarnConfigKeys.ENV_CLIENT_HOME_DIR, home.getAbsolutePath());
	env.put(YarnConfigKeys.ENV_CLIENT_SHIP_FILES, "");
	env.put(YarnConfigKeys.ENV_FLINK_CLASSPATH, "");
	env.put(YarnConfigKeys.ENV_HADOOP_USER_NAME, "foo");
	env.put(YarnConfigKeys.FLINK_JAR_PATH, root.toURI().toString());
	env = Collections.unmodifiableMap(env);

	File credentialFile = temporaryFolder.newFile("container_tokens");
	final Text amRmTokenKind = AMRMTokenIdentifier.KIND_NAME;
	final Text hdfsDelegationTokenKind = new Text("HDFS_DELEGATION_TOKEN");
	final Text service = new Text("test-service");
	Credentials amCredentials = new Credentials();
	amCredentials.addToken(amRmTokenKind, new Token<>(new byte[4], new byte[4], amRmTokenKind, service));
	amCredentials.addToken(hdfsDelegationTokenKind, new Token<>(new byte[4], new byte[4],
		hdfsDelegationTokenKind, service));
	amCredentials.writeTokenStorageFile(new org.apache.hadoop.fs.Path(credentialFile.getAbsolutePath()), yarnConf);

	ContaineredTaskManagerParameters tmParams = new ContaineredTaskManagerParameters(64,
		64, 16, 1, new HashMap<>(1));
	Configuration taskManagerConf = new Configuration();

	String workingDirectory = root.getAbsolutePath();
	Class<?> taskManagerMainClass = YarnTaskExecutorRunner.class;
	ContainerLaunchContext ctx;

	final Map<String, String> originalEnv = System.getenv();
	try {
		Map<String, String> systemEnv = new HashMap<>(originalEnv);
		systemEnv.put("HADOOP_TOKEN_FILE_LOCATION", credentialFile.getAbsolutePath());
		CommonTestUtils.setEnv(systemEnv);
		ctx = Utils.createTaskExecutorContext(flinkConf, yarnConf, env, tmParams,
			taskManagerConf, workingDirectory, taskManagerMainClass, LOG);
	} finally {
		CommonTestUtils.setEnv(originalEnv);
	}

	Credentials credentials = new Credentials();
	try (DataInputStream dis = new DataInputStream(new ByteArrayInputStream(ctx.getTokens().array()))) {
		credentials.readTokenStorageStream(dis);
	}
	Collection<Token<? extends TokenIdentifier>> tokens = credentials.getAllTokens();
	boolean hasHdfsDelegationToken = false;
	boolean hasAmRmToken = false;
	for (Token<? extends TokenIdentifier> token : tokens) {
		if (token.getKind().equals(amRmTokenKind)) {
			hasAmRmToken = true;
		} else if (token.getKind().equals(hdfsDelegationTokenKind)) {
			hasHdfsDelegationToken = true;
		}
	}
	assertTrue(hasHdfsDelegationToken);
	assertFalse(hasAmRmToken);
}
 

/**
 * Basic idea of the test:
 * 0. Setup token KEEP_ALIVE
 * 1. create tokens.
 * 2. register them for renewal - to be cancelled on app complete
 * 3. Complete app.
 * 4. Verify token is alive within the KEEP_ALIVE time
 * 5. Send an explicity KEEP_ALIVE_REQUEST
 * 6. Verify token KEEP_ALIVE time is renewed.
 * 7. Verify token has been cancelled after the renewed KEEP_ALIVE_TIME.
 * @throws IOException
 * @throws URISyntaxException
 */
@Test(timeout=60000)
public void testDTKeepAlive2() throws Exception {
  Configuration lconf = new Configuration(conf);
  lconf.setBoolean(YarnConfiguration.LOG_AGGREGATION_ENABLED, true);
  //Keep tokens alive for 6 seconds.
  lconf.setLong(YarnConfiguration.RM_NM_EXPIRY_INTERVAL_MS, 6000l);
  //Try removing tokens every second.
  lconf.setLong(
      YarnConfiguration.RM_DELAYED_DELEGATION_TOKEN_REMOVAL_INTERVAL_MS,
      1000l);
  DelegationTokenRenewer localDtr =
      createNewDelegationTokenRenewer(conf, counter);
  RMContext mockContext = mock(RMContext.class);
  when(mockContext.getSystemCredentialsForApps()).thenReturn(
    new ConcurrentHashMap<ApplicationId, ByteBuffer>());
  ClientRMService mockClientRMService = mock(ClientRMService.class);
  when(mockContext.getClientRMService()).thenReturn(mockClientRMService);
  when(mockContext.getDelegationTokenRenewer()).thenReturn(
      localDtr);
  when(mockContext.getDispatcher()).thenReturn(dispatcher);
  InetSocketAddress sockAddr =
      InetSocketAddress.createUnresolved("localhost", 1234);
  when(mockClientRMService.getBindAddress()).thenReturn(sockAddr);
  localDtr.setRMContext(mockContext);
  localDtr.init(lconf);
  localDtr.start();
  
  MyFS dfs = (MyFS)FileSystem.get(lconf);
  LOG.info("dfs="+(Object)dfs.hashCode() + ";conf="+lconf.hashCode());

  Credentials ts = new Credentials();
  // get the delegation tokens
  MyToken token1 = dfs.getDelegationToken("user1");
  
  String nn1 = DelegationTokenRenewer.SCHEME + "://host1:0";
  ts.addToken(new Text(nn1), token1);

  // register the tokens for renewal
  ApplicationId applicationId_0 =  BuilderUtils.newApplicationId(0, 0);
  localDtr.addApplicationAsync(applicationId_0, ts, true, "user");
  localDtr.applicationFinished(applicationId_0);
  waitForEventsToGetProcessed(delegationTokenRenewer);
  //Send another keep alive.
  localDtr.updateKeepAliveApplications(Collections
      .singletonList(applicationId_0));
  //Renewal should not fail.
  token1.renew(lconf);
  //Token should be around after this. 
  Thread.sleep(4500l);
  //Renewal should not fail. - ~1.5 seconds for keepalive timeout.
  token1.renew(lconf);
  //Allow the keepalive time to run out
  Thread.sleep(3000l);
  //The token should have been cancelled at this point. Renewal will fail.
  try {
    token1.renew(lconf);
    fail("Renewal of cancelled token should have failed");
  } catch (InvalidToken ite) {}
}
 

/**
 * Write a {@link Token} to a given file.
 *
 * @param token the token to write
 * @param tokenFilePath the token file path
 * @param configuration a {@link Configuration} object carrying Hadoop configuration properties
 * @throws IOException
 */
public static void writeTokenToFile(Token<? extends TokenIdentifier> token, Path tokenFilePath,
    Configuration configuration) throws IOException {
  Credentials credentials = new Credentials();
  credentials.addToken(token.getService(), token);
  credentials.writeTokenStorageFile(tokenFilePath, configuration);
}