/*
* 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 gobblin.tunnel;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotEquals;
import static org.testng.Assert.assertTrue;
import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.commons.codec.binary.Hex;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
/**
* Tests for Tunnel with arbitrary TCP traffic.
*
* @author kkandekar@linkedin.com
*/
@Test(singleThreaded = true, groups = { "gobblin.tunnel", "disabledOnTravis" })
public class TestTunnelWithArbitraryTCPTraffic {
private static final Logger LOG = LoggerFactory.getLogger(TestTunnelWithArbitraryTCPTraffic.class);
MockServer doubleEchoServer;
MockServer delayedDoubleEchoServer;
MockServer talkFirstEchoServer;
@BeforeClass
void setUp() throws IOException {
doubleEchoServer = startDoubleEchoServer();
LOG.info("Double Echo Server on " + doubleEchoServer.getServerSocketPort());
delayedDoubleEchoServer = startDoubleEchoServer(1000);
LOG.info("Delayed DoubleEchoServer on " + delayedDoubleEchoServer.getServerSocketPort());
talkFirstEchoServer = startTalkFirstEchoServer();
LOG.info("TalkFirstEchoServer on " + talkFirstEchoServer.getServerSocketPort());
}
@AfterClass
void cleanUp() {
doubleEchoServer.stopServer();
delayedDoubleEchoServer.stopServer();
talkFirstEchoServer.stopServer();
MockServer.sleepQuietly(100);
int nAlive = 0;
for (EasyThread t : EasyThread.ALL_THREADS) {
if (t.isAlive()) {
LOG.warn(t + " IS ALIVE!");
nAlive++;
}
}
LOG.warn("Threads left alive " + nAlive);
}
private MockServer startConnectProxyServer(final boolean largeResponse,
final boolean mixServerAndProxyResponse,
final int nBytesToCloseSocketAfter) throws IOException {
return new ConnectProxyServer(mixServerAndProxyResponse, largeResponse, nBytesToCloseSocketAfter).start();
}
private MockServer startConnectProxyServer(final boolean largeResponse,
final boolean mixServerAndProxyResponse) throws IOException {
return startConnectProxyServer(largeResponse, mixServerAndProxyResponse, -1);
}
private MockServer startConnectProxyServer() throws IOException {
return startConnectProxyServer(false, false);
}
private MockServer startDoubleEchoServer() throws IOException {
return startDoubleEchoServer(0);
}
private MockServer startDoubleEchoServer(final long delay) throws IOException {
return new DoubleEchoServer(delay).start();
}
private static String readFromSocket(SocketChannel client) throws IOException {
ByteBuffer readBuf = ByteBuffer.allocate(256);
LOG.info("Reading from client");
client.read(readBuf);
readBuf.flip();
return StandardCharsets.US_ASCII.decode(readBuf).toString();
}
private static void writeToSocket(SocketChannel client, byte [] bytes) throws IOException {
client.write(ByteBuffer.wrap(bytes));
client.socket().getOutputStream().flush();
}
// Baseline test to ensure clients work without tunnel
@Test(timeOut = 15000)
public void testDirectConnectionToEchoServer() throws IOException {
SocketChannel client = SocketChannel.open();
try {
client.connect(new InetSocketAddress("localhost", doubleEchoServer.getServerSocketPort()));
writeToSocket(client, "Knock\n".getBytes());
String response = readFromSocket(client);
client.close();
assertEquals(response, "Knock Knock\n");
} finally {
client.close();
}
}
@Test(timeOut = 15000)
public void testTunnelToEchoServer() throws IOException {
MockServer proxyServer = startConnectProxyServer();
Tunnel tunnel = Tunnel.build("localhost", doubleEchoServer.getServerSocketPort(), "localhost",
proxyServer.getServerSocketPort());
try {
int tunnelPort = tunnel.getPort();
SocketChannel client = SocketChannel.open();
client.connect(new InetSocketAddress("localhost", tunnelPort));
client.write(ByteBuffer.wrap("Knock\n".getBytes()));
String response = readFromSocket(client);
client.close();
assertEquals(response, "Knock Knock\n");
assertEquals(proxyServer.getNumConnects(), 1);
} finally {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
}
}
@Test(timeOut = 15000)
public void testTunnelToDelayedEchoServer() throws IOException {
MockServer proxyServer = startConnectProxyServer();
Tunnel tunnel = Tunnel.build("localhost", delayedDoubleEchoServer.getServerSocketPort(), "localhost",
proxyServer.getServerSocketPort());
try {
int tunnelPort = tunnel.getPort();
SocketChannel client = SocketChannel.open();
client.connect(new InetSocketAddress("localhost", tunnelPort));
client.write(ByteBuffer.wrap("Knock\n".getBytes()));
String response = readFromSocket(client);
client.close();
assertEquals(response, "Knock Knock\n");
assertEquals(proxyServer.getNumConnects(), 1);
} finally {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
}
}
@Test(timeOut = 15000)
public void testTunnelToEchoServerMultiRequest() throws IOException {
MockServer proxyServer = startConnectProxyServer();
Tunnel tunnel = Tunnel.build("localhost", doubleEchoServer.getServerSocketPort(),
"localhost", proxyServer.getServerSocketPort());
try {
int tunnelPort = tunnel.getPort();
SocketChannel client = SocketChannel.open();
client.connect(new InetSocketAddress("localhost", tunnelPort));
client.write(ByteBuffer.wrap("Knock\n".getBytes()));
String response1 = readFromSocket(client);
client.write(ByteBuffer.wrap("Hello\n".getBytes()));
String response2 = readFromSocket(client);
client.close();
assertEquals(response1, "Knock Knock\n");
assertEquals(response2, "Hello Hello\n");
assertEquals(proxyServer.getNumConnects(), 1);
} finally {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
}
}
private MockServer startTalkFirstEchoServer() throws IOException {
return new TalkFirstDoubleEchoServer().start();
}
private void runClientToTalkFirstServer(int tunnelPort) throws IOException {
SocketChannel client = SocketChannel.open();
client.connect(new InetSocketAddress("localhost", tunnelPort));
String response0 = readFromSocket(client);
LOG.info(response0);
client.write(ByteBuffer.wrap("Knock\n".getBytes()));
String response1 = readFromSocket(client);
LOG.info(response1);
client.write(ByteBuffer.wrap("Hello\n".getBytes()));
String response2 = readFromSocket(client);
LOG.info(response2);
client.close();
assertEquals(response0, "Hello\n");
assertEquals(response1, "Knock Knock\n");
assertEquals(response2, "Hello Hello\n");
}
@Test(timeOut = 15000)
public void testTunnelToEchoServerThatRespondsFirst() throws IOException {
MockServer proxyServer = startConnectProxyServer();
Tunnel tunnel = Tunnel.build("localhost", talkFirstEchoServer.getServerSocketPort(),
"localhost", proxyServer.getServerSocketPort());
try {
int tunnelPort = tunnel.getPort();
runClientToTalkFirstServer(tunnelPort);
assertEquals(proxyServer.getNumConnects(), 1);
} finally {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
}
}
@Test(timeOut = 15000)
public void testTunnelToEchoServerThatRespondsFirstWithMixedProxyAndServerResponseInBuffer() throws IOException {
MockServer proxyServer = startConnectProxyServer(false, true);
Tunnel tunnel = Tunnel.build("localhost", talkFirstEchoServer.getServerSocketPort(),
"localhost", proxyServer.getServerSocketPort());
try {
int tunnelPort = tunnel.getPort();
runClientToTalkFirstServer(tunnelPort);
assertEquals(proxyServer.getNumConnects(), 1);
} finally {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
}
}
@Test(timeOut = 15000)
public void testTunnelToEchoServerThatRespondsFirstAcrossMultipleDrainReads() throws IOException {
MockServer proxyServer = startConnectProxyServer(true, true);
Tunnel tunnel = Tunnel.build("localhost", talkFirstEchoServer.getServerSocketPort(),
"localhost", proxyServer.getServerSocketPort());
try {
int tunnelPort = tunnel.getPort();
runClientToTalkFirstServer(tunnelPort);
assertEquals(proxyServer.getNumConnects(), 1);
} finally {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
}
}
@Test(timeOut = 15000)
public void testTunnelToEchoServerThatRespondsFirstAcrossMultipleDrainReadsWithMultipleClients()
throws IOException, InterruptedException {
MockServer proxyServer = startConnectProxyServer(true, true);
Tunnel tunnel = Tunnel.build("localhost", talkFirstEchoServer.getServerSocketPort(),
"localhost", proxyServer.getServerSocketPort());
try {
final int tunnelPort = tunnel.getPort();
List<EasyThread> threads = new ArrayList<EasyThread>();
for (int i = 0; i < 5; i++) {
threads.add(new EasyThread() {
@Override
void runQuietly() throws Exception {
try {
runClientToTalkFirstServer(tunnelPort);
} catch (IOException e) {
e.printStackTrace();
}
}
}.startThread());
}
for (Thread t : threads) {
t.join();
}
assertEquals(proxyServer.getNumConnects(), 5);
} finally {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
}
}
private void runSimultaneousDataExchange(boolean useTunnel, int nclients)
throws IOException, InterruptedException, NoSuchAlgorithmException {
long t0 = System.currentTimeMillis();
final int nMsgs = 50;
final Map<String, MessageDigest> digestMsgsRecvdAtServer = new HashMap<String, MessageDigest>();
final Map<String, MessageDigest> digestMsgsSentByClients = new HashMap<String, MessageDigest>();
final Map<String, MessageDigest> digestMsgsRecvdAtClients = new HashMap<String, MessageDigest>();
for (int c = 0; c < nclients ; c++) {
digestMsgsRecvdAtServer.put(Integer.toString(c), MessageDigest.getInstance("MD5"));
digestMsgsSentByClients.put(Integer.toString(c), MessageDigest.getInstance("MD5"));
digestMsgsRecvdAtClients.put(Integer.toString(c), MessageDigest.getInstance("MD5"));
}
final MessageDigest digestMsgsSentByServer = MessageDigest.getInstance("MD5");
for (int i = 0; i < nMsgs; i++) {
digestMsgsSentByServer.update(TalkPastServer.generateMsgFromServer(i).getBytes());
}
String hashOfMsgsSentByServer = Hex.encodeHexString(digestMsgsSentByServer.digest());
MockServer talkPastServer = startTalkPastServer(nMsgs, digestMsgsRecvdAtServer);
int targetPort = talkPastServer.getServerSocketPort();
Tunnel tunnel = null;
MockServer proxyServer = null;
if (useTunnel) {
proxyServer = startConnectProxyServer();
tunnel = Tunnel.build("localhost", talkPastServer.getServerSocketPort(),
"localhost", proxyServer.getServerSocketPort());
targetPort = tunnel.getPort();
}
try {
List<EasyThread> clientThreads = new ArrayList<EasyThread>();
final int portToUse = targetPort;
for (int c = 0; c < nclients; c++) {
final int clientId = c;
clientThreads.add(new EasyThread() {
@Override
void runQuietly() throws Exception {
long t = System.currentTimeMillis();
LOG.info("\t" + clientId + ": Client starting");
final MessageDigest digestMsgsRecvdAtClient = digestMsgsRecvdAtClients.get(Integer.toString(clientId));
//final SocketChannel client = SocketChannel.open(); // tunnel test hangs for some reason with SocketChannel
final Socket client = new Socket();
client.connect(new InetSocketAddress("localhost", portToUse));
EasyThread serverReaderThread = new EasyThread() {
@Override
public void runQuietly() {
try {
BufferedReader clientIn = new BufferedReader(new InputStreamReader(client.getInputStream()));
String line = clientIn.readLine();
while (line != null && !line.equals("Goodbye")) {
//LOG.info("\t" + clientId + ": Server said [" + line.substring(0, 32) + "... ]");
digestMsgsRecvdAtClient.update(line.getBytes());
digestMsgsRecvdAtClient.update("\n".getBytes());
line = clientIn.readLine();
}
} catch (IOException e) {
e.printStackTrace();
}
LOG.info("\t" + clientId + ": Client done reading");
}
}.startThread();
MessageDigest hashMsgsFromClient = digestMsgsSentByClients.get(Integer.toString(clientId));
BufferedOutputStream clientOut = new BufferedOutputStream(client.getOutputStream());
for (int i = 0; i < nMsgs; i++) {
String msg = clientId + ":" + i + " " + StringUtils.repeat("Blahhh Blahhh ", 10000) +"\n";
//LOG.info(clientId + " sending " + msg.length() + " bytes");
byte [] bytes = msg.getBytes();
hashMsgsFromClient.update(bytes);
clientOut.write(bytes);
MockServer.sleepQuietly(2);
}
clientOut.write(("Goodbye\n".getBytes()));
clientOut.flush();
LOG.info("\t" + clientId + ": Client done writing in " + (System.currentTimeMillis() - t) + " ms");
serverReaderThread.join();
LOG.info("\t" + clientId + ": Client done in " + (System.currentTimeMillis() - t) + " ms");
client.close();
}
}.startThread());
}
for (Thread clientThread : clientThreads) {
clientThread.join();
}
LOG.info("All data transfer done in " + (System.currentTimeMillis() - t0) + " ms");
} finally {
talkPastServer.stopServer();
if (tunnel != null) {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
assertEquals(proxyServer.getNumConnects(), nclients);
}
Map<String, String> hashOfMsgsRecvdAtServer = new HashMap<String, String>();
Map<String, String> hashOfMsgsSentByClients = new HashMap<String, String>();
Map<String, String> hashOfMsgsRecvdAtClients = new HashMap<String, String>();
for (int c = 0; c < nclients; c++) {
String client = Integer.toString(c);
hashOfMsgsRecvdAtServer.put(client, Hex.encodeHexString(digestMsgsRecvdAtServer.get(client).digest()));
hashOfMsgsSentByClients.put(client, Hex.encodeHexString(digestMsgsSentByClients.get(client).digest()));
hashOfMsgsRecvdAtClients.put(client, Hex.encodeHexString(digestMsgsRecvdAtClients.get(client).digest()));
}
LOG.info("\tComparing client sent to server received");
assertEquals(hashOfMsgsSentByClients, hashOfMsgsRecvdAtServer);
LOG.info("\tComparing server sent to client received");
for (String hashOfMsgsRecvdAtClient : hashOfMsgsRecvdAtClients.values()) {
assertEquals(hashOfMsgsSentByServer, hashOfMsgsRecvdAtClient);
}
LOG.info("\tDone");
}
}
private MockServer startTalkPastServer(final int nMsgs, final Map<String, MessageDigest> digestMsgsRecvdAtServer) throws IOException {
return new TalkPastServer(nMsgs, digestMsgsRecvdAtServer).start();
}
// Baseline tests to ensure simultaneous data transfer protocol is fine (disabled for now)
@Test(enabled = false, timeOut = 15000)
public void testSimultaneousDataExchangeWithDirectConnection()
throws IOException, InterruptedException, NoSuchAlgorithmException {
runSimultaneousDataExchange(false, 1);
}
@Test(enabled = false, timeOut = 15000)
public void testSimultaneousDataExchangeWithDirectConnectionAndMultipleClients()
throws IOException, InterruptedException, NoSuchAlgorithmException {
runSimultaneousDataExchange(false, 3);
}
/*
I wrote this test because I saw this symptom once randomly while testing with Gobblin. Test passes, but occasionally
we see the following warning in the logs:
15/10/29 21:11:17 WARN tunnel.Tunnel: exception handling event on java.nio.channels.SocketChannel[connected local=/127.0.0.1:34669 remote=/127.0.0.1:38578]
java.nio.channels.CancelledKeyException
at sun.nio.ch.SelectionKeyImpl.ensureValid(SelectionKeyImpl.java:73)
at sun.nio.ch.SelectionKeyImpl.readyOps(SelectionKeyImpl.java:87)
at java.nio.channels.SelectionKey.isWritable(SelectionKey.java:312)
at gobblin.tunnel.Tunnel$ReadWriteHandler.write(Tunnel.java:423)
at gobblin.tunnel.Tunnel$ReadWriteHandler.call(Tunnel.java:403)
at gobblin.tunnel.Tunnel$ReadWriteHandler.call(Tunnel.java:365)
at gobblin.tunnel.Tunnel$Dispatcher.dispatch(Tunnel.java:142)
at gobblin.tunnel.Tunnel$Dispatcher.run(Tunnel.java:127)
at java.lang.Thread.run(Thread.java:745)
*/
@Test(timeOut = 20000)
public void testSimultaneousDataExchangeWithTunnel()
throws IOException, InterruptedException, NoSuchAlgorithmException {
runSimultaneousDataExchange(true, 1);
}
@Test(timeOut = 20000)
public void testSimultaneousDataExchangeWithTunnelAndMultipleClients()
throws IOException, InterruptedException, NoSuchAlgorithmException {
runSimultaneousDataExchange(true, 3);
}
@Test(expectedExceptions = IOException.class)
public void testTunnelWhereProxyConnectionToServerFailsWithWriteFirstClient() throws IOException, InterruptedException {
MockServer proxyServer = startConnectProxyServer();
final int nonExistentPort = 54321; // hope this doesn't exist!
Tunnel tunnel = Tunnel.build("localhost", nonExistentPort, "localhost", proxyServer.getServerSocketPort());
try {
int tunnelPort = tunnel.getPort();
SocketChannel client = SocketChannel.open();
client.configureBlocking(true);
client.connect(new InetSocketAddress("localhost", tunnelPort));
// Might have to write multiple times before connection error propagates back from proxy through tunnel to client
for (int i = 0; i < 5; i++) {
client.write(ByteBuffer.wrap("Knock\n".getBytes()));
Thread.sleep(100);
}
String response1 = readFromSocket(client);
LOG.info(response1);
client.close();
} finally {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
assertEquals(proxyServer.getNumConnects(), 1);
}
}
@Test(timeOut = 15000)
public void testTunnelThreadDeadAfterClose() throws IOException, InterruptedException {
MockServer proxyServer = startConnectProxyServer();
Tunnel tunnel = Tunnel.build("localhost", talkFirstEchoServer.getServerSocketPort(),
"localhost", proxyServer.getServerSocketPort());
try {
int tunnelPort = tunnel.getPort();
SocketChannel client = SocketChannel.open();
client.connect(new InetSocketAddress("localhost", tunnelPort));
String response0 = readFromSocket(client);
LOG.info(response0);
// write a lot of data to increase chance of response after close
for (int i = 0; i < 1000; i++) {
client.write(ByteBuffer.wrap("Knock\n".getBytes()));
}
// don't wait for response
client.close();
assertEquals(response0, "Hello\n");
assertEquals(proxyServer.getNumConnects(), 1);
} finally {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
}
}
@Test(timeOut = 15000, expectedExceptions = IOException.class)
public void testTunnelThreadDeadAfterUnexpectedException() throws IOException, InterruptedException {
MockServer proxyServer = startConnectProxyServer(false, false, 8);
Tunnel tunnel = Tunnel.build("localhost", doubleEchoServer.getServerSocketPort(),
"localhost", proxyServer.getServerSocketPort());
String response = "";
try {
int tunnelPort = tunnel.getPort();
SocketChannel client = SocketChannel.open();
client.connect(new InetSocketAddress("localhost", tunnelPort));
client.write(ByteBuffer.wrap("Knock\n".getBytes()));
response = readFromSocket(client);
LOG.info(response);
for (int i = 0; i < 5; i++) {
client.write(ByteBuffer.wrap("Hello\n".getBytes()));
Thread.sleep(100);
}
client.close();
} finally {
proxyServer.stopServer();
tunnel.close();
assertNotEquals(response, "Knock Knock\n");
assertEquals(proxyServer.getNumConnects(), 1);
assertFalse(tunnel.isTunnelThreadAlive());
}
}
/**
* This test demonstrates connecting to a mysql DB through
* and http proxy tunnel to a public data set of genetic data
* http://www.ensembl.org/info/data/mysql.html
*
* Disabled for now because it requires the inclusion of a mysql jdbc driver jar
*
* @throws Exception
*/
@Test(enabled = false, timeOut = 40000)
public void accessEnsembleDB() throws Exception{
MockServer proxyServer = startConnectProxyServer();
Tunnel tunnel = Tunnel.build("useastdb.ensembl.org", 5306,
"localhost", proxyServer.getServerSocketPort());
try {
int port = tunnel.getPort();
Connection connection =
DriverManager.getConnection("jdbc:mysql://localhost:" + port + "/homo_sapiens_core_82_38?user=anonymous");
String query2 = "SELECT DISTINCT gene_id, biotype, source, description from gene LIMIT 1000";
ResultSet resultSet = connection.createStatement().executeQuery(query2);
int row = 0;
while (resultSet.next()) {
row++;
LOG.info(String.format("%s|%s|%s|%s|%s%n", row, resultSet.getString(1), resultSet.getString(2), resultSet.getString(3),
resultSet.getString(4)));
}
assertEquals(row, 1000);
assertTrue(proxyServer.getNumConnects() > 0);
}
finally {
proxyServer.stopServer();
tunnel.close();
assertFalse(tunnel.isTunnelThreadAlive());
}
}
}