UnixUserGroupInformation.java
上传用户:quxuerui
上传日期:2018-01-08
资源大小:41811k
文件大小:14k
源码类别:

网格计算

开发平台:

Java

  1. /**
  2.  * Licensed to the Apache Software Foundation (ASF) under one
  3.  * or more contributor license agreements.  See the NOTICE file
  4.  * distributed with this work for additional information
  5.  * regarding copyright ownership.  The ASF licenses this file
  6.  * to you under the Apache License, Version 2.0 (the
  7.  * "License"); you may not use this file except in compliance
  8.  * with the License.  You may obtain a copy of the License at
  9.  *
  10.  *     http://www.apache.org/licenses/LICENSE-2.0
  11.  *
  12.  * Unless required by applicable law or agreed to in writing, software
  13.  * distributed under the License is distributed on an "AS IS" BASIS,
  14.  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  15.  * See the License for the specific language governing permissions and
  16.  * limitations under the License.
  17.  */
  18. package org.apache.hadoop.security;
  19. import java.io.DataInput;
  20. import java.io.DataOutput;
  21. import java.io.IOException;
  22. import java.util.Arrays;
  23. import java.util.HashMap;
  24. import java.util.StringTokenizer;
  25. import java.util.TreeSet;
  26. import javax.security.auth.login.LoginException;
  27. import org.apache.hadoop.conf.Configuration;
  28. import org.apache.hadoop.util.Shell;
  29. import org.apache.hadoop.io.Text;
  30. import org.apache.hadoop.io.WritableUtils;
  31. /** An implementation of UserGroupInformation in the Unix system */
  32. public class UnixUserGroupInformation extends UserGroupInformation {
  33.   public static final String DEFAULT_USERNAME = "DrWho";
  34.   public static final String DEFAULT_GROUP = "Tardis";
  35.   final static public String UGI_PROPERTY_NAME = "hadoop.job.ugi";
  36.   final static private HashMap<String, UnixUserGroupInformation> user2UGIMap =
  37.     new HashMap<String, UnixUserGroupInformation>();
  38.   /** Create an immutable {@link UnixUserGroupInformation} object. */
  39.   public static UnixUserGroupInformation createImmutable(String[] ugi) {
  40.     return new UnixUserGroupInformation(ugi) {
  41.       public void readFields(DataInput in) throws IOException {
  42.         throw new UnsupportedOperationException();
  43.       }
  44.     };
  45.   }
  46.   private String userName;
  47.   private String[] groupNames;
  48.   /** Default constructor
  49.    */
  50.   public UnixUserGroupInformation() {
  51.   }
  52.   /** Constructor with parameters user name and its group names.
  53.    * The first entry in the groups list is the default  group.
  54.    * 
  55.    * @param userName a user's name
  56.    * @param groupNames groups list, first of which is the default group
  57.    * @exception IllegalArgumentException if any argument is null
  58.    */
  59.   public UnixUserGroupInformation(String userName, String[] groupNames) {
  60.     setUserGroupNames(userName, groupNames);
  61.   }
  62.   /** Constructor with parameter user/group names
  63.    * 
  64.    * @param ugi an array containing user/group names, the first
  65.    *                     element of which is the user name, the second of
  66.    *                     which is the default group name.
  67.    * @exception IllegalArgumentException if the array size is less than 2 
  68.    *                                     or any element is null.
  69.    */
  70.   public UnixUserGroupInformation(String[] ugi) {
  71.     if (ugi==null || ugi.length < 2) {
  72.       throw new IllegalArgumentException( "Parameter does contain at least "+
  73.           "one user name and one group name");
  74.     }
  75.     String[] groupNames = new String[ugi.length-1];
  76.     System.arraycopy(ugi, 1, groupNames, 0, groupNames.length);
  77.     setUserGroupNames(ugi[0], groupNames);
  78.   }
  79.   
  80.   /* Set this object's user name and group names
  81.    * 
  82.    * @param userName a user's name
  83.    * @param groupNames groups list, the first of which is the default group
  84.    * @exception IllegalArgumentException if any argument is null
  85.    */
  86.   private void setUserGroupNames(String userName, String[] groupNames) {
  87.     if (userName==null || userName.length()==0 ||
  88.         groupNames== null || groupNames.length==0) {
  89.       throw new IllegalArgumentException(
  90.           "Parameters should not be null or an empty string/array");
  91.     }
  92.     for (int i=0; i<groupNames.length; i++) {
  93.       if(groupNames[i] == null || groupNames[i].length() == 0) {
  94.         throw new IllegalArgumentException("A null group name at index " + i);
  95.       }
  96.     }
  97.     this.userName = userName;
  98.     this.groupNames = groupNames;
  99.   }
  100.   /** Return an array of group names
  101.    */
  102.   public String[] getGroupNames() {
  103.     return groupNames;
  104.   }
  105.   /** Return the user's name
  106.    */
  107.   public String getUserName() {
  108.     return userName;
  109.   }
  110.   /* The following two methods implements Writable interface */
  111.   final private static String UGI_TECHNOLOGY = "STRING_UGI"; 
  112.   /** Deserialize this object
  113.    * First check if this is a UGI in the string format.
  114.    * If no, throw an IOException; otherwise
  115.    * set this object's fields by reading them from the given data input
  116.    *  
  117.    *  @param in input stream
  118.    *  @exception IOException is thrown if encounter any error when reading
  119.    */
  120.   public void readFields(DataInput in) throws IOException {
  121.     // read UGI type first
  122.     String ugiType = Text.readString(in);
  123.     if (!UGI_TECHNOLOGY.equals(ugiType)) {
  124.       throw new IOException("Expect UGI prefix: " + UGI_TECHNOLOGY +
  125.           ", but receive a prefix: " + ugiType);
  126.     }
  127.     
  128.     // read this object
  129.     userName = Text.readString(in);
  130.     int numOfGroups = WritableUtils.readVInt(in);
  131.     groupNames = new String[numOfGroups];
  132.     for (int i = 0; i < numOfGroups; i++) {
  133.       groupNames[i] = Text.readString(in);
  134.     }
  135.   }
  136.   /** Serialize this object
  137.    * First write a string marking that this is a UGI in the string format,
  138.    * then write this object's serialized form to the given data output
  139.    * 
  140.    * @param out output stream
  141.    * @exception IOException if encounter any error during writing
  142.    */
  143.   public void write(DataOutput out) throws IOException {
  144.     // write a prefix indicating the type of UGI being written
  145.     Text.writeString(out, UGI_TECHNOLOGY);
  146.     // write this object
  147.     Text.writeString(out, userName);
  148.     WritableUtils.writeVInt(out, groupNames.length);
  149.     for (String groupName : groupNames) {
  150.       Text.writeString(out, groupName);
  151.     }
  152.   }
  153.   /* The following two methods deal with transferring UGI through conf. 
  154.    * In this pass of implementation we store UGI as a string in conf. 
  155.    * Later we may change it to be a more general approach that stores 
  156.    * it as a byte array */
  157.   /** Store the given <code>ugi</code> as a comma separated string in
  158.    * <code>conf</code> as a property <code>attr</code>
  159.    * 
  160.    * The String starts with the user name followed by the default group names,
  161.    * and other group names.
  162.    * 
  163.    * @param conf configuration
  164.    * @param attr property name
  165.    * @param ugi a UnixUserGroupInformation
  166.    */
  167.   public static void saveToConf(Configuration conf, String attr, 
  168.       UnixUserGroupInformation ugi ) {
  169.     conf.set(attr, ugi.toString());
  170.   }
  171.   
  172.   /** Read a UGI from the given <code>conf</code>
  173.    * 
  174.    * The object is expected to store with the property name <code>attr</code>
  175.    * as a comma separated string that starts
  176.    * with the user name followed by group names.
  177.    * If the property name is not defined, return null.
  178.    * It's assumed that there is only one UGI per user. If this user already
  179.    * has a UGI in the ugi map, return the ugi in the map.
  180.    * Otherwise, construct a UGI from the configuration, store it in the
  181.    * ugi map and return it.
  182.    * 
  183.    * @param conf configuration
  184.    * @param attr property name
  185.    * @return a UnixUGI
  186.    * @throws LoginException if the stored string is ill-formatted.
  187.    */
  188.   public static UnixUserGroupInformation readFromConf(
  189.       Configuration conf, String attr) throws LoginException {
  190.     String[] ugi = conf.getStrings(attr);
  191.     if(ugi == null) {
  192.       return null;
  193.     }
  194.     UnixUserGroupInformation currentUGI = null;
  195.     if (ugi.length>0 ){
  196.       currentUGI = user2UGIMap.get(ugi[0]);
  197.     }
  198.     if (currentUGI == null) {
  199.       try {
  200.         currentUGI = new UnixUserGroupInformation(ugi);
  201.         user2UGIMap.put(currentUGI.getUserName(), currentUGI);
  202.       } catch (IllegalArgumentException e) {
  203.         throw new LoginException("Login failed: "+e.getMessage());
  204.       }
  205.     }
  206.     
  207.     return currentUGI;
  208.   }
  209.   
  210.   /**
  211.    * Get current user's name and the names of all its groups from Unix.
  212.    * It's assumed that there is only one UGI per user. If this user already
  213.    * has a UGI in the ugi map, return the ugi in the map.
  214.    * Otherwise get the current user's information from Unix, store it
  215.    * in the map, and return it.
  216.    *
  217.    * If the current user's UNIX username or groups are configured in such a way
  218.    * to throw an Exception, for example if the user uses LDAP, then this method
  219.    * will use a the {@link #DEFAULT_USERNAME} and {@link #DEFAULT_GROUP}
  220.    * constants.
  221.    */
  222.   public static UnixUserGroupInformation login() throws LoginException {
  223.     try {
  224.       String userName;
  225.       // if an exception occurs, then uses the
  226.       // default user
  227.       try {
  228.         userName =  getUnixUserName();
  229.       } catch (Exception e) {
  230.         userName = DEFAULT_USERNAME;
  231.       }
  232.       // check if this user already has a UGI object in the ugi map
  233.       UnixUserGroupInformation ugi = user2UGIMap.get(userName);
  234.       if (ugi != null) {
  235.         return ugi;
  236.       }
  237.       /* get groups list from UNIX. 
  238.        * It's assumed that the first group is the default group.
  239.        */
  240.       String[]  groupNames;
  241.       // if an exception occurs, then uses the
  242.       // default group
  243.       try {
  244.         groupNames = getUnixGroups();
  245.       } catch (Exception e) {
  246.         groupNames = new String[1];
  247.         groupNames[0] = DEFAULT_GROUP;
  248.       }
  249.       // construct a Unix UGI
  250.       ugi = new UnixUserGroupInformation(userName, groupNames);
  251.       user2UGIMap.put(ugi.getUserName(), ugi);
  252.       return ugi;
  253.     } catch (Exception e) {
  254.       throw new LoginException("Login failed: "+e.getMessage());
  255.     }
  256.   }
  257.   /** Equivalent to login(conf, false). */
  258.   public static UnixUserGroupInformation login(Configuration conf)
  259.     throws LoginException {
  260.     return login(conf, false);
  261.   }
  262.   
  263.   /** Get a user's name & its group names from the given configuration; 
  264.    * If it is not defined in the configuration, get the current user's
  265.    * information from Unix.
  266.    * If the user has a UGI in the ugi map, return the one in
  267.    * the UGI map.
  268.    * 
  269.    *  @param conf either a job configuration or client's configuration
  270.    *  @param save saving it to conf?
  271.    *  @return UnixUserGroupInformation a user/group information
  272.    *  @exception LoginException if not able to get the user/group information
  273.    */
  274.   public static UnixUserGroupInformation login(Configuration conf, boolean save
  275.       ) throws LoginException {
  276.     UnixUserGroupInformation ugi = readFromConf(conf, UGI_PROPERTY_NAME);
  277.     if (ugi == null) {
  278.       ugi = login();
  279.       LOG.debug("Unix Login: " + ugi);
  280.       if (save) {
  281.         saveToConf(conf, UGI_PROPERTY_NAME, ugi);
  282.       }
  283.     }
  284.     return ugi;
  285.   }
  286.   
  287.   /* Return a string representation of a string array.
  288.    * Two strings are separated by a blank.
  289.    */
  290.   private static String toString(String[] strArray) {
  291.     if (strArray==null || strArray.length==0) {
  292.       return "";
  293.     }
  294.     StringBuilder buf = new StringBuilder(strArray[0]);
  295.     for (int i=1; i<strArray.length; i++) {
  296.       buf.append(' ');
  297.       buf.append(strArray[i]);
  298.     }
  299.     return buf.toString();
  300.   }
  301.   
  302.   /** Get current user's name from Unix by running the command whoami.
  303.    * 
  304.    * @return current user's name
  305.    * @throws IOException if encounter any error while running the command
  306.    */
  307.   static String getUnixUserName() throws IOException {
  308.     String[] result = executeShellCommand(
  309.         new String[]{Shell.USER_NAME_COMMAND});
  310.     if (result.length!=1) {
  311.       throw new IOException("Expect one token as the result of " + 
  312.           Shell.USER_NAME_COMMAND + ": " + toString(result));
  313.     }
  314.     return result[0];
  315.   }
  316.   /** Get the current user's group list from Unix by running the command groups
  317.    * 
  318.    * @return the groups list that the current user belongs to
  319.    * @throws IOException if encounter any error when running the command
  320.    */
  321.   private static String[] getUnixGroups() throws IOException {
  322.     return executeShellCommand(Shell.getGROUPS_COMMAND());
  323.   }
  324.   
  325.   /* Execute a command and return the result as an array of Strings */
  326.   private static String[] executeShellCommand(String[] command)
  327.   throws IOException {
  328.     String groups = Shell.execCommand(command);
  329.     StringTokenizer tokenizer = new StringTokenizer(groups);
  330.     int numOfTokens = tokenizer.countTokens();
  331.     String[] tokens = new String[numOfTokens];
  332.     for (int i=0; tokenizer.hasMoreTokens(); i++) {
  333.       tokens[i] = tokenizer.nextToken();
  334.     }
  335.     return tokens;
  336.   }
  337.   /** Decide if two UGIs are the same
  338.    *
  339.    * @param other other object
  340.    * @return true if they are the same; false otherwise.
  341.    */
  342.   public boolean equals(Object other) {
  343.     if (this == other) {
  344.       return true;
  345.     }
  346.     
  347.     if (!(other instanceof UnixUserGroupInformation)) {
  348.       return false;
  349.     }
  350.     
  351.     UnixUserGroupInformation otherUGI = (UnixUserGroupInformation)other;
  352.     
  353.     // check userName
  354.     if (userName == null) {
  355.       if (otherUGI.getUserName() != null) {
  356.         return false;
  357.       }
  358.     } else {
  359.       if (!userName.equals(otherUGI.getUserName())) {
  360.         return false;
  361.       }
  362.     }
  363.     
  364.     // checkGroupNames
  365.     if (groupNames == otherUGI.groupNames) {
  366.       return true;
  367.     }
  368.     if (groupNames.length != otherUGI.groupNames.length) {
  369.       return false;
  370.     }
  371.     // check default group name
  372.     if (groupNames.length>0 && !groupNames[0].equals(otherUGI.groupNames[0])) {
  373.       return false;
  374.     }
  375.     // check all group names, ignoring the order
  376.     return new TreeSet<String>(Arrays.asList(groupNames)).equals(
  377.            new TreeSet<String>(Arrays.asList(otherUGI.groupNames)));
  378.   }
  379.   /** Returns a hash code for this UGI. 
  380.    * The hash code for a UGI is the hash code of its user name string.
  381.    * 
  382.    * @return  a hash code value for this UGI.
  383.    */
  384.   public int hashCode() {
  385.     return getUserName().hashCode();
  386.   }
  387.   
  388.   /** Convert this object to a string
  389.    * 
  390.    * @return a comma separated string containing the user name and group names
  391.    */
  392.   public String toString() {
  393.     StringBuilder buf = new StringBuilder();
  394.     buf.append(userName);
  395.     for (String groupName : groupNames) {
  396.       buf.append(',');
  397.       buf.append(groupName);
  398.     }
  399.     return buf.toString();
  400.   }
  401.   @Override
  402.   public String getName() {
  403.     return toString();
  404.   }
  405. }