service/auth/TSetIpAddressProcessor.java

0001 /**
0002  * Licensed to the Apache Software Foundation (ASF) under one
0003  * or more contributor license agreements.  See the NOTICE file
0004  * distributed with this work for additional information
0005  * regarding copyright ownership.  The ASF licenses this file
0006  * to you under the Apache License, Version 2.0 (the
0007  * "License"); you may not use this file except in compliance
0008  * with the License.  You may obtain a copy of the License at
0009  *
0010  *     http://www.apache.org/licenses/LICENSE-2.0
0011  *
0012  * Unless required by applicable law or agreed to in writing, software
0013  * distributed under the License is distributed on an "AS IS" BASIS,
0014  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0015  * See the License for the specific language governing permissions and
0016  * limitations under the License.
0017  */
0018
0019 package org.apache.hive.service.auth;
0020
0021 import org.apache.hive.service.rpc.thrift.TCLIService;
0022 import org.apache.hive.service.rpc.thrift.TCLIService.Iface;
0023 import org.apache.thrift.TException;
0024 import org.apache.thrift.protocol.TProtocol;
0025 import org.apache.thrift.transport.TSaslClientTransport;
0026 import org.apache.thrift.transport.TSaslServerTransport;
0027 import org.apache.thrift.transport.TSocket;
0028 import org.apache.thrift.transport.TTransport;
0029 import org.slf4j.Logger;
0030 import org.slf4j.LoggerFactory;
0031
0032 /**
0033  * This class is responsible for setting the ipAddress for operations executed via HiveServer2.
0034  *
0035  * - IP address is only set for operations that calls listeners with hookContext
0036  * - IP address is only set if the underlying transport mechanism is socket
0037  *
0038  * @see org.apache.hadoop.hive.ql.hooks.ExecuteWithHookContext
0039  */
0040 public class TSetIpAddressProcessor<I extends Iface> extends TCLIService.Processor<Iface> {
0041
0042   private static final Logger LOGGER = LoggerFactory.getLogger(TSetIpAddressProcessor.class.getName());
0043
0044   public TSetIpAddressProcessor(Iface iface) {
0045     super(iface);
0046   }
0047
0048   @Override
0049   public boolean process(final TProtocol in, final TProtocol out) throws TException {
0050     setIpAddress(in);
0051     setUserName(in);
0052     try {
0053       return super.process(in, out);
0054     } finally {
0055       THREAD_LOCAL_USER_NAME.remove();
0056       THREAD_LOCAL_IP_ADDRESS.remove();
0057     }
0058   }
0059
0060   private void setUserName(final TProtocol in) {
0061     TTransport transport = in.getTransport();
0062     if (transport instanceof TSaslServerTransport) {
0063       String userName = ((TSaslServerTransport) transport).getSaslServer().getAuthorizationID();
0064       THREAD_LOCAL_USER_NAME.set(userName);
0065     }
0066   }
0067
0068   protected void setIpAddress(final TProtocol in) {
0069     TTransport transport = in.getTransport();
0070     TSocket tSocket = getUnderlyingSocketFromTransport(transport);
0071     if (tSocket == null) {
0072       LOGGER.warn("Unknown Transport, cannot determine ipAddress");
0073     } else {
0074       THREAD_LOCAL_IP_ADDRESS.set(tSocket.getSocket().getInetAddress().getHostAddress());
0075     }
0076   }
0077
0078   private TSocket getUnderlyingSocketFromTransport(TTransport transport) {
0079     while (transport != null) {
0080       if (transport instanceof TSaslServerTransport) {
0081         transport = ((TSaslServerTransport) transport).getUnderlyingTransport();
0082       }
0083       if (transport instanceof TSaslClientTransport) {
0084         transport = ((TSaslClientTransport) transport).getUnderlyingTransport();
0085       }
0086       if (transport instanceof TSocket) {
0087         return (TSocket) transport;
0088       }
0089     }
0090     return null;
0091   }
0092
0093   private static final ThreadLocal<String> THREAD_LOCAL_IP_ADDRESS = new ThreadLocal<String>() {
0094     @Override
0095     protected synchronized String initialValue() {
0096       return null;
0097     }
0098   };
0099
0100   private static final ThreadLocal<String> THREAD_LOCAL_USER_NAME = new ThreadLocal<String>() {
0101     @Override
0102     protected synchronized String initialValue() {
0103       return null;
0104     }
0105   };
0106
0107   public static String getUserIpAddress() {
0108     return THREAD_LOCAL_IP_ADDRESS.get();
0109   }
0110
0111   public static String getUserName() {
0112     return THREAD_LOCAL_USER_NAME.get();
0113   }
0114 }