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

public static void setTokensFor(ContainerLaunchContext amContainer, List<Path> paths, Configuration conf) throws IOException {
	Credentials credentials = new Credentials();
	// for HDFS
	TokenCache.obtainTokensForNamenodes(credentials, paths.toArray(new Path[0]), conf);
	// for HBase
	obtainTokenForHBase(credentials, conf);
	// for user
	UserGroupInformation currUsr = UserGroupInformation.getCurrentUser();

	Collection<Token<? extends TokenIdentifier>> usrTok = currUsr.getTokens();
	for (Token<? extends TokenIdentifier> token : usrTok) {
		final Text id = new Text(token.getIdentifier());
		LOG.info("Adding user token " + id + " with " + token);
		credentials.addToken(id, token);
	}
	try (DataOutputBuffer dob = new DataOutputBuffer()) {
		credentials.writeTokenStorageToStream(dob);

		if (LOG.isDebugEnabled()) {
			LOG.debug("Wrote tokens. Credentials buffer length: " + dob.getLength());
		}

		ByteBuffer securityTokens = ByteBuffer.wrap(dob.getData(), 0, dob.getLength());
		amContainer.setTokens(securityTokens);
	}
}
 

@Test (timeout = 30000)
public void testAppRecoverPath() throws IOException {
  LOG.info("--- START: testAppRecoverPath ---");
  ApplicationSubmissionContext sub =
      Records.newRecord(ApplicationSubmissionContext.class);
  ContainerLaunchContext clc =
      Records.newRecord(ContainerLaunchContext.class);
  Credentials credentials = new Credentials();
  DataOutputBuffer dob = new DataOutputBuffer();
  credentials.writeTokenStorageToStream(dob);
  ByteBuffer securityTokens =
      ByteBuffer.wrap(dob.getData(), 0, dob.getLength());
  clc.setTokens(securityTokens);
  sub.setAMContainerSpec(clc);
  testCreateAppSubmittedRecovery(sub);
}
 

@SuppressWarnings({ "rawtypes", "unchecked" })
static DataInputBuffer createFakeCredentials(Random r, int nTok)
      throws IOException {
    Credentials creds = new Credentials();
    byte[] password = new byte[20];
    Text kind = new Text();
    Text service = new Text();
    Text alias = new Text();
    for (int i = 0; i < nTok; ++i) {
      byte[] identifier = ("idef" + i).getBytes();
      r.nextBytes(password);
      kind.set("kind" + i);
      service.set("service" + i);
      alias.set("token" + i);
      Token token = new Token(identifier, password, kind, service);
      creds.addToken(alias, token);
    }
    DataOutputBuffer buf = new DataOutputBuffer();
    creds.writeTokenStorageToStream(buf);
    DataInputBuffer ret = new DataInputBuffer();
    ret.reset(buf.getData(), 0, buf.getLength());
    return ret;
  }
 

private static void createBinaryTokenFile(Configuration conf) {
  // Fetch delegation tokens and store in binary token file.
  try {
    Credentials cred1 = new Credentials();
    Credentials cred2 = new Credentials();
    TokenCache.obtainTokensForNamenodesInternal(cred1, new Path[] { p1 },
        conf);
    for (Token<? extends TokenIdentifier> t : cred1.getAllTokens()) {
      cred2.addToken(new Text(DELEGATION_TOKEN_KEY), t);
    }
    DataOutputStream os = new DataOutputStream(new FileOutputStream(
        binaryTokenFileName.toString()));
    try {
      cred2.writeTokenStorageToStream(os);
    } finally {
      os.close();
    }
  } catch (IOException e) {
    Assert.fail("Exception " + e);
  }
}
 

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);
  }
}
 

public static void setTokensFor(ContainerLaunchContext amContainer, List<Path> paths, Configuration conf) throws IOException {
	Credentials credentials = new Credentials();
	// for HDFS
	TokenCache.obtainTokensForNamenodes(credentials, paths.toArray(new Path[0]), conf);
	// for HBase
	obtainTokenForHBase(credentials, conf);
	// for user
	UserGroupInformation currUsr = UserGroupInformation.getCurrentUser();

	Collection<Token<? extends TokenIdentifier>> usrTok = currUsr.getTokens();
	for (Token<? extends TokenIdentifier> token : usrTok) {
		final Text id = new Text(token.getIdentifier());
		LOG.info("Adding user token " + id + " with " + token);
		credentials.addToken(id, token);
	}
	try (DataOutputBuffer dob = new DataOutputBuffer()) {
		credentials.writeTokenStorageToStream(dob);

		if (LOG.isDebugEnabled()) {
			LOG.debug("Wrote tokens. Credentials buffer length: " + dob.getLength());
		}

		ByteBuffer securityTokens = ByteBuffer.wrap(dob.getData(), 0, dob.getLength());
		amContainer.setTokens(securityTokens);
	}
}
 

