Java Code Examples for org.apache.hadoop.yarn.util.ConverterUtils#convertFromYarn()

@Private
@VisibleForTesting
protected ContainerManagementProtocol newProxy(final YarnRPC rpc,
    String containerManagerBindAddr, ContainerId containerId, Token token)
    throws InvalidToken {

  if (token == null) {
    throw new InvalidToken("No NMToken sent for "
        + containerManagerBindAddr);
  }
  
  final InetSocketAddress cmAddr =
      NetUtils.createSocketAddr(containerManagerBindAddr);
  LOG.info("Opening proxy : " + containerManagerBindAddr);
  // the user in createRemoteUser in this context has to be ContainerID
  UserGroupInformation user =
      UserGroupInformation.createRemoteUser(containerId
          .getApplicationAttemptId().toString());

  org.apache.hadoop.security.token.Token<NMTokenIdentifier> nmToken =
      ConverterUtils.convertFromYarn(token, cmAddr);
  user.addToken(nmToken);

  return NMProxy.createNMProxy(conf, ContainerManagementProtocol.class,
    user, rpc, cmAddr);
}
 

public void addRMDelegationToken(final String renewer, final Credentials credentials) throws IOException, YarnException
{
  // Get the ResourceManager delegation rmToken
  final org.apache.hadoop.yarn.api.records.Token rmDelegationToken = clientRM.getRMDelegationToken(new Text(renewer));

  Token<RMDelegationTokenIdentifier> token;
  // TODO: Use the utility method getRMDelegationTokenService in ClientRMProxy to remove the separate handling of
  // TODO: HA and non-HA cases when hadoop dependency is changed to hadoop 2.4 or above
  if (ConfigUtils.isRMHAEnabled(conf)) {
    LOG.info("Yarn Resource Manager HA is enabled");
    token = getRMHAToken(rmDelegationToken);
  } else {
    LOG.info("Yarn Resource Manager HA is not enabled");
    InetSocketAddress rmAddress = conf.getSocketAddr(YarnConfiguration.RM_ADDRESS,
        YarnConfiguration.DEFAULT_RM_ADDRESS,
        YarnConfiguration.DEFAULT_RM_PORT);

    token = ConverterUtils.convertFromYarn(rmDelegationToken, rmAddress);
  }

  LOG.info("RM dt {}", token);

  credentials.addToken(token.getService(), token);
}
 

@Private
@VisibleForTesting
protected ContainerManagementProtocol newProxy(final YarnRPC rpc,
    String containerManagerBindAddr, ContainerId containerId, Token token)
    throws InvalidToken {

  if (token == null) {
    throw new InvalidToken("No NMToken sent for "
        + containerManagerBindAddr);
  }
  
  final InetSocketAddress cmAddr =
      NetUtils.createSocketAddr(containerManagerBindAddr);
  LOG.info("Opening proxy : " + containerManagerBindAddr);
  // the user in createRemoteUser in this context has to be ContainerID
  UserGroupInformation user =
      UserGroupInformation.createRemoteUser(containerId
          .getApplicationAttemptId().toString());

  org.apache.hadoop.security.token.Token<NMTokenIdentifier> nmToken =
      ConverterUtils.convertFromYarn(token, cmAddr);
  user.addToken(nmToken);

  return NMProxy.createNMProxy(conf, ContainerManagementProtocol.class,
    user, rpc, cmAddr);
}
 

public void addRMDelegationToken(final String renewer, final Credentials credentials) throws IOException, YarnException
{
  // Get the ResourceManager delegation rmToken
  final org.apache.hadoop.yarn.api.records.Token rmDelegationToken = clientRM.getRMDelegationToken(new Text(renewer));

  Token<RMDelegationTokenIdentifier> token;
  // TODO: Use the utility method getRMDelegationTokenService in ClientRMProxy to remove the separate handling of
  // TODO: HA and non-HA cases when hadoop dependency is changed to hadoop 2.4 or above
  if (ConfigUtils.isRMHAEnabled(conf)) {
    LOG.info("Yarn Resource Manager HA is enabled");
    token = getRMHAToken(rmDelegationToken);
  } else {
    LOG.info("Yarn Resource Manager HA is not enabled");
    InetSocketAddress rmAddress = conf.getSocketAddr(YarnConfiguration.RM_ADDRESS,
        YarnConfiguration.DEFAULT_RM_ADDRESS,
        YarnConfiguration.DEFAULT_RM_PORT);

    token = ConverterUtils.convertFromYarn(rmDelegationToken, rmAddress);
  }

  LOG.info("RM dt {}", token);

  credentials.addToken(token.getService(), token);
}
 

