package backtype.storm.security.auth;
import java.io.IOException;
import java.net.InetAddress;
import java.net.Socket;
import java.net.UnknownHostException;
import java.util.Map;
import javax.security.auth.login.Configuration;
import org.apache.thrift7.TException;
import org.apache.thrift7.TProcessor;
import org.apache.thrift7.protocol.TBinaryProtocol;
import org.apache.thrift7.protocol.TProtocol;
import org.apache.thrift7.server.THsHaServer;
import org.apache.thrift7.server.TServer;
import org.apache.thrift7.transport.TFramedTransport;
import org.apache.thrift7.transport.TMemoryInputTransport;
import org.apache.thrift7.transport.TNonblockingServerSocket;
import org.apache.thrift7.transport.TSocket;
import org.apache.thrift7.transport.TTransport;
import org.apache.thrift7.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Simple transport for Thrift plugin.
*
* This plugin is designed to be backward compatible with existing Storm code.
*/
public class SimpleTransportPlugin implements ITransportPlugin {
protected Configuration login_conf;
private static final Logger LOG = LoggerFactory
.getLogger(SimpleTransportPlugin.class);
/**
* Invoked once immediately after construction
*
* @param conf
* Storm configuration
* @param login_conf
* login configuration
*/
public void prepare(Map storm_conf, Configuration login_conf) {
this.login_conf = login_conf;
}
/**
* We will let Thrift to apply default transport factory
*/
public TServer getServer(int port, TProcessor processor)
throws IOException, TTransportException {
TNonblockingServerSocket serverTransport = new TNonblockingServerSocket(
port);
THsHaServer.Args server_args = new THsHaServer.Args(serverTransport)
.processor(new SimpleWrapProcessor(processor))
.workerThreads(64)
.protocolFactory(new TBinaryProtocol.Factory());
// construct THsHaServer
return new THsHaServer(server_args);
}
/**
* Connect to the specified server via framed transport
*
* @param transport
* The underlying Thrift transport.
*/
public TTransport connect(TTransport transport, String serverHost)
throws TTransportException {
// create a framed transport
TTransport conn = new TFramedTransport(transport);
// connect
conn.open();
LOG.debug("Simple client transport has been established");
return conn;
}
/**
* Processor that populate simple transport info into ReqContext, and then
* invoke a service handler
*/
private class SimpleWrapProcessor implements TProcessor {
final TProcessor wrapped;
SimpleWrapProcessor(TProcessor wrapped) {
this.wrapped = wrapped;
}
public boolean process(final TProtocol inProt, final TProtocol outProt)
throws TException {
// populating request context
ReqContext req_context = ReqContext.context();
TTransport trans = inProt.getTransport();
if (trans instanceof TMemoryInputTransport) {
try {
req_context.setRemoteAddress(InetAddress.getLocalHost());
} catch (UnknownHostException e) {
throw new RuntimeException(e);
}
} else if (trans instanceof TSocket) {
TSocket tsocket = (TSocket) trans;
// remote address
Socket socket = tsocket.getSocket();
req_context.setRemoteAddress(socket.getInetAddress());
}
// anonymous user
req_context.setSubject(null);
// invoke service handler
return wrapped.process(inProt, outProt);
}
}
}