@Override
public byte[] getTokens() {

  try {

    final UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
    final Credentials credentials = ugi.getCredentials();

    LOG.log(Level.FINEST, "Got {0} tokens for user {1}", new Object[] {credentials.numberOfTokens(), ugi});

    if (credentials.numberOfTokens() > 0) {
      try (DataOutputBuffer dob = new DataOutputBuffer()) {
        credentials.writeTokenStorageToStream(dob);
        return dob.getData();
      }
    }
  } catch (final IOException e) {
    LOG.log(Level.WARNING, "Could not access tokens in user credentials.", e);
  }

  LOG.log(Level.FINE, "No security token found.");

  return null;
}
 

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()));
}
 

private ByteBuffer serializeCredentials(Credentials credentials) throws IOException {
  Credentials containerCredentials = new Credentials();
  containerCredentials.addAll(credentials);
  DataOutputBuffer containerTokens_dob = new DataOutputBuffer();
  containerCredentials.writeTokenStorageToStream(containerTokens_dob);
  ByteBuffer containerCredentialsBuffer = ByteBuffer.wrap(containerTokens_dob.getData(), 0,
      containerTokens_dob.getLength());
  return containerCredentialsBuffer;
}
 

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);
}
 

/**
 * Encodes the given {@link Credentials} as bytes.
 */
public static ByteBuffer encodeCredentials(Credentials credentials) {
  try {
    DataOutputBuffer out = new DataOutputBuffer();
    credentials.writeTokenStorageToStream(out);
    return ByteBuffer.wrap(out.getData(), 0, out.getLength());
  } catch (IOException e) {
    // Shouldn't throw
    LOG.error("Failed to encode Credentials.", e);
    throw Throwables.propagate(e);
  }
}
 

@Test
public void testNodeHeartBeatResponse() throws IOException {
  NodeHeartbeatResponse record =
      Records.newRecord(NodeHeartbeatResponse.class);
  Map<ApplicationId, ByteBuffer> appCredentials =
      new HashMap<ApplicationId, ByteBuffer>();
  Credentials app1Cred = new Credentials();

  Token<DelegationTokenIdentifier> token1 =
      new Token<DelegationTokenIdentifier>();
  token1.setKind(new Text("kind1"));
  app1Cred.addToken(new Text("token1"), token1);
  Token<DelegationTokenIdentifier> token2 =
      new Token<DelegationTokenIdentifier>();
  token2.setKind(new Text("kind2"));
  app1Cred.addToken(new Text("token2"), token2);

  DataOutputBuffer dob = new DataOutputBuffer();
  app1Cred.writeTokenStorageToStream(dob);
  ByteBuffer byteBuffer1 = ByteBuffer.wrap(dob.getData(), 0, dob.getLength());
  appCredentials.put(ApplicationId.newInstance(1234, 1), byteBuffer1);
  record.setSystemCredentialsForApps(appCredentials);

  NodeHeartbeatResponse proto =
      new NodeHeartbeatResponsePBImpl(
        ((NodeHeartbeatResponsePBImpl) record).getProto());
  Assert.assertEquals(appCredentials, proto.getSystemCredentialsForApps());
}
 

private ContainerManagerApplicationProto buildAppProto(ApplicationId appId,
    String user, Credentials credentials,
    Map<ApplicationAccessType, String> appAcls,
    LogAggregationContext logAggregationContext) {

  ContainerManagerApplicationProto.Builder builder =
      ContainerManagerApplicationProto.newBuilder();
  builder.setId(((ApplicationIdPBImpl) appId).getProto());
  builder.setUser(user);

  if (logAggregationContext != null) {
    builder.setLogAggregationContext((
        (LogAggregationContextPBImpl)logAggregationContext).getProto());
  }

  builder.clearCredentials();
  if (credentials != null) {
    DataOutputBuffer dob = new DataOutputBuffer();
    try {
      credentials.writeTokenStorageToStream(dob);
      builder.setCredentials(ByteString.copyFrom(dob.getData()));
    } catch (IOException e) {
      // should not occur
      LOG.error("Cannot serialize credentials", e);
    }
  }

  builder.clearAcls();
  if (appAcls != null) {
    for (Map.Entry<ApplicationAccessType, String> acl : appAcls.entrySet()) {
      ApplicationACLMapProto p = ApplicationACLMapProto.newBuilder()
          .setAccessType(ProtoUtils.convertToProtoFormat(acl.getKey()))
          .setAcl(acl.getValue())
          .build();
      builder.addAcls(p);
    }
  }

  return builder.build();
}
 

