Search in sources :

Example 1 with PortForwarding

use of io.cdap.cdap.runtime.spi.ssh.PortForwarding in project cdap by caskdata.

the class SSHSessionTest method testLocalPortForwarding.

@Test
public void testLocalPortForwarding() throws Exception {
    // Starts an echo server for testing the port forwarding
    EchoServer echoServer = new EchoServer();
    echoServer.startAndWait();
    try {
        // Creates the DataConsumer for receiving data and validating the lifecycle
        StringBuilder received = new StringBuilder();
        AtomicBoolean finished = new AtomicBoolean();
        PortForwarding.DataConsumer dataConsumer = new PortForwarding.DataConsumer() {

            private final List<String> messages = new ArrayList<>();

            @Override
            public void received(ByteBuffer buffer) {
                messages.add(StandardCharsets.UTF_8.decode(buffer).toString());
            }

            @Override
            public synchronized void flushed() {
                messages.forEach(received::append);
            }

            @Override
            public void finished() {
                finished.set(true);
            }
        };
        SSHConfig sshConfig = getSSHConfig();
        // Creates a SSH session.
        try (SSHSession session = new DefaultSSHSession(sshConfig)) {
            InetSocketAddress bindAddress = echoServer.getBindAddress();
            // Creates local port forward and send data to the echo server through that forwarding channel
            try (PortForwarding portForwarding = session.createLocalPortForward(bindAddress.getHostName(), bindAddress.getPort(), 12345, dataConsumer)) {
                List<String> messages = new ArrayList<>();
                for (int i = 0; i < 10; i++) {
                    String msg = "Testing" + i;
                    portForwarding.write(StandardCharsets.UTF_8.encode(msg));
                    portForwarding.write(StandardCharsets.UTF_8.encode("\n"));
                    messages.add(msg);
                }
                portForwarding.flush();
                Iterable<String> splits = Splitter.on("\n").omitEmptyStrings().split(received);
                Assert.assertEquals(messages, StreamSupport.stream(splits.spliterator(), false).collect(Collectors.toList()));
            }
            // After closing the port forwarding, the data consumer should have finished.
            Assert.assertTrue(finished.get());
        }
    } finally {
        echoServer.stopAndWait();
    }
}
Also used : InetSocketAddress(java.net.InetSocketAddress) ArrayList(java.util.ArrayList) ByteBuffer(java.nio.ByteBuffer) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) SSHSession(io.cdap.cdap.runtime.spi.ssh.SSHSession) ArrayList(java.util.ArrayList) List(java.util.List) PortForwarding(io.cdap.cdap.runtime.spi.ssh.PortForwarding) RemotePortForwarding(io.cdap.cdap.runtime.spi.ssh.RemotePortForwarding) Test(org.junit.Test)

Example 2 with PortForwarding

use of io.cdap.cdap.runtime.spi.ssh.PortForwarding in project cdap by caskdata.

the class SSHSessionTest method testForwardingOnSessionClose.

@Test
public void testForwardingOnSessionClose() throws Exception {
    EchoServer echoServer = new EchoServer();
    echoServer.startAndWait();
    try {
        SSHConfig sshConfig = getSSHConfig();
        AtomicBoolean finished = new AtomicBoolean(false);
        PortForwarding portForwarding;
        // Creates a SSH session
        try (SSHSession session = new DefaultSSHSession(sshConfig)) {
            InetSocketAddress bindAddress = echoServer.getBindAddress();
            // Creates a port forwarding and send some data
            BlockingQueue<String> received = new LinkedBlockingQueue<>();
            portForwarding = session.createLocalPortForward(bindAddress.getHostName(), bindAddress.getPort(), 12345, new PortForwarding.DataConsumer() {

                @Override
                public void received(ByteBuffer buffer) {
                    received.add(StandardCharsets.UTF_8.decode(buffer).toString());
                }

                @Override
                public void finished() {
                    finished.set(true);
                }
            });
            portForwarding.write(StandardCharsets.UTF_8.encode("Testing"));
            portForwarding.flush();
            Assert.assertEquals("Testing", received.poll(5, TimeUnit.SECONDS));
        }
        // After closing of the SSH session, the port forwarding should be closed as well
        Assert.assertTrue(finished.get());
        // Writing to a closed port forwarding should fails.
        try {
            portForwarding.write(StandardCharsets.UTF_8.encode("Testing 2"));
            Assert.fail("Expected failure when writing to closed PortForwarding");
        } catch (IOException e) {
        // expected
        }
    } finally {
        echoServer.stopAndWait();
    }
}
Also used : InetSocketAddress(java.net.InetSocketAddress) IOException(java.io.IOException) LinkedBlockingQueue(java.util.concurrent.LinkedBlockingQueue) ByteBuffer(java.nio.ByteBuffer) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) SSHSession(io.cdap.cdap.runtime.spi.ssh.SSHSession) PortForwarding(io.cdap.cdap.runtime.spi.ssh.PortForwarding) RemotePortForwarding(io.cdap.cdap.runtime.spi.ssh.RemotePortForwarding) Test(org.junit.Test)

Aggregations

PortForwarding (io.cdap.cdap.runtime.spi.ssh.PortForwarding)2 RemotePortForwarding (io.cdap.cdap.runtime.spi.ssh.RemotePortForwarding)2 SSHSession (io.cdap.cdap.runtime.spi.ssh.SSHSession)2 InetSocketAddress (java.net.InetSocketAddress)2 ByteBuffer (java.nio.ByteBuffer)2 AtomicBoolean (java.util.concurrent.atomic.AtomicBoolean)2 Test (org.junit.Test)2 IOException (java.io.IOException)1 ArrayList (java.util.ArrayList)1 List (java.util.List)1 LinkedBlockingQueue (java.util.concurrent.LinkedBlockingQueue)1