/**
 * Adds RM delegation token to the given {@link ContainerLaunchContext} so that the AM can authenticate itself
 * with RM using the delegation token.
 */
protected void addRMToken(ContainerLaunchContext context, YarnClient yarnClient, ApplicationId appId) {
  if (!UserGroupInformation.isSecurityEnabled()) {
    return;
  }

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

    Configuration config = yarnClient.getConfig();
    Token<TokenIdentifier> token = ConverterUtils.convertFromYarn(
      yarnClient.getRMDelegationToken(new Text(YarnUtils.getYarnTokenRenewer(config))),
      YarnUtils.getRMAddress(config));

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

    context.setTokens(YarnUtils.encodeCredentials(credentials));
  } catch (YarnException | IOException e) {
    throw new RuntimeException("Failed to acquire RM delegation token", e);
  }
}
 

@VisibleForTesting
Token<?> getDelegationTokenFromHS(MRClientProtocol hsProxy)
    throws IOException, InterruptedException {
  GetDelegationTokenRequest request = recordFactory
    .newRecordInstance(GetDelegationTokenRequest.class);
  request.setRenewer(Master.getMasterPrincipal(conf));
  org.apache.hadoop.yarn.api.records.Token mrDelegationToken;
  mrDelegationToken = hsProxy.getDelegationToken(request)
      .getDelegationToken();
  return ConverterUtils.convertFromYarn(mrDelegationToken,
      hsProxy.getConnectAddress());
}
 

private Token<?> getDelegationTokenFromHS(HSClientProtocol hsProxy)
    throws IOException, InterruptedException {
  GetDelegationTokenRequest request =
      recordFactory.newRecordInstance(GetDelegationTokenRequest.class);
  request.setRenewer(Master.getMasterPrincipal(conf));
  org.apache.hadoop.yarn.api.records.Token mrDelegationToken;
  mrDelegationToken =
      hsProxy.getDelegationToken(request).getDelegationToken();
  return ConverterUtils.convertFromYarn(mrDelegationToken,
      hsProxy.getConnectAddress());
}
 

@Test (timeout = 60000)
public void testAppSubmissionWithOldDelegationTokenAfterRMRestart()
    throws Exception {
  conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS, 2);
  conf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHENTICATION,
      "kerberos");
  conf.set(YarnConfiguration.RM_ADDRESS, "localhost:8032");
  UserGroupInformation.setConfiguration(conf);
  MemoryRMStateStore memStore = new MemoryRMStateStore();
  memStore.init(conf);

  MockRM rm1 = new TestSecurityMockRM(conf, memStore);
  rm1.start();

  GetDelegationTokenRequest request1 =
      GetDelegationTokenRequest.newInstance("renewer1");
  UserGroupInformation.getCurrentUser().setAuthenticationMethod(
      AuthMethod.KERBEROS);
  GetDelegationTokenResponse response1 =
      rm1.getClientRMService().getDelegationToken(request1);
  Token<RMDelegationTokenIdentifier> token1 =
      ConverterUtils.convertFromYarn(response1.getRMDelegationToken(), rmAddr);

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

  // submit an app with the old delegation token got from previous RM.
  Credentials ts = new Credentials();
  ts.addToken(token1.getService(), token1);
  RMApp app = rm2.submitApp(200, "name", "user",
      new HashMap<ApplicationAccessType, String>(), false, "default", 1, ts);
  rm2.waitForState(app.getApplicationId(), RMAppState.ACCEPTED);
}
 

@Private
public static DAGClientAMProtocolBlockingPB getAMProxy(final Configuration conf, String amHost,
    int amRpcPort, org.apache.hadoop.yarn.api.records.Token clientToAMToken) throws IOException {

  final InetSocketAddress serviceAddr = new InetSocketAddress(amHost, amRpcPort);
  UserGroupInformation userUgi = UserGroupInformation.createRemoteUser(UserGroupInformation
      .getCurrentUser().getUserName());
  if (clientToAMToken != null) {
    Token<ClientToAMTokenIdentifier> token = ConverterUtils.convertFromYarn(clientToAMToken,
        serviceAddr);
    userUgi.addToken(token);
  }
  if (LOG.isDebugEnabled()) {
    LOG.debug("Connecting to Tez AM at " + serviceAddr);
  }
  DAGClientAMProtocolBlockingPB proxy = null;
  try {
    proxy = userUgi.doAs(new PrivilegedExceptionAction<DAGClientAMProtocolBlockingPB>() {
      @Override
      public DAGClientAMProtocolBlockingPB run() throws IOException {
        RPC.setProtocolEngine(conf, DAGClientAMProtocolBlockingPB.class, ProtobufRpcEngine.class);
        return (DAGClientAMProtocolBlockingPB) RPC.getProxy(DAGClientAMProtocolBlockingPB.class,
            0, serviceAddr, conf);
      }
    });
  } catch (InterruptedException e) {
    throw new IOException("Failed to connect to AM", e);
  }
  return proxy;
}
 