private static ByteBuffer getTokens() throws IOException {

    final UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
    final Credentials credentials = ugi.getCredentials();

    try (DataOutputBuffer dob = new DataOutputBuffer()) {
      credentials.writeTokenStorageToStream(dob);
      return ByteBuffer.wrap(dob.getData());
    }
  }
 

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()));
}
 

/**
 * Create the common {@link ContainerLaunchContext} for all attempts.
 *
 * @param applicationACLs
 */
private static ContainerLaunchContext createCommonContainerLaunchContext(
    Map<ApplicationAccessType, String> applicationACLs,
    Credentials credentials) {

  // Application resources
  Map<String, LocalResource> localResources =
      new HashMap<String, LocalResource>();

  // Application environment
  Map<String, String> environment = new HashMap<String, String>();

  // Service data
  Map<String, ByteBuffer> serviceData = new HashMap<String, ByteBuffer>();

  // Tokens
  
  // Setup up task credentials buffer
  ByteBuffer containerCredentialsBuffer = ByteBuffer.wrap(new byte[] {});
  try {
    Credentials containerCredentials = new Credentials();
    
    // All Credentials need to be set so that YARN can localize the resources
    // correctly, even though they may not be used by all tasks which will run
    // on this container.

    LOG.info("Adding #" + credentials.numberOfTokens() + " tokens and #"
        + credentials.numberOfSecretKeys() + " secret keys for NM use for launching container");
    containerCredentials.addAll(credentials);

    DataOutputBuffer containerTokens_dob = new DataOutputBuffer();
    containerCredentials.writeTokenStorageToStream(containerTokens_dob);
    containerCredentialsBuffer = ByteBuffer.wrap(containerTokens_dob.getData(), 0,
        containerTokens_dob.getLength());

    // Add shuffle token
    LOG.info("Putting shuffle token in serviceData");
    serviceData.put(TezConfiguration.TEZ_SHUFFLE_HANDLER_SERVICE_ID,
        serializeServiceData(TokenCache.getSessionToken(containerCredentials)));
  } catch (IOException e) {
    throw new TezUncheckedException(e);
  }
  // Construct the actual Container
  // The null fields are per-container and will be constructed for each
  // container separately.
  ContainerLaunchContext container =
      ContainerLaunchContext.newInstance(localResources, environment, null,
          serviceData, containerCredentialsBuffer, applicationACLs);
  return container;
}
 

