Java Code Examples for org.apache.hadoop.security.UserGroupInformation#addCredentials()

/**
 * Add serialized token to teh credentials.
 * @param tokens ByteBuffer containing token.
 */
@Override
public void addTokens(final byte[] tokens) {

  try (DataInputBuffer buf = new DataInputBuffer()) {

    buf.reset(tokens, tokens.length);
    final Credentials credentials = new Credentials();
    credentials.readTokenStorageStream(buf);

    final UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
    ugi.addCredentials(credentials);
    LOG.log(Level.FINEST, "Added {0} tokens for user {1}", new Object[] {credentials.numberOfTokens(), ugi});

  } catch (final IOException ex) {
    LOG.log(Level.SEVERE, "Could not access tokens in user credentials.", ex);
    throw new RuntimeException(ex);
  }
}
 

/**
 * Setup
 * 
 * @param containerTask
 *          the new task specification. Must be a valid task
 * @param childUGI
 *          the old UGI instance being used
 * @return
 */
UserGroupInformation handleNewTaskCredentials(ContainerTask containerTask,
    UserGroupInformation childUGI) {
  // Re-use the UGI only if the Credentials have not changed.
  Preconditions.checkState(containerTask.shouldDie() != true);
  Preconditions.checkState(containerTask.getTaskSpec() != null);
  if (containerTask.haveCredentialsChanged()) {
    LOG.info("Refreshing UGI since Credentials have changed");
    Credentials taskCreds = containerTask.getCredentials();
    if (taskCreds != null) {
      LOG.info("Credentials : #Tokens=" + taskCreds.numberOfTokens() + ", #SecretKeys="
          + taskCreds.numberOfSecretKeys());
      childUGI = UserGroupInformation.createRemoteUser(System
          .getenv(ApplicationConstants.Environment.USER.toString()));
      childUGI.addCredentials(containerTask.getCredentials());
    } else {
      LOG.info("Not loading any credentials, since no credentials provided");
    }
  }
  return childUGI;
}
 

/**
 * Setup
 * 
 * @param containerTask
 *          the new task specification. Must be a valid task
 * @param childUGI
 *          the old UGI instance being used
 * @return childUGI
 */
UserGroupInformation handleNewTaskCredentials(ContainerTask containerTask,
    UserGroupInformation childUGI) {
  // Re-use the UGI only if the Credentials have not changed.
  Preconditions.checkState(!containerTask.shouldDie());
  Preconditions.checkState(containerTask.getTaskSpec() != null);
  if (containerTask.haveCredentialsChanged()) {
    Credentials taskCreds = containerTask.getCredentials();
    if (taskCreds != null) {
      LOG.info("Refreshing UGI since Credentials have changed. Credentials : #Tokens=" +
          taskCreds.numberOfTokens() + ", #SecretKeys="
          + taskCreds.numberOfSecretKeys());
      childUGI = UserGroupInformation.createRemoteUser(user);
      childUGI.addCredentials(containerTask.getCredentials());
    } else {
      LOG.info("Not loading any credentials, since no credentials provided");
    }
  }
  return childUGI;
}
 

@Override
public ContainerExecutionResult call() throws Exception {

  // TODO Consolidate this code with TezChild.
  StopWatch sw = new StopWatch().start();
  UserGroupInformation taskUgi = UserGroupInformation.createRemoteUser(request.getUser());
  taskUgi.addCredentials(credentials);

  Token<JobTokenIdentifier> jobToken = TokenCache.getSessionToken(credentials);
  Map<String, ByteBuffer> serviceConsumerMetadata = new HashMap<String, ByteBuffer>();
  String auxiliaryService = conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID,
      TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT);
  serviceConsumerMetadata.put(auxiliaryService,
      TezCommonUtils.convertJobTokenToBytes(jobToken));
  Multimap<String, String> startedInputsMap = HashMultimap.create();

  UserGroupInformation taskOwner =
      UserGroupInformation.createRemoteUser(request.getTokenIdentifier());
  final InetSocketAddress address =
      NetUtils.createSocketAddrForHost(request.getAmHost(), request.getAmPort());
  SecurityUtil.setTokenService(jobToken, address);
  taskOwner.addToken(jobToken);
  umbilical = taskOwner.doAs(new PrivilegedExceptionAction<TezTaskUmbilicalProtocol>() {
    @Override
    public TezTaskUmbilicalProtocol run() throws Exception {
      return RPC.getProxy(TezTaskUmbilicalProtocol.class,
          TezTaskUmbilicalProtocol.versionID, address, conf);
    }
  });
  // TODO Stop reading this on each request.
  taskReporter = new TaskReporter(
      umbilical,
      conf.getInt(TezConfiguration.TEZ_TASK_AM_HEARTBEAT_INTERVAL_MS,
          TezConfiguration.TEZ_TASK_AM_HEARTBEAT_INTERVAL_MS_DEFAULT),
      conf.getLong(
          TezConfiguration.TEZ_TASK_AM_HEARTBEAT_COUNTER_INTERVAL_MS,
          TezConfiguration.TEZ_TASK_AM_HEARTBEAT_COUNTER_INTERVAL_MS_DEFAULT),
      conf.getInt(TezConfiguration.TEZ_TASK_MAX_EVENTS_PER_HEARTBEAT,
          TezConfiguration.TEZ_TASK_MAX_EVENTS_PER_HEARTBEAT_DEFAULT),
      new AtomicLong(0),
      request.getContainerIdString());

  TezCommonUtils.logCredentials(LOG, taskUgi.getCredentials(), "taskUgi");
  taskRunner = new TezTaskRunner2(conf, taskUgi, localDirs,
      ProtoConverters.getTaskSpecfromProto(request.getTaskSpec()),
      request.getAppAttemptNumber(),
      serviceConsumerMetadata, envMap, startedInputsMap, taskReporter, executor, objectRegistry,
      pid,
      executionContext, memoryAvailable, false, new DefaultHadoopShim(), sharedExecutor);

  boolean shouldDie;
  try {
    TaskRunner2Result result = taskRunner.run();
    LOG.info("TaskRunner2Result: {}", result);
    shouldDie = result.isContainerShutdownRequested();
    if (shouldDie) {
      LOG.info("Got a shouldDie notification via heartbeats. Shutting down");
      return new ContainerExecutionResult(ContainerExecutionResult.ExitStatus.SUCCESS, null,
          "Asked to die by the AM");
    }
    if (result.getError() != null) {
      Throwable e = result.getError();
      return new ContainerExecutionResult(
          ContainerExecutionResult.ExitStatus.EXECUTION_FAILURE,
          e, "TaskExecutionFailure: " + e.getMessage());
    }
  } finally {
    FileSystem.closeAllForUGI(taskUgi);
  }
  LOG.info("ExecutionTime for Container: " + request.getContainerIdString() + "=" +
      sw.stop().now(TimeUnit.MILLISECONDS));
  return new ContainerExecutionResult(ContainerExecutionResult.ExitStatus.SUCCESS, null,
      null);
}
 

/**
 * Create UserGroupInformation for specified user and credentials.
 *
 * @param user User.
 * @param credentialsBytes Credentials byte array.
 */
public static UserGroupInformation createUGI(String user, byte[] credentialsBytes) throws IOException {
    Credentials credentials = new Credentials();

    HadoopUtils.deserialize(credentials, credentialsBytes);

    UserGroupInformation ugi = UserGroupInformation.createRemoteUser(user);

    ugi.addCredentials(credentials);

    if (credentials.numberOfTokens() > 0)
        ugi.setAuthenticationMethod(UserGroupInformation.AuthenticationMethod.TOKEN);

    return ugi;
}