protected ContainerManagementProtocol getCMProxy(ContainerId containerID,
                                                 final String containerManagerBindAddr,
                                                 Token containerToken)
    throws IOException {
  String [] hosts = containerManagerBindAddr.split(":");
  final InetSocketAddress cmAddr =
      new InetSocketAddress(hosts[0], Integer.parseInt(hosts[1]));
  UserGroupInformation user = UserGroupInformation.getCurrentUser();

  if (UserGroupInformation.isSecurityEnabled()) {
    org.apache.hadoop.security.token.Token<ContainerTokenIdentifier> token =
        ConverterUtils.convertFromYarn(containerToken, cmAddr);
    // the user in createRemoteUser in this context has to be ContainerID
    user = UserGroupInformation.createRemoteUser(containerID.toString());
    user.addToken(token);
  }

  ContainerManagementProtocol proxy = user.doAs(new PrivilegedAction<ContainerManagementProtocol>() {
    @Override
    public ContainerManagementProtocol run() {
      return (ContainerManagementProtocol) yarnRPC.getProxy(ContainerManagementProtocol.class,
          cmAddr, conf);
    }
  });

  return proxy;
}
 

@VisibleForTesting
Token<?> getDelegationTokenFromHS(MRClientProtocol hsProxy)
    throws IOException, InterruptedException {
  GetDelegationTokenRequest request = recordFactory
    .newRecordInstance(GetDelegationTokenRequest.class);
  request.setRenewer(Master.getMasterPrincipal(conf));
  org.apache.hadoop.yarn.api.records.Token mrDelegationToken;
  mrDelegationToken = hsProxy.getDelegationToken(request)
      .getDelegationToken();
  return ConverterUtils.convertFromYarn(mrDelegationToken,
      hsProxy.getConnectAddress());
}
 

@SuppressWarnings("rawtypes")
public Token getDelegationToken(Text renewer) throws IOException,
    InterruptedException {
  try {
    return ConverterUtils.convertFromYarn(
        client.getRMDelegationToken(renewer), getRMDelegationTokenService());
  } catch (YarnException e) {
    throw new IOException(e);
  }
}
 

@Test (timeout = 60000)
public void testAppSubmissionWithOldDelegationTokenAfterRMRestart()
    throws Exception {
  conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS, 2);
  conf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHENTICATION,
      "kerberos");
  conf.set(YarnConfiguration.RM_ADDRESS, "localhost:8032");
  UserGroupInformation.setConfiguration(conf);
  MemoryRMStateStore memStore = new MemoryRMStateStore();
  memStore.init(conf);

  MockRM rm1 = new TestSecurityMockRM(conf, memStore);
  rm1.start();

  GetDelegationTokenRequest request1 =
      GetDelegationTokenRequest.newInstance("renewer1");
  UserGroupInformation.getCurrentUser().setAuthenticationMethod(
      AuthMethod.KERBEROS);
  GetDelegationTokenResponse response1 =
      rm1.getClientRMService().getDelegationToken(request1);
  Token<RMDelegationTokenIdentifier> token1 =
      ConverterUtils.convertFromYarn(response1.getRMDelegationToken(), rmAddr);

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

  // submit an app with the old delegation token got from previous RM.
  Credentials ts = new Credentials();
  ts.addToken(token1.getService(), token1);
  RMApp app = rm2.submitApp(200, "name", "user",
      new HashMap<ApplicationAccessType, String>(), false, "default", 1, ts);
  rm2.waitForState(app.getApplicationId(), RMAppState.ACCEPTED);
}
 

@SuppressWarnings("rawtypes")
public Token getDelegationToken(Text renewer) throws IOException,
    InterruptedException {
  try {
    // Remove rmAddress after YARN-868 is addressed
    return ConverterUtils.convertFromYarn(
      client.getRMDelegationToken(renewer), rmAddress);
  } catch (YarnException e) {
    throw new IOException(e);
  }
}
 