@SuppressWarnings({ "unchecked" })
public void run() throws YarnException, IOException, InterruptedException {
  LOG.info("Starting ApplicationMaster");

  // Note: Credentials, Token, UserGroupInformation, DataOutputBuffer class
  // are marked as LimitedPrivate
  Credentials credentials =
      UserGroupInformation.getCurrentUser().getCredentials();
  DataOutputBuffer dob = new DataOutputBuffer();
  credentials.writeTokenStorageToStream(dob);
  // Now remove the AM->RM token so that containers cannot access it.
  Iterator<Token<?>> iter = credentials.getAllTokens().iterator();
  LOG.info("Executing with tokens:");
  while (iter.hasNext()) {
    Token<?> token = iter.next();
    LOG.info(token);
    if (token.getKind().equals(AMRMTokenIdentifier.KIND_NAME)) {
      iter.remove();
    }
  }
  allTokens = ByteBuffer.wrap(dob.getData(), 0, dob.getLength());

  // Create appSubmitterUgi and add original tokens to it
  String appSubmitterUserName =
      System.getenv(ApplicationConstants.Environment.USER.name());
  appSubmitterUgi =
      UserGroupInformation.createRemoteUser(appSubmitterUserName);
  appSubmitterUgi.addCredentials(credentials);


  AMRMClientAsync.CallbackHandler allocListener = new RMCallbackHandler();
  amRMClient = AMRMClientAsync.createAMRMClientAsync(1000, allocListener);
  amRMClient.init(conf);
  amRMClient.start();

  containerListener = createNMCallbackHandler();
  nmClientAsync = new NMClientAsyncImpl(containerListener);
  nmClientAsync.init(conf);
  nmClientAsync.start();

  startTimelineClient(conf);
  if(timelineClient != null) {
    publishApplicationAttemptEvent(timelineClient, appAttemptID.toString(),
        DSEvent.DS_APP_ATTEMPT_START, domainId, appSubmitterUgi);
  }

  // Setup local RPC Server to accept status requests directly from clients
  // TODO need to setup a protocol for client to be able to communicate to
  // the RPC server
  // TODO use the rpc port info to register with the RM for the client to
  // send requests to this app master

  // Register self with ResourceManager
  // This will start heartbeating to the RM
  appMasterHostname = NetUtils.getHostname();
  RegisterApplicationMasterResponse response = amRMClient
      .registerApplicationMaster(appMasterHostname, appMasterRpcPort,
          appMasterTrackingUrl);
  // Dump out information about cluster capability as seen by the
  // resource manager
  int maxMem = response.getMaximumResourceCapability().getMemory();
  LOG.info("Max mem capabililty of resources in this cluster " + maxMem);
  
  int maxVCores = response.getMaximumResourceCapability().getVirtualCores();
  LOG.info("Max vcores capabililty of resources in this cluster " + maxVCores);

  // A resource ask cannot exceed the max.
  if (containerMemory > maxMem) {
    LOG.info("Container memory specified above max threshold of cluster."
        + " Using max value." + ", specified=" + containerMemory + ", max="
        + maxMem);
    containerMemory = maxMem;
  }

  if (containerVirtualCores > maxVCores) {
    LOG.info("Container virtual cores specified above max threshold of cluster."
        + " Using max value." + ", specified=" + containerVirtualCores + ", max="
        + maxVCores);
    containerVirtualCores = maxVCores;
  }

  List<Container> previousAMRunningContainers =
      response.getContainersFromPreviousAttempts();
  LOG.info(appAttemptID + " received " + previousAMRunningContainers.size()
    + " previous attempts' running containers on AM registration.");
  numAllocatedContainers.addAndGet(previousAMRunningContainers.size());

  int numTotalContainersToRequest =
      numTotalContainers - previousAMRunningContainers.size();
  // Setup ask for containers from RM
  // Send request for containers to RM
  // Until we get our fully allocated quota, we keep on polling RM for
  // containers
  // Keep looping until all the containers are launched and shell script
  // executed on them ( regardless of success/failure).
  for (int i = 0; i < numTotalContainersToRequest; ++i) {
    ContainerRequest containerAsk = setupContainerAskForRM();
    amRMClient.addContainerRequest(containerAsk);
  }
  numRequestedContainers.set(numTotalContainers);
}
 

@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());
}
 

/**
 * Create the common {@link ContainerLaunchContext} for all attempts.
 *
 * @param applicationACLs
 * @param auxiliaryService
 */
private static ContainerLaunchContext createCommonContainerLaunchContext(
    Map<ApplicationAccessType, String> applicationACLs,
    Credentials credentials, String auxiliaryService) {

  // Application environment
  Map<String, String> environment = new HashMap<String, String>();

  // Service data
  Map<String, ByteBuffer> serviceData = new HashMap<String, ByteBuffer>();

  // Tokens
  
  // Setup up task credentials buffer
  ByteBuffer containerCredentialsBuffer = ByteBuffer.wrap(new byte[] {});
  try {
    Credentials containerCredentials = new Credentials();
    
    // All Credentials need to be set so that YARN can localize the resources
    // correctly, even though they may not be used by all tasks which will run
    // on this container.

    if (LOG.isDebugEnabled()) {
      LOG.debug("Adding #" + credentials.numberOfTokens() + " tokens and #"
          + credentials.numberOfSecretKeys() + " secret keys for NM use for launching container in common CLC");
    }
    containerCredentials.addAll(credentials);

    DataOutputBuffer containerTokens_dob = new DataOutputBuffer();
    containerCredentials.writeTokenStorageToStream(containerTokens_dob);
    containerCredentialsBuffer = ByteBuffer.wrap(containerTokens_dob.getData(), 0,
        containerTokens_dob.getLength());

    // Add shuffle token
    if (LOG.isDebugEnabled()) {
      LOG.debug("Putting shuffle token in serviceData in common CLC");
    }
    serviceData.put(auxiliaryService,
        TezCommonUtils.serializeServiceData(TokenCache.getSessionToken(containerCredentials)));
  } catch (IOException e) {
    throw new TezUncheckedException(e);
  }
  // Construct the actual Container
  // The null fields are per-container and will be constructed for each
  // container separately.
  ContainerLaunchContext container =
      ContainerLaunchContext.newInstance(null, environment, null,
          serviceData, containerCredentialsBuffer, applicationACLs);
  return container;
}
 

@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());
}