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

private Credentials createCredentials() {
  Credentials credentials = new Credentials();
  if (!UserGroupInformation.isSecurityEnabled()) {
    return credentials;
  }

  try {
    credentials.addAll(UserGroupInformation.getCurrentUser().getCredentials());

    // Remove the AM->RM tokens
    Iterator<Token<?>> iter = credentials.getAllTokens().iterator();
    while (iter.hasNext()) {
      Token<?> token = iter.next();
      if (token.getKind().equals(AMRM_TOKEN_KIND_NAME)) {
        iter.remove();
      }
    }
  } catch (IOException e) {
    LOG.warn("Failed to get current user. No credentials will be provided to containers.", e);
  }

  return credentials;
}
 

/**
 * Creates a {@link Credentials} for the application submission.
 */
private Credentials createSubmissionCredentials() {
  Credentials credentials = new Credentials();
  try {
    // Acquires delegation token for the location
    List<Token<?>> tokens = YarnUtils.addDelegationTokens(config, appLocation.getLocationFactory(), credentials);
    if (LOG.isDebugEnabled()) {
      for (Token<?> token : tokens) {
        LOG.debug("Delegation token acquired for {}, {}", appLocation, token);
      }
    }
  } catch (IOException e) {
    LOG.warn("Failed to acquire delegation token for location {}", appLocation);
  }

  // Copy the user provided credentials.
  // It will override the location delegation tokens acquired above if user supplies it.
  credentials.addAll(this.credentials);
  return credentials;
}
 

private SubmitWorkRequestProto constructSubmitWorkRequest(ContainerId containerId,
                                                          TaskSpec taskSpec) throws
    IOException {
  SubmitWorkRequestProto.Builder builder =
      SubmitWorkRequestProto.newBuilder(BASE_SUBMIT_WORK_REQUEST);
  builder.setContainerIdString(containerId.toString());
  builder.setAmHost(getAddress().getHostName());
  builder.setAmPort(getAddress().getPort());
  Credentials taskCredentials = new Credentials();
  // Credentials can change across DAGs. Ideally construct only once per DAG.
  taskCredentials.addAll(getContext().getAMCredentials());

  ByteBuffer credentialsBinary = credentialMap.get(taskSpec.getDAGName());
  if (credentialsBinary == null) {
    credentialsBinary = serializeCredentials(getContext().getAMCredentials());
    credentialMap.putIfAbsent(taskSpec.getDAGName(), credentialsBinary.duplicate());
  } else {
    credentialsBinary = credentialsBinary.duplicate();
  }
  builder.setCredentialsBinary(ByteString.copyFrom(credentialsBinary));
  builder.setTaskSpec(ProtoConverters.convertTaskSpecToProto(taskSpec));
  return builder.build();
}
 

@Test
public void addAll() {
  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.addAll(credsToAdd);
  assertEquals(3, creds.numberOfTokens());
  assertEquals(3, creds.numberOfSecretKeys());
  // existing token & secret should be overwritten
  assertEquals(token[3], creds.getToken(service[0]));
  assertEquals(secret[3], 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])));
}
 

/**
 * Creates an {@link Credentials} by copying the {@link Credentials} of the current user.
 */
private Credentials createCredentials() {
  Credentials credentials = new Credentials();

  try {
    credentials.addAll(UserGroupInformation.getCurrentUser().getCredentials());
  } catch (IOException e) {
    LOG.warn("Failed to get current user UGI. Current user credentials not added.", e);
  }
  return credentials;
}
 

@Test
public void addAll() {
  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.addAll(credsToAdd);
  assertEquals(3, creds.numberOfTokens());
  assertEquals(3, creds.numberOfSecretKeys());
  // existing token & secret should be overwritten
  assertEquals(token[3], creds.getToken(service[0]));
  assertEquals(secret[3], 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])));
}
 

private void setupSecurityTokens(ContainerLaunchContext containerLaunchContext) throws IOException {
  Credentials credentials = UserGroupInformation.getCurrentUser().getCredentials();

  // Pass on the credentials from the hadoop token file if present.
  // The value in the token file takes precedence.
  if (System.getenv(HADOOP_TOKEN_FILE_LOCATION) != null) {
    Credentials tokenFileCredentials = Credentials.readTokenStorageFile(new File(System.getenv(HADOOP_TOKEN_FILE_LOCATION)),
        new Configuration());
    credentials.addAll(tokenFileCredentials);
  }

  String tokenRenewer = this.yarnConfiguration.get(YarnConfiguration.RM_PRINCIPAL);
  if (tokenRenewer == null || tokenRenewer.length() == 0) {
    throw new IOException("Failed to get master Kerberos principal for the RM to use as renewer");
  }

  // For now, only getting tokens for the default file-system.
  Token<?> tokens[] = this.fs.addDelegationTokens(tokenRenewer, credentials);
  if (tokens != null) {
    for (Token<?> token : tokens) {
      LOGGER.info("Got delegation token for " + this.fs.getUri() + "; " + token);
    }
  }

  Closer closer = Closer.create();
  try {
    DataOutputBuffer dataOutputBuffer = closer.register(new DataOutputBuffer());
    credentials.writeTokenStorageToStream(dataOutputBuffer);
    ByteBuffer fsTokens = ByteBuffer.wrap(dataOutputBuffer.getData(), 0, dataOutputBuffer.getLength());
    containerLaunchContext.setTokens(fsTokens);
  } catch (Throwable t) {
    throw closer.rethrow(t);
  } finally {
    closer.close();
  }
}
 

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

/**
 * 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;
}
 

/**
 * 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;
}