Merge pull request #6 from fogbow/develop

Merge branch 'develop' of https://github.com/fogbow/fogbow-reverse-tunnel.git into develop

pom.xml

@@ -14,7 +14,7 @@
 					<target>1.6</target>
 				</configuration>
 			</plugin>
-                        <plugin>
+			<plugin>
 				<artifactId>maven-assembly-plugin</artifactId>
 				<configuration>
 					<archive>
@@ -39,10 +39,15 @@
 		</plugins>
 	</build>
 	<dependencies>
+		<dependency>
+			<groupId>org.apache.mina</groupId>
+			<artifactId>mina-core</artifactId>
+			<version>2.0.9</version>
+		</dependency>
 		<dependency>
 			<groupId>org.apache.sshd</groupId>
 			<artifactId>sshd-core</artifactId>
-			<version>0.12.0</version>
+			<version>0.14.0</version>
 		</dependency>
 		<dependency>
 			<groupId>org.slf4j</groupId>
@@ -69,6 +74,11 @@
 			<artifactId>commons-io</artifactId>
 			<version>2.4</version>
 		</dependency>
+		<dependency>
+			<groupId>org.json</groupId>
+			<artifactId>json</artifactId>
+			<version>20140107</version>
+		</dependency>
 	</dependencies>
 	<repositories>
 		<repository>

reverse-tunnel.conf.example

@@ -3,3 +3,4 @@ tunnel_host=0.0.0.0
 http_port=2223
 external_port_range=20000:30000
 host_key_path=hostkey.ser
+idle_token_timeout=86400

src/main/java/org/fogbowcloud/ssh/Main.java

@@ -17,13 +17,19 @@ public static void main(String[] args) throws IOException {
 		String externalPortRange = properties.getProperty("external_port_range");
 		String[] externalRangeSplit = externalPortRange.split(":");
 		String externalHostKeyPath = properties.getProperty("host_key_path");
+		String idleTokenTimeoutStr = properties.getProperty("idle_token_timeout");
+		Long idleTokenTimeout = null;
+		if (idleTokenTimeoutStr != null) {
+			idleTokenTimeout = Long.parseLong(idleTokenTimeoutStr) * 1000;
+		}
 
 		TunnelHttpServer tunnelHttpServer = new TunnelHttpServer(
 				Integer.parseInt(httpPort),
 				tunnelHost,
 				Integer.parseInt(tunnelPort),
 				Integer.parseInt(externalRangeSplit[0]), 
 				Integer.parseInt(externalRangeSplit[1]),
+				idleTokenTimeout,
 				externalHostKeyPath);
 		tunnelHttpServer.start();
 	}

src/main/java/org/fogbowcloud/ssh/ReverseTunnelForwarder.java

