content: /**
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   * http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package backtype.storm.security.auth;
  
  import java.io.IOException;
  import java.security.InvalidParameterException;
  import java.util.List;
  import java.util.Map;
  
  import javax.security.auth.login.Configuration;
  
  import org.apache.commons.lang.StringUtils;
  import org.apache.curator.framework.CuratorFramework;
  import org.apache.thrift.protocol.TBinaryProtocol;
  import org.apache.thrift.protocol.TProtocol;
  import org.apache.thrift.transport.TSocket;
  import org.apache.thrift.transport.TTransport;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import com.alibaba.jstorm.cluster.Cluster;
  
  import backtype.storm.Config;
  import backtype.storm.utils.Utils;
  
  public class ThriftClient {
      private static final Logger LOG = LoggerFactory.getLogger(ThriftClient.class);
      private TTransport _transport;
      protected TProtocol _protocol;
      private String hostPort;
      private String host;
      private Integer port;
  
      private Map<Object, Object> conf;
  
      private Integer timeout;
      private ThriftConnectionType type;
      private String asUser;
  
      public ThriftClient(Map conf, ThriftConnectionType type) throws Exception {
          this(conf, type, null, null, null, null);
      }
  
      @SuppressWarnings("unchecked")
      public ThriftClient(Map conf, ThriftConnectionType type, Integer timeout) throws Exception {
          this(conf, type, null, null, timeout, null);
      }
  
      /**
       * This is only for be compatible for Storm
       * 
       * @param conf
       * @param type
       * @param host
       */
      public ThriftClient(Map conf, ThriftConnectionType type, String host) {
          this(conf, type, host, null, null, null);
      }
  
      public ThriftClient(Map conf, ThriftConnectionType type, String host, Integer port, Integer timeout) {
          this(conf, type, host, port, timeout, null);
      }
  
      public ThriftClient(Map conf, ThriftConnectionType type, String host, Integer port, Integer timeout, String asUser) {
          // create a socket with server
  
          this.timeout = timeout;
          this.conf = conf;
          this.type = type;
          this.asUser = asUser;
  
          getMaster(conf, host, port);
          reconnect();
      }
  
      public static String getMasterByZk(Map conf) throws Exception {
  
          CuratorFramework zkobj = null;
          String masterHost = null;
  
          try {
              String root = String.valueOf(conf.get(Config.STORM_ZOOKEEPER_ROOT));
              String zkMasterDir = root + Cluster.MASTER_SUBTREE;
  
              zkobj = Utils.newCurator(conf, (List<String>) conf.get(Config.STORM_ZOOKEEPER_SERVERS), conf.get(Config.STORM_ZOOKEEPER_PORT), zkMasterDir);
              zkobj.start();
              if (zkobj.checkExists().forPath("/") == null) {
                  throw new RuntimeException("!!!!!!!!!!! \n\n\n No alive nimbus !!!!!!!!!!! \n\n\n");
              }
  
              masterHost = new String(zkobj.getData().forPath("/"));
  
              LOG.info("masterHost:" + masterHost);
              return masterHost;
          } finally {
              if (zkobj != null) {
                  zkobj.close();
                  zkobj = null;
              }
          }
      }
  
      public void getMaster(Map conf, String host, Integer port) {
          if (StringUtils.isBlank(host) == false) {
              this.host = host;
              if (port == null) {
                  port = type.getPort(conf);
              }
              this.port = port;
              this.hostPort = host + ":" + port;
          } else {
              try {
                  hostPort = ThriftClient.getMasterByZk(conf);
              } catch (Exception e) {
                  // TODO Auto-generated catch block
                  String errMsg = "Failed to get master.\n\n" + e.getMessage();
                  errMsg += "\n\n\nConfiguration is " + conf.toString();
                  throw new RuntimeException(errMsg, e);
              }
              String[] host_port = hostPort.split(":");
              if (host_port.length != 2) {
                  throw new InvalidParameterException("Host format error: " + hostPort);
              }
              this.host = host_port[0];
              this.port = Integer.parseInt(host_port[1]);
          }
  
          // create a socket with server
          if (this.host == null) {
              throw new IllegalArgumentException("host is not set");
          }
          if (this.port <= 0) {
              throw new IllegalArgumentException("invalid port: " + port);
          }
      }
  
      public synchronized TTransport transport() {
          return _transport;
      }
  
      public synchronized void reconnect() {
          close();
          try {
              TSocket socket = new TSocket(host, port);
              if (timeout != null) {
                  socket.setTimeout(timeout);
              } else {
                  // @@@ Todo
                  // set the socket default Timeout as xxxx
              }
  
              // locate login configuration
              Configuration login_conf = AuthUtils.GetConfiguration(conf);
  
              // construct a transport plugin
              ITransportPlugin transportPlugin = AuthUtils.GetTransportPlugin(type, conf, login_conf);
  
              final TTransport underlyingTransport = socket;
  
              // TODO get this from type instead of hardcoding to Nimbus.
              // establish client-server transport via plugin
              // do retries if the connect fails
              TBackoffConnect connectionRetry =
                      new TBackoffConnect(Utils.getInt(conf.get(Config.STORM_NIMBUS_RETRY_TIMES)), Utils.getInt(conf.get(Config.STORM_NIMBUS_RETRY_INTERVAL)),
                              Utils.getInt(conf.get(Config.STORM_NIMBUS_RETRY_INTERVAL_CEILING)));
              _transport = connectionRetry.doConnectWithRetry(transportPlugin, underlyingTransport, host, asUser);
          } catch (IOException ex) {
              throw new RuntimeException(ex);
          }
          _protocol = null;
          if (_transport != null) {
              _protocol = new TBinaryProtocol(_transport);
          }
      }
  
      public synchronized void close() {
          if (_transport != null) {
              _transport.close();
              _transport = null;
              _protocol = null;
          }
      }
  
  }
  