Java AMRMTokenIdentifier类代码示例

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
public Token<AMRMTokenIdentifier> createAndGetAMRMToken(
    ApplicationAttemptId appAttemptId) {
  this.writeLock.lock();
  try {
    LOG.info("Create AMRMToken for ApplicationAttempt: " + appAttemptId);
    AMRMTokenIdentifier identifier =
        new AMRMTokenIdentifier(appAttemptId, getMasterKey().getMasterKey()
          .getKeyId());
    byte[] password = this.createPassword(identifier);
    appAttemptSet.add(appAttemptId);
    return new Token<AMRMTokenIdentifier>(identifier.getBytes(), password,
      identifier.getKind(), new Text());
  } finally {
    this.writeLock.unlock();
  }
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
public MockAM sendAMLaunched(ApplicationAttemptId appAttemptId)
    throws Exception {
  MockAM am = new MockAM(getRMContext(), masterService, appAttemptId);
  am.waitForState(RMAppAttemptState.ALLOCATED);
  //create and set AMRMToken
  Token<AMRMTokenIdentifier> amrmToken =
      this.rmContext.getAMRMTokenSecretManager().createAndGetAMRMToken(
        appAttemptId);
  ((RMAppAttemptImpl) this.rmContext.getRMApps()
    .get(appAttemptId.getApplicationId()).getRMAppAttempt(appAttemptId))
    .setAMRMToken(amrmToken);
  getRMContext()
      .getDispatcher()
      .getEventHandler()
      .handle(
          new RMAppAttemptEvent(appAttemptId, RMAppAttemptEventType.LAUNCHED));
  return am;
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
public void unregisterAppAttempt(final FinishApplicationMasterRequest req,
    boolean waitForStateRunning) throws Exception {
  if (waitForStateRunning) {
    waitForState(RMAppAttemptState.RUNNING);
  }
  if (ugi == null) {
    ugi =  UserGroupInformation.createRemoteUser(attemptId.toString());
    Token<AMRMTokenIdentifier> token =
        context.getRMApps()
            .get(attemptId.getApplicationId())
            .getRMAppAttempt(attemptId).getAMRMToken();
    ugi.addTokenIdentifier(token.decodeIdentifier());
  }
  try {
    ugi.doAs(new PrivilegedExceptionAction<Object>() {
      @Override
      public Object run() throws Exception {
        amRMProtocol.finishApplicationMaster(req);
        return null;
      }
    });
  } catch (UndeclaredThrowableException e) {
    throw (Exception) e.getCause();
  }
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
protected ContainerId storeAttempt(RMStateStore store,
    ApplicationAttemptId attemptId,
    String containerIdStr, Token<AMRMTokenIdentifier> appToken,
    SecretKey clientTokenMasterKey, TestDispatcher dispatcher)
    throws Exception {

  RMAppAttemptMetrics mockRmAppAttemptMetrics = 
      mock(RMAppAttemptMetrics.class);
  Container container = new ContainerPBImpl();
  container.setId(ConverterUtils.toContainerId(containerIdStr));
  RMAppAttempt mockAttempt = mock(RMAppAttempt.class);
  when(mockAttempt.getAppAttemptId()).thenReturn(attemptId);
  when(mockAttempt.getMasterContainer()).thenReturn(container);
  when(mockAttempt.getAMRMToken()).thenReturn(appToken);
  when(mockAttempt.getClientTokenMasterKey())
      .thenReturn(clientTokenMasterKey);
  when(mockAttempt.getRMAppAttemptMetrics())
      .thenReturn(mockRmAppAttemptMetrics);
  when(mockRmAppAttemptMetrics.getAggregateAppResourceUsage())
      .thenReturn(new AggregateAppResourceUsage(0, 0, 0));
  dispatcher.attemptId = attemptId;
  store.storeNewApplicationAttempt(mockAttempt);
  waitNotify(dispatcher);
  return container.getId();
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
@Before
public void initialize() throws Exception {
  startHACluster(0, false, false, true);
  attemptId = this.cluster.createFakeApplicationAttemptId();
  amClient = ClientRMProxy
      .createRMProxy(this.conf, ApplicationMasterProtocol.class);

  Token<AMRMTokenIdentifier> appToken =
      this.cluster.getResourceManager().getRMContext()
        .getAMRMTokenSecretManager().createAndGetAMRMToken(attemptId);
  appToken.setService(ClientRMProxy.getAMRMTokenService(conf));
  UserGroupInformation.setLoginUser(UserGroupInformation
      .createRemoteUser(UserGroupInformation.getCurrentUser()
          .getUserName()));
  UserGroupInformation.getCurrentUser().addToken(appToken);
  syncToken(appToken);
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
private RequestInterceptorChainWrapper getInterceptorChain(
    AMRMTokenIdentifier tokenIdentifier) throws YarnException {
  ApplicationAttemptId appAttemptId =
      tokenIdentifier.getApplicationAttemptId();

  synchronized (this.applPipelineMap) {
    if (!this.applPipelineMap.containsKey(appAttemptId
        .getApplicationId())) {
      throw new YarnException(
          "The AM request processing pipeline is not initialized for app: "
              + appAttemptId.getApplicationId().toString());
    }

    return this.applPipelineMap.get(appAttemptId.getApplicationId());
  }
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
public Token<AMRMTokenIdentifier> createAndGetAMRMToken(
    ApplicationAttemptId appAttemptId) {
  this.writeLock.lock();
  try {
    LOG.info("Create AMRMToken for ApplicationAttempt: " + appAttemptId);
    AMRMTokenIdentifier identifier =
        new AMRMTokenIdentifier(appAttemptId, getMasterKey()
            .getMasterKey().getKeyId());
    byte[] password = this.createPassword(identifier);
    appAttemptSet.add(appAttemptId);
    return new Token<AMRMTokenIdentifier>(identifier.getBytes(),
        password, identifier.getKind(), new Text());
  } finally {
    this.writeLock.unlock();
  }
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
protected ContainerId storeAttempt(RMStateStore store,
    ApplicationAttemptId attemptId,
    String containerIdStr, Token<AMRMTokenIdentifier> appToken,
    SecretKey clientTokenMasterKey, TestDispatcher dispatcher)
    throws Exception {

  RMAppAttemptMetrics mockRmAppAttemptMetrics = 
      mock(RMAppAttemptMetrics.class);
  Container container = new ContainerPBImpl();
  container.setId(ConverterUtils.toContainerId(containerIdStr));
  RMAppAttempt mockAttempt = mock(RMAppAttempt.class);
  when(mockAttempt.getAppAttemptId()).thenReturn(attemptId);
  when(mockAttempt.getMasterContainer()).thenReturn(container);
  when(mockAttempt.getAMRMToken()).thenReturn(appToken);
  when(mockAttempt.getClientTokenMasterKey())
      .thenReturn(clientTokenMasterKey);
  when(mockAttempt.getRMAppAttemptMetrics())
      .thenReturn(mockRmAppAttemptMetrics);
  when(mockRmAppAttemptMetrics.getAggregateAppResourceUsage())
      .thenReturn(new AggregateAppResourceUsage(0, 0));
  dispatcher.attemptId = attemptId;
  store.storeNewApplicationAttempt(mockAttempt);
  waitNotify(dispatcher);
  return container.getId();
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
@Before
public void initialize() throws Exception {
  startHACluster(0, false, false, true);
  attemptId = this.cluster.createFakeApplicationAttemptId();

  Token<AMRMTokenIdentifier> appToken =
      this.cluster.getResourceManager().getRMContext()
        .getAMRMTokenSecretManager().createAndGetAMRMToken(attemptId);
  appToken.setService(ClientRMProxy.getAMRMTokenService(this.conf));
  UserGroupInformation.setLoginUser(UserGroupInformation
      .createRemoteUser(UserGroupInformation.getCurrentUser().getUserName()));
  UserGroupInformation.getCurrentUser().addToken(appToken);
  syncToken(appToken);

  amClient = ClientRMProxy
      .createRMProxy(this.conf, ApplicationMasterProtocol.class);
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
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);
  }
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
public AllocateResponse allocate(AllocateRequest allocateRequest)
          throws Exception {
  final AllocateRequest req = allocateRequest;
  req.setResponseId(++responseId);

  UserGroupInformation ugi =
      UserGroupInformation.createRemoteUser(attemptId.toString());
  Token<AMRMTokenIdentifier> token =
      context.getRMApps().get(attemptId.getApplicationId())
          .getRMAppAttempt(attemptId).getAMRMToken();
  ugi.addTokenIdentifier(token.decodeIdentifier());
  try {
    return ugi.doAs(new PrivilegedExceptionAction<AllocateResponse>() {
      @Override
      public AllocateResponse run() throws Exception {
        return amRMProtocol.allocate(req);
      }
    });
  } catch (UndeclaredThrowableException e) {
    throw (Exception) e.getCause();
  }
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
private ByteBuffer getSecurityTokens() throws IOException {
  Credentials credentials = UserGroupInformation.getCurrentUser().getCredentials();
  Closer closer = Closer.create();
  try {
    DataOutputBuffer dataOutputBuffer = closer.register(new DataOutputBuffer());
    credentials.writeTokenStorageToStream(dataOutputBuffer);

    // Remove the AM->RM token so that containers cannot access it
    Iterator<Token<?>> tokenIterator = credentials.getAllTokens().iterator();
    while (tokenIterator.hasNext()) {
      Token<?> token = tokenIterator.next();
      if (token.getKind().equals(AMRMTokenIdentifier.KIND_NAME)) {
        tokenIterator.remove();
      }
    }

    return ByteBuffer.wrap(dataOutputBuffer.getData(), 0, dataOutputBuffer.getLength());
  } catch (Throwable t) {
    throw closer.rethrow(t);
  } finally {
    closer.close();
  }
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
/**
 * Retrieve the password for the given {@link AMRMTokenIdentifier}.
 * Used by RPC layer to validate a remote {@link AMRMTokenIdentifier}.
 */
@Override
public synchronized byte[] retrievePassword(
    AMRMTokenIdentifier identifier) throws InvalidToken {
  ApplicationAttemptId applicationAttemptId =
      identifier.getApplicationAttemptId();
  if (LOG.isDebugEnabled()) {
    LOG.debug("Trying to retrieve password for " + applicationAttemptId);
  }
  byte[] password = this.passwords.get(applicationAttemptId);
  if (password == null) {
    throw new InvalidToken("Password not found for ApplicationAttempt "
        + applicationAttemptId);
  }
  return password;
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
private void recoverAppAttemptTokens(Credentials appAttemptTokens) {
  if (appAttemptTokens == null) {
    return;
  }
  if (UserGroupInformation.isSecurityEnabled()) {

    ClientToAMTokenSelector clientToAMTokenSelector =
        new ClientToAMTokenSelector();
    this.clientToAMToken =
        clientToAMTokenSelector.selectToken(new Text(),
          appAttemptTokens.getAllTokens());
  }

  // Only one AMRMToken is stored per-attempt, so this should be fine. Can't
  // use TokenSelector as service may change - think fail-over.
  this.amrmToken =
      (Token<AMRMTokenIdentifier>) appAttemptTokens
        .getToken(RMStateStore.AM_RM_TOKEN_SERVICE);

  // For now, no need to populate tokens back to AMRMTokenSecretManager,
  // because running attempts are rebooted. Later in work-preserve restart,
  // we'll create NEW->RUNNING transition in which the restored tokens will be
  // added to the secret manager
}
 

import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; //导入依赖的package包/类
public AllocateResponse allocate(
    List<ResourceRequest> resourceRequest, List<ContainerId> releases)
    throws Exception {
  final AllocateRequest req =
      AllocateRequest.newInstance(++responseId, 0F, resourceRequest,
        releases, null);
  UserGroupInformation ugi =
      UserGroupInformation.createRemoteUser(attemptId.toString());
  Token<AMRMTokenIdentifier> token =
      context.getRMApps().get(attemptId.getApplicationId())
        .getRMAppAttempt(attemptId).getAMRMToken();
  ugi.addTokenIdentifier(token.decodeIdentifier());
  try {
    return ugi.doAs(new PrivilegedExceptionAction<AllocateResponse>() {
      @Override
      public AllocateResponse run() throws Exception {
        return amRMProtocol.allocate(req);
      }
    });
  } catch (UndeclaredThrowableException e) {
    throw (Exception) e.getCause();
  }
}