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.curator.framework.CuratorFramework;
import org.apache.thrift7.protocol.TBinaryProtocol;
import org.apache.thrift7.protocol.TProtocol;
import org.apache.thrift7.transport.TSocket;
import org.apache.thrift7.transport.TTransport;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import backtype.storm.Config;
import backtype.storm.utils.Utils;
public class ThriftClient {
private static final Logger LOG = LoggerFactory
.getLogger(ThriftClient.class);
private static final String MASTER_PATH = "/nimbus_master";
private TTransport _transport;
protected TProtocol _protocol;
private CuratorFramework zkobj;
private String masterHost;
private final String zkMasterDir;
private Map<Object, Object> conf;
public ThriftClient(Map storm_conf) throws Exception {
this(storm_conf, null);
}
@SuppressWarnings("unchecked")
public ThriftClient(Map storm_conf, Integer timeout) throws Exception {
conf = storm_conf;
String root = String.valueOf(storm_conf
.get(Config.STORM_ZOOKEEPER_ROOT));
zkMasterDir = root + MASTER_PATH;
LOG.info("zkServer:"
+ (List<String>) storm_conf.get(Config.STORM_ZOOKEEPER_SERVERS)
+ ", zkPort:"
+ (Integer) storm_conf.get(Config.STORM_ZOOKEEPER_PORT));
zkobj = Utils.newCurator(storm_conf,
(List<String>) storm_conf.get(Config.STORM_ZOOKEEPER_SERVERS),
storm_conf.get(Config.STORM_ZOOKEEPER_PORT), zkMasterDir);
zkobj.start();
if (zkobj.checkExists().forPath("/") == null)
throw new RuntimeException("No alive nimbus ");
flushClient(storm_conf, timeout);
}
public TTransport transport() {
return _transport;
}
protected void flushClient(Map storm_conf, Integer timeout)
throws Exception {
try {
flushHost();
String[] host_port = masterHost.split(":");
if (host_port.length != 2) {
throw new InvalidParameterException("Host format error: "
+ masterHost);
}
String host = host_port[0];
int port = Integer.parseInt(host_port[1]);
LOG.info("Begin to connect " + host + ":" + port);
// locate login configuration
Configuration login_conf = AuthUtils.GetConfiguration(storm_conf);
// construct a transport plugin
ITransportPlugin transportPlugin = AuthUtils.GetTransportPlugin(
storm_conf, login_conf);
// create a socket with server
if (host == null) {
throw new IllegalArgumentException("host is not set");
}
if (port <= 0) {
throw new IllegalArgumentException("invalid port: " + port);
}
// /***************only test for daily *************/
// if (host.endsWith("bja")) {
// host += ".tbsite.net";
// }
// /***************only test for daily *************/
TSocket socket = new TSocket(host, port);
if (timeout != null) {
socket.setTimeout(timeout);
}
final TTransport underlyingTransport = socket;
// establish client-server transport via plugin
_transport = transportPlugin.connect(underlyingTransport, host);
} catch (IOException ex) {
throw new RuntimeException("Create transport error");
}
_protocol = null;
if (_transport != null)
_protocol = new TBinaryProtocol(_transport);
}
private void flushHost() throws Exception {
masterHost = new String(zkobj.getData().forPath("/"));
}
public void close() {
if (_transport != null)
_transport.close();
if (zkobj != null)
zkobj.close();
}
public String getMasterHost() {
return masterHost;
}
public Map<Object, Object> getConf() {
return conf;
}
protected void flush() {
}
}