Source code

001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *   http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing,
013 * software distributed under the License is distributed on an
014 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
015 * KIND, either express or implied.  See the License for the
016 * specific language governing permissions and limitations
017 * under the License.
018 */
019package org.apache.reef.runtime.yarn.client;
020
021import org.apache.hadoop.io.DataInputBuffer;
022import org.apache.hadoop.io.DataOutputBuffer;
023import org.apache.hadoop.security.Credentials;
024import org.apache.hadoop.security.UserGroupInformation;
025import org.apache.hadoop.security.token.Token;
026import org.apache.hadoop.yarn.security.AMRMTokenIdentifier;
027
028import javax.inject.Inject;
029import java.io.IOException;
030import java.util.logging.Level;
031import java.util.logging.Logger;
032
033/**
034 * Reads security token from user credentials.
035 */
036public final class UserCredentialSecurityTokenProvider implements SecurityTokenProvider {
037
038  private static final Logger LOG = Logger.getLogger(UserCredentialSecurityTokenProvider.class.getName());
039
040  @Inject
041  private UserCredentialSecurityTokenProvider() { }
042
043  @Override
044  public byte[] getTokens() {
045
046    try {
047
048      final UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
049      final Credentials credentials = ugi.getCredentials();
050
051      LOG.log(Level.FINEST, "Got {0} tokens for user {1}", new Object[] {credentials.numberOfTokens(), ugi});
052
053      if (credentials.numberOfTokens() > 0) {
054        try (final DataOutputBuffer dob = new DataOutputBuffer()) {
055          credentials.writeTokenStorageToStream(dob);
056          return dob.getData();
057        }
058      }
059    } catch (final IOException e) {
060      LOG.log(Level.WARNING, "Could not access tokens in user credentials.", e);
061    }
062
063    LOG.log(Level.FINE, "No security token found.");
064
065    return null;
066  }
067
068  /**
069   * Add serialized token to teh credentials.
070   * @param tokens ByteBuffer containing token.
071   */
072  @Override
073  public void addTokens(final byte[] tokens) {
074
075    try (final DataInputBuffer buf = new DataInputBuffer()) {
076
077      buf.reset(tokens, tokens.length);
078      final Credentials credentials = new Credentials();
079      credentials.readTokenStorageStream(buf);
080
081      final UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
082      ugi.addCredentials(credentials);
083      LOG.log(Level.FINEST, "Added {0} tokens for user {1}", new Object[] {credentials.numberOfTokens(), ugi});
084
085    } catch (final IOException ex) {
086      LOG.log(Level.SEVERE, "Could not access tokens in user credentials.", ex);
087      throw new RuntimeException(ex);
088    }
089  }
090
091  /**
092   * Helper method to serialize a security token.
093   * @param token AM security token.
094   * @return ByteBuffer that contains the token. It is compatible with addTokens() method.
095   */
096  public static byte[] serializeToken(final Token<AMRMTokenIdentifier> token) {
097    try (final DataOutputBuffer dob = new DataOutputBuffer()) {
098      final Credentials credentials = new Credentials();
099      credentials.addToken(token.getService(), token);
100      credentials.writeTokenStorageToStream(dob);
101      return dob.getData();
102    } catch (final IOException ex) {
103      LOG.log(Level.SEVERE, "Could not write credentials to the buffer.", ex);
104      throw new RuntimeException(ex);
105    }
106  }
107}