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, software
013    * distributed under the License is distributed on an "AS IS" BASIS,
014    * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015    * See the License for the specific language governing permissions and
016    * limitations under the License.
017    */
018    
019    package org.apache.hadoop.yarn.security;
020    
021    import java.io.DataInput;
022    import java.io.DataInputStream;
023    import java.io.DataOutput;
024    import java.io.IOException;
025    
026    import org.apache.hadoop.classification.InterfaceAudience;
027    import org.apache.hadoop.classification.InterfaceAudience.Private;
028    import org.apache.hadoop.classification.InterfaceAudience.Public;
029    import org.apache.hadoop.classification.InterfaceStability.Evolving;
030    import org.apache.hadoop.io.Text;
031    import org.apache.hadoop.security.UserGroupInformation;
032    import org.apache.hadoop.security.token.Token;
033    import org.apache.hadoop.security.token.TokenIdentifier;
034    import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
035    import org.apache.hadoop.yarn.api.records.impl.pb.ApplicationAttemptIdPBImpl;
036    import org.apache.hadoop.yarn.proto.YarnSecurityTokenProtos.AMRMTokenIdentifierProto;
037    
038    import com.google.protobuf.TextFormat;
039    
040    /**
041     * AMRMTokenIdentifier is the TokenIdentifier to be used by
042     * ApplicationMasters to authenticate to the ResourceManager.
043     */
044    @Public
045    @Evolving
046    public class AMRMTokenIdentifier extends TokenIdentifier {
047    
048      public static final Text KIND_NAME = new Text("YARN_AM_RM_TOKEN");
049      private AMRMTokenIdentifierProto proto;
050    
051      public AMRMTokenIdentifier() {
052      }
053      
054      public AMRMTokenIdentifier(ApplicationAttemptId appAttemptId,
055          int masterKeyId) {
056        AMRMTokenIdentifierProto.Builder builder = 
057            AMRMTokenIdentifierProto.newBuilder();
058        if (appAttemptId != null) {
059          builder.setAppAttemptId(
060              ((ApplicationAttemptIdPBImpl)appAttemptId).getProto());
061        }
062        builder.setKeyId(masterKeyId);
063        proto = builder.build();
064      }
065    
066      @Private
067      public ApplicationAttemptId getApplicationAttemptId() {
068        if (!proto.hasAppAttemptId()) {
069          return null;
070        }
071        return new ApplicationAttemptIdPBImpl(proto.getAppAttemptId());
072      }
073    
074      @Override
075      public void write(DataOutput out) throws IOException {
076        out.write(proto.toByteArray());
077      }
078    
079      @Override
080      public void readFields(DataInput in) throws IOException {
081        proto = AMRMTokenIdentifierProto.parseFrom((DataInputStream)in);
082      }
083    
084      @Override
085      public Text getKind() {
086        return KIND_NAME;
087      }
088    
089      @Override
090      public UserGroupInformation getUser() {
091        String appAttemptId = null;
092        if (proto.hasAppAttemptId()) {
093          appAttemptId = 
094              new ApplicationAttemptIdPBImpl(proto.getAppAttemptId()).toString();
095        }
096        return UserGroupInformation.createRemoteUser(appAttemptId);
097      }
098    
099      public int getKeyId() {
100        return proto.getKeyId();
101      }
102      
103      public AMRMTokenIdentifierProto getProto() {
104        return this.proto;
105      }
106    
107      // TODO: Needed?
108      @InterfaceAudience.Private
109      public static class Renewer extends Token.TrivialRenewer {
110        @Override
111        protected Text getKind() {
112          return KIND_NAME;
113        }
114      }
115      
116      @Override
117      public int hashCode() {
118        return getProto().hashCode();
119      }
120    
121      @Override
122      public boolean equals(Object other) {
123        if (other == null)
124          return false;
125        if (other.getClass().isAssignableFrom(this.getClass())) {
126          return this.getProto().equals(this.getClass().cast(other).getProto());
127        }
128        return false;
129      }
130    
131      @Override
132      public String toString() {
133        return TextFormat.shortDebugString(getProto());
134      }
135    }