@@ -85,8 +85,10 @@ public synchronized SshdSocketAddress localPortForwardingRequested(SshdSocketAdd
 
     public synchronized void localPortForwardingCancelled(SshdSocketAddress local) throws IOException {
         if (localForwards.remove(local) && acceptor != null) {
-        	Nio2Acceptor a = (Nio2Acceptor) acceptor;
-        	a.doCloseImmediately();
+        	if (acceptor instanceof Nio2Acceptor) {
+        		Nio2Acceptor a = (Nio2Acceptor) acceptor;
+        		a.doCloseImmediately();
+        	}
             acceptor.unbind(local.toInetSocketAddress());
             if (acceptor.getBoundAddresses().isEmpty()) {
                 acceptor.close(true);
@@ -195,4 +197,16 @@ private SshdSocketAddress doBind(SshdSocketAddress address) throws IOException {
         }
     }
 
+	@Override
+	public SshdSocketAddress startDynamicPortForwarding(SshdSocketAddress arg0)
+			throws IOException {
+		return arg0;
+	}
+
+	@Override
+	public void stopDynamicPortForwarding(SshdSocketAddress arg0)
+			throws IOException {
+
+	}
+
 }

src/main/java/org/fogbowcloud/ssh/ReverseTunnelSession.java

@@ -1,5 +1,6 @@
 package org.fogbowcloud.ssh;
 
+import org.apache.sshd.common.Service;
 import org.apache.sshd.common.io.IoSession;
 import org.apache.sshd.common.util.Buffer;
 import org.apache.sshd.server.ServerFactoryManager;
@@ -17,5 +18,9 @@ protected void doHandleMessage(Buffer buffer) throws Exception {
 		super.doHandleMessage(buffer);
 		resetIdleTimeout();
 	}
+
+	public Service getService() {
+		return currentService;
+	}
 
 }

src/main/java/org/fogbowcloud/ssh/ReverseTunnelTcpipChannel.java

@@ -13,7 +13,6 @@
 import org.apache.sshd.common.channel.ChannelOutputStream;
 import org.apache.sshd.common.io.IoSession;
 import org.apache.sshd.common.util.Buffer;
-import org.apache.sshd.common.util.CloseableUtils;
 
 public class ReverseTunnelTcpipChannel extends AbstractClientChannel {
 
@@ -67,7 +66,7 @@ protected synchronized void doOpen() throws IOException {
 
     @Override
     protected Closeable getInnerCloseable() {
-        return CloseableUtils.sequential(serverSession, super.getInnerCloseable());
+        return builder().sequential(serverSession, super.getInnerCloseable()).build();
     }
 
     protected synchronized void doWriteData(byte[] data, int off, int len) throws IOException {

src/main/java/org/fogbowcloud/ssh/TunnelHttpServer.java

@@ -5,8 +5,10 @@
 import java.io.ObjectInputStream;
 import java.io.UnsupportedEncodingException;
 import java.security.KeyPair;
+import java.util.Map;
 
-import org.apache.mina.util.Base64;
+import org.apache.sshd.common.util.Base64;
+import org.json.JSONObject;
 
 import fi.iki.elonen.NanoHTTPD;
 import fi.iki.elonen.NanoHTTPD.Response.Status;
@@ -18,12 +20,12 @@ public class TunnelHttpServer extends NanoHTTPD {
 	private KeyPair kp;
 
 	public TunnelHttpServer(int httpPort, String sshTunnelHost, int sshTunnelPort, 
-			int lowerPort, int higherPort, String hostKeyPath) {
+			int lowerPort, int higherPort, Long idleTokenTimeout, String hostKeyPath) {
 		super(httpPort);
 		this.hostKeyPath = hostKeyPath;
 		try {
 			this.tunneling = new TunnelServer(sshTunnelHost, sshTunnelPort, 
-					lowerPort, higherPort, hostKeyPath);
+					lowerPort, higherPort, idleTokenTimeout, hostKeyPath);
 			this.tunneling.start();
 		} catch (IOException e) {
 			e.printStackTrace();
@@ -40,26 +42,30 @@ public Response serve(IHTTPSession session) {
 		}
 		if (splitUri[1].equals("token")) {
 
-			if (splitUri.length != 3) {
+			if (splitUri.length > 4) {
 				return new NanoHTTPD.Response(Status.METHOD_NOT_ALLOWED, MIME_PLAINTEXT, "");
 			}
 
 			String tokenId = splitUri[2];
-
+			
 			if (method.equals(Method.GET)) {
-				Integer port = this.tunneling.getPort(tokenId);
-				if (port == null) {
-					return new NanoHTTPD.Response(Status.NOT_FOUND, 
-							MIME_PLAINTEXT, "404 Port Not Found");
+				if (splitUri.length == 4 && splitUri[3].equals("all")) {
+					Map<String, Integer> ports = this.tunneling.getPortByPrefix(tokenId);
+					return new NanoHTTPD.Response(new JSONObject(ports).toString());
+				} else {
+					Integer port = this.tunneling.getPort(tokenId);
+					if (port == null) {
+						return new NanoHTTPD.Response(Status.NOT_FOUND, 
+								MIME_PLAINTEXT, "404 Port Not Found");
+					}
+					return new NanoHTTPD.Response(port.toString());
 				}
-				return new NanoHTTPD.Response(port.toString());
 			}
 
 			if (method.equals(Method.POST)) {
 				Integer port = this.tunneling.createPort(tokenId);
 				if (port == null) {
-					return new NanoHTTPD.Response(Status.INTERNAL_ERROR, 
-							MIME_PLAINTEXT, "Internal error");
+					return new NanoHTTPD.Response(Status.INTERNAL_ERROR, MIME_PLAINTEXT, "");
 				}
 				return new NanoHTTPD.Response(port.toString());
 			}

src/main/java/org/fogbowcloud/ssh/TunnelServer.java

@@ -1,6 +1,7 @@
 package org.fogbowcloud.ssh;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
@@ -16,6 +17,7 @@
 import org.apache.sshd.SshServer;
 import org.apache.sshd.common.ForwardingFilter;
 import org.apache.sshd.common.NamedFactory;
+import org.apache.sshd.common.Service;
 import org.apache.sshd.common.Session;
 import org.apache.sshd.common.Session.AttributeKey;
 import org.apache.sshd.common.SshdSocketAddress;
@@ -55,13 +57,16 @@ public Token(Integer port) {
 	private int lowerPort;
 	private int higherPort;
 	private String hostKeyPath;
+	private Long idleTokenTimeout;
 
 	public TunnelServer(String sshTunnelHost, int sshTunnelPort, int lowerPort, 
-			int higherPort, String hostKeyPath) {
+			int higherPort, Long idleTokenTimeout, String hostKeyPath) {
 		this.sshTunnelHost = sshTunnelHost;
 		this.sshTunnelPort = sshTunnelPort;
 		this.lowerPort = lowerPort;
 		this.higherPort = higherPort;
+		this.idleTokenTimeout = idleTokenTimeout == null ? TOKEN_EXPIRATION_TIMEOUT
+				: idleTokenTimeout;
 		this.hostKeyPath = hostKeyPath;
 	}
 
@@ -77,6 +82,11 @@ public synchronized Integer createPort(String token) {
 			newPort = port;
 			break;
 		}
+		if (newPort == null) {
+			LOGGER.debug("Token [" + token + "] didn't get any port. All ports are busy.");
+			return null;
+		}
+
 		LOGGER.debug("Token [" + token + "] got port [" + newPort + "].");
 		tokens.put(token, new Token(newPort));
 		return newPort;
@@ -94,7 +104,14 @@ private boolean isTaken(int port) {
 	private ReverseTunnelForwarder getActiveSession(int port) {
 		List<AbstractSession> activeSessions = sshServer.getActiveSessions();
 		for (AbstractSession session : activeSessions) {
-			ServerConnectionService service = session.getService(ServerConnectionService.class);
+			Service rawService = ((ReverseTunnelSession)session).getService();
+			if (rawService == null) {
+				continue;
+			}
+			if (!(rawService instanceof ServerConnectionService)) {
+				continue;
+			}
+			ServerConnectionService service = (ServerConnectionService) rawService;
 			ReverseTunnelForwarder f = (ReverseTunnelForwarder) service.getTcpipForwarder();
 			for (SshdSocketAddress address : f.getLocalForwards()) {
 				if (address.getPort() == port) {
@@ -152,7 +169,7 @@ public void run() {
 						if (token.lastIdleCheck == 0) {
 							token.lastIdleCheck = now;
 						}
-						if (now - token.lastIdleCheck > TOKEN_EXPIRATION_TIMEOUT) {
+						if (now - token.lastIdleCheck > idleTokenTimeout) {
 							tokensToExpire.add(tokenEntry.getKey());
 						}
 					} else {
@@ -223,4 +240,31 @@ public Integer getPort(String tokenId) {
 		return token.port;
 	}
 
+	public Map<String, Integer> getAllPorts() {
+		Map<String, Integer> portsByPrefix = new HashMap<String, Integer>();
+		for (Entry<String, Token> tokenEntry : tokens.entrySet()) {
+			portsByPrefix.put(
+					tokenEntry.getKey(), 
+					tokenEntry.getValue().port);
+		}
+		return portsByPrefix;
+	}
+
+	public Map<String, Integer> getPortByPrefix(String tokenId) {
+		Map<String, Integer> portsByPrefix = new HashMap<String, Integer>();
+		Integer sshPort = getPort(tokenId);
+		if (sshPort != null) {
+			portsByPrefix.put("ssh", sshPort);
+		}
+		for (Entry<String, Token> tokenEntry : tokens.entrySet()) {
+			String tokenPrefix = tokenId + "-";
+			if (tokenEntry.getKey().startsWith(tokenPrefix)) {
+				portsByPrefix.put(
+						tokenEntry.getKey().substring(tokenPrefix.length()), 
+						tokenEntry.getValue().port);
+			}
+		}
+		return portsByPrefix;
+	}
+
 }