From 7b53ff287de99e84e2ab7f6b21763f24194ba13e Mon Sep 17 00:00:00 2001
From: Simon Rettberg
Date: Thu, 20 Nov 2014 18:44:34 +0100
Subject: Enable TLS support for thrift connection aswell

---
 .../imagemaster/thrift/server/BinaryListener.java  | 97 +++++++++++++++++++---
 1 file changed, 84 insertions(+), 13 deletions(-)

(limited to 'src/main/java/org/openslx/imagemaster/thrift/server/BinaryListener.java')

diff --git a/src/main/java/org/openslx/imagemaster/thrift/server/BinaryListener.java b/src/main/java/org/openslx/imagemaster/thrift/server/BinaryListener.java
index ebacbfc..d7a3c12 100644
--- a/src/main/java/org/openslx/imagemaster/thrift/server/BinaryListener.java
+++ b/src/main/java/org/openslx/imagemaster/thrift/server/BinaryListener.java
@@ -1,37 +1,108 @@
 package org.openslx.imagemaster.thrift.server;
 
+import java.security.NoSuchAlgorithmException;
+import java.util.concurrent.TimeUnit;
+
+import javax.net.ssl.SSLContext;
+import javax.net.ssl.SSLSocketFactory;
+
 import org.apache.log4j.Logger;
 import org.apache.thrift.protocol.TProtocolFactory;
 import org.apache.thrift.server.THsHaServer;
 import org.apache.thrift.server.TServer;
+import org.apache.thrift.server.TThreadPoolServer;
+import org.apache.thrift.transport.TFramedTransport;
 import org.apache.thrift.transport.TNonblockingServerSocket;
 import org.apache.thrift.transport.TNonblockingServerTransport;
+import org.apache.thrift.transport.TSSLTransportFactory;
+import org.apache.thrift.transport.TSSLTransportFactory.TSSLTransportParameters;
+import org.apache.thrift.transport.TServerTransport;
 import org.apache.thrift.transport.TTransportException;
+import org.openslx.imagemaster.Globals;
 import org.openslx.imagemaster.thrift.iface.ImageServer;
 
 public class BinaryListener implements Runnable
 {
+	private static final int MAX_MSG_LEN = 30 * 1000 * 1000;
+
+	private final ImageServer.Processor<ImageServerHandler> processor = new ImageServer.Processor<ImageServerHandler>( new ImageServerHandler() );
+	final TProtocolFactory protFactory = new TBinaryProtocolSafe.Factory( true, true );
+
 	private static Logger log = Logger.getLogger( BinaryListener.class );
+	final TServer server;
 
 	@Override
 	public void run()
 	{
-		final ImageServerHandler handler = new ImageServerHandler();
-		final ImageServer.Processor<ImageServerHandler> processor = new ImageServer.Processor<ImageServerHandler>( handler );
-		final TProtocolFactory protFactory = new TBinaryProtocolSafe.Factory( true, true );
-		final TNonblockingServerTransport transport;
+		log.info( "Starting Binary Thrift" );
+		server.serve();
+		log.info( "Stopped Binary Thrift" );
+		System.exit( 1 ); // Exit so the server can fully restart
+	}
+
+	public BinaryListener( int port, boolean secure ) throws TTransportException, NoSuchAlgorithmException
+	{
+		if ( secure )
+			server = initSecure( port );
+		else
+			server = initNormal( port );
+	}
+
+	/**
+	 * Listen with TLS wrapping - has to use the threadpool server, since encrypted
+	 * servers cannot use nonblocking sockets :(
+	 * 
+	 * @param port listen port
+	 * @return the server
+	 * @throws NoSuchAlgorithmException
+	 * @throws TTransportException
+	 */
+	private TServer initSecure( int port ) throws NoSuchAlgorithmException, TTransportException
+	{
+		SSLContext context = SSLContext.getDefault();
+		SSLSocketFactory sf = context.getSocketFactory();
+		String[] cipherSuites = sf.getSupportedCipherSuites();
+		// TODO: Remove insecure ones
+		final TSSLTransportParameters params = new TSSLTransportParameters( "TLS", cipherSuites );
+		params.setKeyStore( Globals.getSslKeystoreFile(), Globals.getSslKeystorePassword() );
+		TServerTransport serverTransport;
 		try {
-			transport = new TNonblockingServerSocket( 9090 );
+			serverTransport = TSSLTransportFactory.getServerSocket( port, 0, null, params );
 		} catch ( TTransportException e ) {
-			log.fatal( "Could not listen on port 9090" );
-			return;
+			log.fatal( "Could not listen on port " + port );
+			throw e;
 		}
-		THsHaServer.Args args = new THsHaServer.Args( transport ).protocolFactory( protFactory ).processor( processor ).workerThreads( 8 );
-		args.maxReadBufferBytes = 30l * 1000l * 1000l;
-		TServer server = new THsHaServer( args );
-		log.info( "Starting Binary Thrift" );
-		server.serve();
-		System.exit(1);
+		TThreadPoolServer.Args args = new TThreadPoolServer.Args( serverTransport );
+		args.protocolFactory( protFactory );
+		args.processor( processor );
+		args.minWorkerThreads( 4 ).maxWorkerThreads( 256 );
+		args.requestTimeout( 30 ).requestTimeoutUnit( TimeUnit.SECONDS );
+		args.transportFactory( new TFramedTransport.Factory( MAX_MSG_LEN ) );
+		return new TThreadPoolServer( args );
+	}
+
+	/**
+	 * Create normal plain server, no encryption.
+	 * 
+	 * @param port listen port
+	 * @return server instance
+	 * @throws TTransportException
+	 */
+	public TServer initNormal( int port ) throws TTransportException
+	{
+		final TNonblockingServerTransport serverTransport;
+		try {
+			serverTransport = new TNonblockingServerSocket( port );
+		} catch ( TTransportException e ) {
+			log.fatal( "Could not listen on port " + port );
+			throw e;
+		}
+		THsHaServer.Args args = new THsHaServer.Args( serverTransport );
+		args.protocolFactory( protFactory );
+		args.processor( processor );
+		args.workerThreads( 8 );
+		args.maxReadBufferBytes = MAX_MSG_LEN;
+		return new THsHaServer( args );
 	}
 
 }
-- 
cgit v1.2.3-55-g7522