private void storeOrUpdateRMDT(
RMDelegationTokenIdentifier tokenId, Long renewDate, boolean isUpdate) throws IOException {
String tokenKey = getRMDTTokenNodeKey(tokenId);
RMDelegationTokenIdentifierData tokenData =
new RMDelegationTokenIdentifierData(tokenId, renewDate);
if (LOG.isDebugEnabled()) {
LOG.debug("Storing token to " + tokenKey);
}
try {
WriteBatch batch = db.createWriteBatch();
try {
batch.put(bytes(tokenKey), tokenData.toByteArray());
if (!isUpdate) {
ByteArrayOutputStream bs = new ByteArrayOutputStream();
try (DataOutputStream ds = new DataOutputStream(bs)) {
ds.writeInt(tokenId.getSequenceNumber());
}
if (LOG.isDebugEnabled()) {
LOG.debug(
"Storing " + tokenId.getSequenceNumber() + " to " + RM_DT_SEQUENCE_NUMBER_KEY);
}
batch.put(bytes(RM_DT_SEQUENCE_NUMBER_KEY), bs.toByteArray());
}
db.write(batch);
} finally {
batch.close();
}
} catch (DBException e) {
throw new IOException(e);
}
}
private void storeOrUpdateRMDelegationTokenState(
RMDelegationTokenIdentifier identifier, Long renewDate, boolean isUpdate) throws Exception {
Path nodeCreatePath =
getNodePath(
rmDTSecretManagerRoot, DELEGATION_TOKEN_PREFIX + identifier.getSequenceNumber());
RMDelegationTokenIdentifierData identifierData =
new RMDelegationTokenIdentifierData(identifier, renewDate);
if (isUpdate) {
LOG.info("Updating RMDelegationToken_" + identifier.getSequenceNumber());
updateFile(nodeCreatePath, identifierData.toByteArray(), true);
} else {
LOG.info("Storing RMDelegationToken_" + identifier.getSequenceNumber());
writeFileWithRetries(nodeCreatePath, identifierData.toByteArray(), true);
// store sequence number
Path latestSequenceNumberPath =
getNodePath(
rmDTSecretManagerRoot,
DELEGATION_TOKEN_SEQUENCE_NUMBER_PREFIX + identifier.getSequenceNumber());
LOG.info(
"Storing " + DELEGATION_TOKEN_SEQUENCE_NUMBER_PREFIX + identifier.getSequenceNumber());
if (dtSequenceNumberPath == null) {
if (!createFileWithRetries(latestSequenceNumberPath)) {
throw new Exception("Failed to create " + latestSequenceNumberPath);
}
} else {
if (!renameFileWithRetries(dtSequenceNumberPath, latestSequenceNumberPath)) {
throw new Exception("Failed to rename " + dtSequenceNumberPath);
}
}
dtSequenceNumberPath = latestSequenceNumberPath;
}
}
@Override
public synchronized void storeRMDelegationTokenAndSequenceNumberState(
RMDelegationTokenIdentifier identifier, Long renewDate, int latestSequenceNumber)
throws Exception {
Path nodeCreatePath =
getNodePath(
rmDTSecretManagerRoot, DELEGATION_TOKEN_PREFIX + identifier.getSequenceNumber());
ByteArrayOutputStream os = new ByteArrayOutputStream();
DataOutputStream fsOut = new DataOutputStream(os);
LOG.info("Storing RMDelegationToken_" + identifier.getSequenceNumber());
identifier.write(fsOut);
fsOut.writeLong(renewDate);
writeFile(nodeCreatePath, os.toByteArray());
fsOut.close();
// store sequence number
Path latestSequenceNumberPath =
getNodePath(
rmDTSecretManagerRoot, DELEGATION_TOKEN_SEQUENCE_NUMBER_PREFIX + latestSequenceNumber);
LOG.info("Storing " + DELEGATION_TOKEN_SEQUENCE_NUMBER_PREFIX + latestSequenceNumber);
if (dtSequenceNumberPath == null) {
if (!createFile(latestSequenceNumberPath)) {
throw new Exception("Failed to create " + latestSequenceNumberPath);
}
} else {
if (!renameFile(dtSequenceNumberPath, latestSequenceNumberPath)) {
throw new Exception("Failed to rename " + dtSequenceNumberPath);
}
}
dtSequenceNumberPath = latestSequenceNumberPath;
}
private void loadRMDTSecretManagerState(RMState rmState) throws Exception {
FileStatus[] childNodes = fs.listStatus(rmDTSecretManagerRoot);
for (FileStatus childNodeStatus : childNodes) {
assert childNodeStatus.isFile();
String childNodeName = childNodeStatus.getPath().getName();
if (childNodeName.startsWith(DELEGATION_TOKEN_SEQUENCE_NUMBER_PREFIX)) {
rmState.rmSecretManagerState.dtSequenceNumber =
Integer.parseInt(childNodeName.split("_")[1]);
continue;
}
Path childNodePath = getNodePath(rmDTSecretManagerRoot, childNodeName);
byte[] childData = readFile(childNodePath, childNodeStatus.getLen());
ByteArrayInputStream is = new ByteArrayInputStream(childData);
DataInputStream fsIn = new DataInputStream(is);
if (childNodeName.startsWith(DELEGATION_KEY_PREFIX)) {
DelegationKey key = new DelegationKey();
key.readFields(fsIn);
rmState.rmSecretManagerState.masterKeyState.add(key);
} else if (childNodeName.startsWith(DELEGATION_TOKEN_PREFIX)) {
RMDelegationTokenIdentifier identifier = new RMDelegationTokenIdentifier();
identifier.readFields(fsIn);
long renewDate = fsIn.readLong();
rmState.rmSecretManagerState.delegationTokenState.put(identifier, renewDate);
} else {
LOG.warn("Unknown file for recovering RMDelegationTokenSecretManager");
}
fsIn.close();
}
}
@Override
public synchronized void removeRMDelegationTokenState(RMDelegationTokenIdentifier identifier)
throws Exception {
Path nodeCreatePath =
getNodePath(
rmDTSecretManagerRoot, DELEGATION_TOKEN_PREFIX + identifier.getSequenceNumber());
LOG.info("Removing RMDelegationToken_" + identifier.getSequenceNumber());
deleteFileWithRetries(nodeCreatePath);
}
private String getRMDTTokenNodeKey(RMDelegationTokenIdentifier tokenId) {
return RM_DT_TOKEN_KEY_PREFIX + tokenId.getSequenceNumber();
}
@Test(timeout = 20000)
public void testGetHSDelegationToken() throws Exception {
try {
Configuration conf = new Configuration();
// Setup mock service
InetSocketAddress mockRmAddress = new InetSocketAddress("localhost", 4444);
Text rmTokenSevice = SecurityUtil.buildTokenService(mockRmAddress);
InetSocketAddress mockHsAddress = new InetSocketAddress("localhost", 9200);
Text hsTokenSevice = SecurityUtil.buildTokenService(mockHsAddress);
// Setup mock rm token
RMDelegationTokenIdentifier tokenIdentifier =
new RMDelegationTokenIdentifier(new Text("owner"), new Text("renewer"), new Text("real"));
Token<RMDelegationTokenIdentifier> token =
new Token<RMDelegationTokenIdentifier>(
new byte[0], new byte[0], tokenIdentifier.getKind(), rmTokenSevice);
token.setKind(RMDelegationTokenIdentifier.KIND_NAME);
// Setup mock history token
DelegationToken historyToken =
BuilderUtils.newDelegationToken(
new byte[0], MRDelegationTokenIdentifier.KIND_NAME.toString(),
new byte[0], hsTokenSevice.toString());
GetDelegationTokenResponse getDtResponse =
Records.newRecord(GetDelegationTokenResponse.class);
getDtResponse.setDelegationToken(historyToken);
// mock services
MRClientProtocol mockHsProxy = mock(MRClientProtocol.class);
doReturn(mockHsAddress).when(mockHsProxy).getConnectAddress();
doReturn(getDtResponse)
.when(mockHsProxy)
.getDelegationToken(any(GetDelegationTokenRequest.class));
ResourceMgrDelegate rmDelegate = mock(ResourceMgrDelegate.class);
doReturn(mockRmAddress).when(rmDelegate).getConnectAddress();
ClientCache clientCache = mock(ClientCache.class);
doReturn(mockHsProxy).when(clientCache).getInitializedHSProxy();
Credentials creds = new Credentials();
YARNRunner yarnRunner = new YARNRunner(conf, rmDelegate, clientCache);
// No HS token if no RM token
yarnRunner.addHistoyToken(creds);
verify(mockHsProxy, times(0)).getDelegationToken(any(GetDelegationTokenRequest.class));
// No HS token if RM token, but secirity disabled.
creds.addToken(new Text("rmdt"), token);
yarnRunner.addHistoyToken(creds);
verify(mockHsProxy, times(0)).getDelegationToken(any(GetDelegationTokenRequest.class));
conf.set(CommonConfigurationKeys.HADOOP_SECURITY_AUTHENTICATION, "kerberos");
UserGroupInformation.setConfiguration(conf);
creds = new Credentials();
// No HS token if no RM token, security enabled
yarnRunner.addHistoyToken(creds);
verify(mockHsProxy, times(0)).getDelegationToken(any(GetDelegationTokenRequest.class));
// HS token if RM token present, security enabled
creds.addToken(new Text("rmdt"), token);
yarnRunner.addHistoyToken(creds);
verify(mockHsProxy, times(1)).getDelegationToken(any(GetDelegationTokenRequest.class));
// No additional call to get HS token if RM and HS token present
yarnRunner.addHistoyToken(creds);
verify(mockHsProxy, times(1)).getDelegationToken(any(GetDelegationTokenRequest.class));
} finally {
// Back to defaults.
UserGroupInformation.setConfiguration(new Configuration());
}
}