@Private
public static DAGClientAMProtocolBlockingPB getAMProxy(final Configuration conf, String amHost,
    int amRpcPort, org.apache.hadoop.yarn.api.records.Token clientToAMToken,
    UserGroupInformation userUgi) throws IOException {

  final InetSocketAddress serviceAddr = NetUtils.createSocketAddrForHost(amHost, amRpcPort);
  if (clientToAMToken != null) {
    Token<ClientToAMTokenIdentifier> token = ConverterUtils.convertFromYarn(clientToAMToken,
        serviceAddr);
    userUgi.addToken(token);
  }
  if (LOG.isDebugEnabled()) {
    LOG.debug("Connecting to Tez AM at " + serviceAddr);
  }
  DAGClientAMProtocolBlockingPB proxy = null;
  try {
    proxy = userUgi.doAs(new PrivilegedExceptionAction<DAGClientAMProtocolBlockingPB>() {
      @Override
      public DAGClientAMProtocolBlockingPB run() throws IOException {
        RPC.setProtocolEngine(conf, DAGClientAMProtocolBlockingPB.class, ProtobufRpcEngine.class);
        return (DAGClientAMProtocolBlockingPB) RPC.getProxy(DAGClientAMProtocolBlockingPB.class,
            0, serviceAddr, conf);
      }
    });
  } catch (InterruptedException e) {
    throw new IOException("Failed to connect to AM", e);
  }
  return proxy;
}
 

private static Token<?> getDelegationTokenFromHS(HSClientProtocol hsProxy, Configuration conf) throws IOException {
  GetDelegationTokenRequest request =
      RecordFactoryProvider.getRecordFactory(null).newRecordInstance(GetDelegationTokenRequest.class);
  request.setRenewer(Master.getMasterPrincipal(conf));
  org.apache.hadoop.yarn.api.records.Token mrDelegationToken;
  mrDelegationToken = hsProxy.getDelegationToken(request).getDelegationToken();
  return ConverterUtils.convertFromYarn(mrDelegationToken, hsProxy.getConnectAddress());
}
 

private NMTokenIdentifier getNMTokenId(
    org.apache.hadoop.yarn.api.records.Token token) throws IOException {
  Token<NMTokenIdentifier> convertedToken =
      ConverterUtils.convertFromYarn(token, (Text) null);
  return convertedToken.decodeIdentifier();
}
 

@Test(timeout = 15000)
public void testRMDTMasterKeyStateOnRollingMasterKey() throws Exception {
  MemoryRMStateStore memStore = new MemoryRMStateStore();
  memStore.init(conf);
  RMState rmState = memStore.getState();

  Map<RMDelegationTokenIdentifier, Long> rmDTState =
      rmState.getRMDTSecretManagerState().getTokenState();
  Set<DelegationKey> rmDTMasterKeyState =
      rmState.getRMDTSecretManagerState().getMasterKeyState();

  MockRM rm1 = new MyMockRM(conf, memStore);
  rm1.start();
  // on rm start, two master keys are created.
  // One is created at RMDTSecretMgr.startThreads.updateCurrentKey();
  // the other is created on the first run of
  // tokenRemoverThread.rollMasterKey()

  RMDelegationTokenSecretManager dtSecretManager =
      rm1.getRMContext().getRMDelegationTokenSecretManager();
  // assert all master keys are saved
  Assert.assertEquals(dtSecretManager.getAllMasterKeys(), rmDTMasterKeyState);
  Set<DelegationKey> expiringKeys = new HashSet<DelegationKey>();
  expiringKeys.addAll(dtSecretManager.getAllMasterKeys());


  // request to generate a RMDelegationToken
  GetDelegationTokenRequest request = mock(GetDelegationTokenRequest.class);
  when(request.getRenewer()).thenReturn("renewer1");
  GetDelegationTokenResponse response =
      rm1.getClientRMService().getDelegationToken(request);
  org.apache.hadoop.yarn.api.records.Token delegationToken =
      response.getRMDelegationToken();
  Token<RMDelegationTokenIdentifier> token1 =
      ConverterUtils.convertFromYarn(delegationToken, (Text) null);
  RMDelegationTokenIdentifier dtId1 = token1.decodeIdentifier();

  // For all keys that still remain in memory, we should have them stored
  // in state-store also.
  while (((TestRMDelegationTokenSecretManager) dtSecretManager).numUpdatedKeys
    .get() < 3) {
    ((TestRMDelegationTokenSecretManager) dtSecretManager)
      .checkCurrentKeyInStateStore(rmDTMasterKeyState);
    Thread.sleep(100);
  }

  // wait for token to expire and remove from state-store
  // rollMasterKey is called every 1 second.
  int count = 0;
  while (rmDTState.containsKey(dtId1) && count < 100) {
    Thread.sleep(100);
    count++;
  }
  rm1.stop();
}
 

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

private NMTokenIdentifier getNMTokenId(
    org.apache.hadoop.yarn.api.records.Token token) throws IOException {
  Token<NMTokenIdentifier> convertedToken =
      ConverterUtils.convertFromYarn(token, (Text) null);
  return convertedToken.decodeIdentifier();
}