Repository: reef Updated Branches: refs/heads/master 0fb3f8d9a -> 1396fb3dc
[REEF-1705] implement SecurityTokenProvider.addToken() * Added the `.addToken()` method to the `SecurityTokenProvider` interface and its implementations * Implemented token serialization for YARN JIRA: [REEF-1705](https://issues.apache.org/jira/browse/REEF-1705) Closes PR #1219 Project: http://git-wip-us.apache.org/repos/asf/reef/repo Commit: http://git-wip-us.apache.org/repos/asf/reef/commit/1396fb3d Tree: http://git-wip-us.apache.org/repos/asf/reef/tree/1396fb3d Diff: http://git-wip-us.apache.org/repos/asf/reef/diff/1396fb3d Branch: refs/heads/master Commit: 1396fb3dc7a4b9c739e245d260320eb0d3096357 Parents: 0fb3f8d Author: Sergiy Matusevych <[email protected]> Authored: Wed Jan 11 18:12:24 2017 -0800 Committer: Julia Wang <[email protected]> Committed: Fri Jan 13 14:20:15 2017 -0800 ---------------------------------------------------------------------- .../yarn/client/SecurityTokenProvider.java | 6 +++ .../UserCredentialSecurityTokenProvider.java | 57 ++++++++++++++++++-- 2 files changed, 59 insertions(+), 4 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/reef/blob/1396fb3d/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/SecurityTokenProvider.java ---------------------------------------------------------------------- diff --git a/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/SecurityTokenProvider.java b/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/SecurityTokenProvider.java index 37cfa76..f91cedf 100644 --- a/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/SecurityTokenProvider.java +++ b/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/SecurityTokenProvider.java @@ -31,4 +31,10 @@ public interface SecurityTokenProvider { * @return a ByteBuffer */ byte[] getTokens(); + + /** + * Add serialized tokens to the credentials. + * @param tokens ByteBuffer containing tokens. + */ + void addTokens(final byte[] tokens); } http://git-wip-us.apache.org/repos/asf/reef/blob/1396fb3d/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/UserCredentialSecurityTokenProvider.java ---------------------------------------------------------------------- diff --git a/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/UserCredentialSecurityTokenProvider.java b/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/UserCredentialSecurityTokenProvider.java index 5fde825..6735550 100644 --- a/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/UserCredentialSecurityTokenProvider.java +++ b/lang/java/reef-runtime-yarn/src/main/java/org/apache/reef/runtime/yarn/client/UserCredentialSecurityTokenProvider.java @@ -18,9 +18,12 @@ */ package org.apache.reef.runtime.yarn.client; +import org.apache.hadoop.io.DataInputBuffer; import org.apache.hadoop.io.DataOutputBuffer; import org.apache.hadoop.security.Credentials; import org.apache.hadoop.security.UserGroupInformation; +import org.apache.hadoop.security.token.Token; +import org.apache.hadoop.yarn.security.AMRMTokenIdentifier; import javax.inject.Inject; import java.io.IOException; @@ -30,29 +33,75 @@ import java.util.logging.Logger; /** * Reads security token from user credentials. */ -final class UserCredentialSecurityTokenProvider implements SecurityTokenProvider { +public final class UserCredentialSecurityTokenProvider implements SecurityTokenProvider { private static final Logger LOG = Logger.getLogger(UserCredentialSecurityTokenProvider.class.getName()); @Inject - private UserCredentialSecurityTokenProvider(){} + private UserCredentialSecurityTokenProvider() { } @Override public byte[] getTokens() { + try { + final UserGroupInformation ugi = UserGroupInformation.getCurrentUser(); final Credentials credentials = ugi.getCredentials(); + + LOG.log(Level.FINEST, "Got {0} tokens for user {1}", new Object[] {credentials.numberOfTokens(), ugi}); + if (credentials.numberOfTokens() > 0) { - try(final DataOutputBuffer dob = new DataOutputBuffer()) { + try (final DataOutputBuffer dob = new DataOutputBuffer()) { credentials.writeTokenStorageToStream(dob); return dob.getData(); } } - } catch (IOException e) { + } catch (final IOException e) { LOG.log(Level.WARNING, "Could not access tokens in user credentials.", e); } LOG.log(Level.FINE, "No security token found."); + return null; } + + /** + * Add serialized token to teh credentials. + * @param tokens ByteBuffer containing token. + */ + @Override + public void addTokens(final byte[] tokens) { + + try (final 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); + } + } + + /** + * Helper method to serialize a security token. + * @param token AM security token. + * @return ByteBuffer that contains the token. It is compatible with addTokens() method. + */ + public static byte[] serializeToken(final Token<AMRMTokenIdentifier> token) { + try (final DataOutputBuffer dob = new DataOutputBuffer()) { + final Credentials credentials = new Credentials(); + credentials.addToken(token.getService(), token); + credentials.writeTokenStorageToStream(dob); + return dob.getData(); + } catch (final IOException ex) { + LOG.log(Level.SEVERE, "Could not write credentials to the buffer.", ex); + throw new RuntimeException(ex); + } + } }
