sshd-core/src/test/java/org/apache/sshd/WindowAdjustTest.java - mina-sshd - Git at Google

 /*
  * 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 org.apache.sshd;

 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.util.Collection;
 import java.util.Deque;
 import java.util.EnumSet;
 import java.util.LinkedList;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;

 import org.apache.sshd.client.SshClient;
 import org.apache.sshd.client.channel.ClientChannel;
 import org.apache.sshd.client.channel.ClientChannelEvent;
 import org.apache.sshd.client.session.ClientSession;
 import org.apache.sshd.common.io.IoInputStream;
 import org.apache.sshd.common.io.IoOutputStream;
 import org.apache.sshd.common.io.WritePendingException;
 import org.apache.sshd.common.util.buffer.Buffer;
 import org.apache.sshd.common.util.buffer.ByteArrayBuffer;
 import org.apache.sshd.common.util.io.IoUtils;
 import org.apache.sshd.common.util.io.NoCloseOutputStream;
 import org.apache.sshd.common.util.logging.AbstractLoggingBean;
 import org.apache.sshd.common.util.threads.ThreadUtils;
 import org.apache.sshd.server.Environment;
 import org.apache.sshd.server.ExitCallback;
 import org.apache.sshd.server.SshServer;
 import org.apache.sshd.server.channel.ChannelSession;
 import org.apache.sshd.server.command.AsyncCommand;
 import org.apache.sshd.server.keyprovider.SimpleGeneratorHostKeyProvider;
 import org.apache.sshd.util.test.BaseTestSupport;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.FixMethodOrder;
 import org.junit.Test;
 import org.junit.runners.MethodSorters;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 /**
  * <p>
  * This test simulates heavy traffic coming from the server towards the client making sure the traffic does not get
  * stuck. Especially if the server receives window adjust message while it tries to transfer all the data.
  * </p>
  * {@link AsyncInPendingWrapper} in this test serves as a handler for {@link WritePendingException}, which can occur
  * when sending too many messages one after another.
  */
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 public class WindowAdjustTest extends BaseTestSupport {

     public static final byte END_FILE = '#';
     public static final int BIG_MSG_SEND_COUNT = 10000;

     private SshServer sshServer;
     private int port;

     public WindowAdjustTest() {
         super();
     }

     @Before
     public void setUp() throws Exception {
         sshServer = setupTestServer();

         byte[] msg = IoUtils.toByteArray(getClass().getResourceAsStream("/big-msg.txt"));
         sshServer.setShellFactory(
                 channel -> new FloodingAsyncCommand(msg, BIG_MSG_SEND_COUNT, END_FILE));

         sshServer.setKeyPairProvider(new SimpleGeneratorHostKeyProvider());
         sshServer.start();
         port = sshServer.getPort();
     }

     @After
     public void tearDown() throws Exception {
         if (sshServer != null) {
             sshServer.stop();
             sshServer.close(true);
         }
     }

     @Test(timeout = 6L * 60L * 1000L)
     public void testTrafficHeavyLoad() throws Exception {
         try (SshClient client = setupTestClient()) {
             client.start();

             try (ClientSession session = client.connect(getCurrentTestName(), TEST_LOCALHOST, port)
                     .verify(CONNECT_TIMEOUT).getSession()) {
                 session.addPasswordIdentity(getCurrentTestName());
                 session.auth().verify(AUTH_TIMEOUT);

                 try (ClientChannel channel = session.createShellChannel()) {
                     channel.setOut(new VerifyingOutputStream(channel, END_FILE));
                     channel.setErr(new NoCloseOutputStream(System.err));
                     channel.open().verify(OPEN_TIMEOUT);

                     Collection<ClientChannelEvent> result
                             = channel.waitFor(EnumSet.of(ClientChannelEvent.CLOSED), TimeUnit.MINUTES.toMillis(2L));
                     assertFalse("Timeout while waiting for channel closure", result.contains(ClientChannelEvent.TIMEOUT));
                 }
             } finally {
                 client.stop();
             }
         }
     }

     /**
      * Read all incoming data and if END_FILE symbol is detected, kill client session to end test
      */
     private static class VerifyingOutputStream extends OutputStream {
         private final Logger log;
         private final ClientChannel channel;
         private final byte eofSignal;

         VerifyingOutputStream(ClientChannel channel, final byte eofSignal) {
             this.log = LoggerFactory.getLogger(getClass());
             this.channel = channel;
             this.eofSignal = eofSignal;
         }

         @Override
         public void write(int b) throws IOException {
             if (channel.isClosed() || channel.isClosing()) {
                 throw new IOException("Channel (" + channel + ") is closing / closed on write single byte");
             }

             if (b == (eofSignal & 0xff)) {
                 log.info("Closing channel (" + channel + ") due to single byte EOF");
                 channel.close(true);
             }
         }

         @Override
         public void write(byte[] b, int off, int len) throws IOException {
             if (channel.isClosed() || channel.isClosing()) {
                 throw new IOException("Channel (" + channel + ") is closing / closed on write " + len + " bytes");
             }

             if (len <= 0) {
                 return;
             }

             int lastPos = off + len - 1;
             if ((b[lastPos] & 0xFF) == (eofSignal & 0xFF)) {
                 log.info("Closing channel (" + channel + ") due to last byte EOF");
                 channel.close(true);
             }
         }
     }

     public static final class FloodingAsyncCommand extends AbstractLoggingBean implements AsyncCommand {
         private static final AtomicInteger POOL_COUNT = new AtomicInteger(0);

         private final AtomicReference<ExecutorService> executorHolder = new AtomicReference<>();
         private final AtomicReference<Future<?>> futureHolder = new AtomicReference<>();

         private AsyncInPendingWrapper pendingWrapper;
         private byte[] msg;
         private int sendCount;
         private byte eofSignal;

         public FloodingAsyncCommand(final byte[] msg, final int sendCount, final byte eofSignal) {
             this.msg = msg;
             this.sendCount = sendCount;
             this.eofSignal = eofSignal;
         }

         @Override
         public void setIoInputStream(IoInputStream in) {
             // ignored
         }

         @Override
         public void setIoOutputStream(IoOutputStream out) {
             pendingWrapper = new AsyncInPendingWrapper(out);
         }

         @Override
         public void setIoErrorStream(IoOutputStream err) {
             // ignored
         }

         @Override
         public void setInputStream(InputStream in) {
             // ignored
         }

         @Override
         public void setOutputStream(OutputStream out) {
             // ignored
         }

         @Override
         public void setErrorStream(OutputStream err) {
             // ignored
         }

         @Override
         public void setExitCallback(ExitCallback callback) {
             // ignored
         }

         @Override
         public void start(ChannelSession channel, Environment env) throws IOException {
             log.info("Starting");

             ExecutorService service
                     = ThreadUtils.newSingleThreadExecutor(getClass().getSimpleName() + "-" + POOL_COUNT.incrementAndGet());
             executorHolder.set(service);

             futureHolder.set(service.submit((Runnable) () -> {
                 log.info("Start heavy load sending " + sendCount + " messages of " + msg.length + " bytes");
                 for (int i = 0; i < sendCount; i++) {
                     try {
                         pendingWrapper.write(new ByteArrayBuffer(msg));
                     } catch (IOException e) {
                         log.error("Failed ({}) to send message #{}/{}: {}",
                                 e.getClass().getSimpleName(), i + 1, sendCount, e.getMessage());
                         throw new RuntimeException(e);
                     }
                 }
                 log.info("Sending EOF signal");

                 try {
                     pendingWrapper.write(new ByteArrayBuffer(new byte[] { eofSignal }));
                 } catch (IOException e) {
                     log.error("Failed ({}) to send EOF message after {} messages: {}",
                             e.getClass().getSimpleName(), sendCount, e.getMessage());
                     throw new RuntimeException(e);
                 }
             }));
             log.info("Started");
         }

         @Override
         public void destroy(ChannelSession channel) {
             log.info("Destroying");

             Future<?> future = futureHolder.getAndSet(null);
             if ((future != null) && (!future.isDone())) {
                 log.info("Cancelling");
                 future.cancel(true);
             }

             ExecutorService service = executorHolder.getAndSet(null);
             if ((service != null) && (!service.isShutdown())) {
                 log.info("Shutdown");
                 service.shutdownNow();
             }
         }
     }

     /**
      * Wrapper for asyncIn stream that catches Pending exception and queues the pending messages for later retry (send
      * after previous messages were fully transfered)
      */
     private static class AsyncInPendingWrapper extends AbstractLoggingBean {
         private IoOutputStream asyncIn;

         // Order has to be preserved for queued writes
         private final Deque<Buffer> pending = new LinkedList<>();

         AsyncInPendingWrapper(IoOutputStream out) {
             this.asyncIn = out;
         }

         public synchronized void write(Object msg) throws IOException {
             if ((asyncIn != null) && (!asyncIn.isClosed()) && (!asyncIn.isClosing())) {
                 Buffer byteBufferMsg = (Buffer) msg;
                 if (!pending.isEmpty()) {
                     queueRequest(byteBufferMsg);
                     return;
                 }

                 writeWithPendingDetection(byteBufferMsg, false);
             }
         }

         private void writeWithPendingDetection(Buffer msg, boolean wasPending) throws IOException {
             try {
                 asyncIn.writeBuffer(msg).addListener(future -> {
                     if (future.isWritten()) {
                         if (wasPending) {
                             pending.remove();
                         }

                         try {
                             writePendingIfAny();
                         } catch (IOException e) {
                             log.error("Failed ({}) to re-write pending: {}", e.getClass().getSimpleName(), e.getMessage());
                         }
                     } else {
                         Throwable t = future.getException();
                         log.warn("Failed to write message", t);
                     }
                 });
             } catch (final WritePendingException e) {
                 if (!wasPending) {
                     queueRequest(msg);
                 }
             }
         }

         private synchronized void writePendingIfAny() throws IOException {
             if (pending.peek() == null) {
                 return;
             }

             Buffer msg = pending.peek();
             writeWithPendingDetection(msg, true);
         }

         private void queueRequest(Buffer msg) {
             msg.rpos(0);
             pending.add(msg);
         }
     }
 }
	/*
	* 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 org.apache.sshd;

	import java.io.IOException;
	import java.io.InputStream;
	import java.io.OutputStream;
	import java.util.Collection;
	import java.util.Deque;
	import java.util.EnumSet;
	import java.util.LinkedList;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.Future;
	import java.util.concurrent.TimeUnit;
	import java.util.concurrent.atomic.AtomicInteger;
	import java.util.concurrent.atomic.AtomicReference;

	import org.apache.sshd.client.SshClient;
	import org.apache.sshd.client.channel.ClientChannel;
	import org.apache.sshd.client.channel.ClientChannelEvent;
	import org.apache.sshd.client.session.ClientSession;
	import org.apache.sshd.common.io.IoInputStream;
	import org.apache.sshd.common.io.IoOutputStream;
	import org.apache.sshd.common.io.WritePendingException;
	import org.apache.sshd.common.util.buffer.Buffer;
	import org.apache.sshd.common.util.buffer.ByteArrayBuffer;
	import org.apache.sshd.common.util.io.IoUtils;
	import org.apache.sshd.common.util.io.NoCloseOutputStream;
	import org.apache.sshd.common.util.logging.AbstractLoggingBean;
	import org.apache.sshd.common.util.threads.ThreadUtils;
	import org.apache.sshd.server.Environment;
	import org.apache.sshd.server.ExitCallback;
	import org.apache.sshd.server.SshServer;
	import org.apache.sshd.server.channel.ChannelSession;
	import org.apache.sshd.server.command.AsyncCommand;
	import org.apache.sshd.server.keyprovider.SimpleGeneratorHostKeyProvider;
	import org.apache.sshd.util.test.BaseTestSupport;
	import org.junit.After;
	import org.junit.Before;
	import org.junit.FixMethodOrder;
	import org.junit.Test;
	import org.junit.runners.MethodSorters;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	/**
	* <p>
	* This test simulates heavy traffic coming from the server towards the client making sure the traffic does not get
	* stuck. Especially if the server receives window adjust message while it tries to transfer all the data.
	* </p>
	* {@link AsyncInPendingWrapper} in this test serves as a handler for {@link WritePendingException}, which can occur
	* when sending too many messages one after another.
	*/
	@FixMethodOrder(MethodSorters.NAME_ASCENDING)
	public class WindowAdjustTest extends BaseTestSupport {

	public static final byte END_FILE = '#';
	public static final int BIG_MSG_SEND_COUNT = 10000;

	private SshServer sshServer;
	private int port;

	public WindowAdjustTest() {
	super();
	}

	@Before
	public void setUp() throws Exception {
	sshServer = setupTestServer();

	byte[] msg = IoUtils.toByteArray(getClass().getResourceAsStream("/big-msg.txt"));
	sshServer.setShellFactory(
	channel -> new FloodingAsyncCommand(msg, BIG_MSG_SEND_COUNT, END_FILE));

	sshServer.setKeyPairProvider(new SimpleGeneratorHostKeyProvider());
	sshServer.start();
	port = sshServer.getPort();
	}

	@After
	public void tearDown() throws Exception {
	if (sshServer != null) {
	sshServer.stop();
	sshServer.close(true);
	}
	}

	@Test(timeout = 6L * 60L * 1000L)
	public void testTrafficHeavyLoad() throws Exception {
	try (SshClient client = setupTestClient()) {
	client.start();

	try (ClientSession session = client.connect(getCurrentTestName(), TEST_LOCALHOST, port)
	.verify(CONNECT_TIMEOUT).getSession()) {
	session.addPasswordIdentity(getCurrentTestName());
	session.auth().verify(AUTH_TIMEOUT);

	try (ClientChannel channel = session.createShellChannel()) {
	channel.setOut(new VerifyingOutputStream(channel, END_FILE));
	channel.setErr(new NoCloseOutputStream(System.err));
	channel.open().verify(OPEN_TIMEOUT);

	Collection<ClientChannelEvent> result
	= channel.waitFor(EnumSet.of(ClientChannelEvent.CLOSED), TimeUnit.MINUTES.toMillis(2L));
	assertFalse("Timeout while waiting for channel closure", result.contains(ClientChannelEvent.TIMEOUT));
	}
	} finally {
	client.stop();
	}
	}
	}

	/**
	* Read all incoming data and if END_FILE symbol is detected, kill client session to end test
	*/
	private static class VerifyingOutputStream extends OutputStream {
	private final Logger log;
	private final ClientChannel channel;
	private final byte eofSignal;

	VerifyingOutputStream(ClientChannel channel, final byte eofSignal) {
	this.log = LoggerFactory.getLogger(getClass());
	this.channel = channel;
	this.eofSignal = eofSignal;
	}

	@Override
	public void write(int b) throws IOException {
	if (channel.isClosed() \|\| channel.isClosing()) {
	throw new IOException("Channel (" + channel + ") is closing / closed on write single byte");
	}

	if (b == (eofSignal & 0xff)) {
	log.info("Closing channel (" + channel + ") due to single byte EOF");
	channel.close(true);
	}
	}

	@Override
	public void write(byte[] b, int off, int len) throws IOException {
	if (channel.isClosed() \|\| channel.isClosing()) {
	throw new IOException("Channel (" + channel + ") is closing / closed on write " + len + " bytes");
	}

	if (len <= 0) {
	return;
	}

	int lastPos = off + len - 1;
	if ((b[lastPos] & 0xFF) == (eofSignal & 0xFF)) {
	log.info("Closing channel (" + channel + ") due to last byte EOF");
	channel.close(true);
	}
	}
	}

	public static final class FloodingAsyncCommand extends AbstractLoggingBean implements AsyncCommand {
	private static final AtomicInteger POOL_COUNT = new AtomicInteger(0);

	private final AtomicReference<ExecutorService> executorHolder = new AtomicReference<>();
	private final AtomicReference<Future<?>> futureHolder = new AtomicReference<>();

	private AsyncInPendingWrapper pendingWrapper;
	private byte[] msg;
	private int sendCount;
	private byte eofSignal;

	public FloodingAsyncCommand(final byte[] msg, final int sendCount, final byte eofSignal) {
	this.msg = msg;
	this.sendCount = sendCount;
	this.eofSignal = eofSignal;
	}

	@Override
	public void setIoInputStream(IoInputStream in) {
	// ignored
	}

	@Override
	public void setIoOutputStream(IoOutputStream out) {
	pendingWrapper = new AsyncInPendingWrapper(out);
	}

	@Override
	public void setIoErrorStream(IoOutputStream err) {
	// ignored
	}

	@Override
	public void setInputStream(InputStream in) {
	// ignored
	}

	@Override
	public void setOutputStream(OutputStream out) {
	// ignored
	}

	@Override
	public void setErrorStream(OutputStream err) {
	// ignored
	}

	@Override
	public void setExitCallback(ExitCallback callback) {
	// ignored
	}

	@Override
	public void start(ChannelSession channel, Environment env) throws IOException {
	log.info("Starting");

	ExecutorService service
	= ThreadUtils.newSingleThreadExecutor(getClass().getSimpleName() + "-" + POOL_COUNT.incrementAndGet());
	executorHolder.set(service);

	futureHolder.set(service.submit((Runnable) () -> {
	log.info("Start heavy load sending " + sendCount + " messages of " + msg.length + " bytes");
	for (int i = 0; i < sendCount; i++) {
	try {
	pendingWrapper.write(new ByteArrayBuffer(msg));
	} catch (IOException e) {
	log.error("Failed ({}) to send message #{}/{}: {}",
	e.getClass().getSimpleName(), i + 1, sendCount, e.getMessage());
	throw new RuntimeException(e);
	}
	}
	log.info("Sending EOF signal");

	try {
	pendingWrapper.write(new ByteArrayBuffer(new byte[] { eofSignal }));
	} catch (IOException e) {
	log.error("Failed ({}) to send EOF message after {} messages: {}",
	e.getClass().getSimpleName(), sendCount, e.getMessage());
	throw new RuntimeException(e);
	}
	}));
	log.info("Started");
	}

	@Override
	public void destroy(ChannelSession channel) {
	log.info("Destroying");

	Future<?> future = futureHolder.getAndSet(null);
	if ((future != null) && (!future.isDone())) {
	log.info("Cancelling");
	future.cancel(true);
	}

	ExecutorService service = executorHolder.getAndSet(null);
	if ((service != null) && (!service.isShutdown())) {
	log.info("Shutdown");
	service.shutdownNow();
	}
	}
	}

	/**
	* Wrapper for asyncIn stream that catches Pending exception and queues the pending messages for later retry (send
	* after previous messages were fully transfered)
	*/
	private static class AsyncInPendingWrapper extends AbstractLoggingBean {
	private IoOutputStream asyncIn;

	// Order has to be preserved for queued writes
	private final Deque<Buffer> pending = new LinkedList<>();

	AsyncInPendingWrapper(IoOutputStream out) {
	this.asyncIn = out;
	}

	public synchronized void write(Object msg) throws IOException {
	if ((asyncIn != null) && (!asyncIn.isClosed()) && (!asyncIn.isClosing())) {
	Buffer byteBufferMsg = (Buffer) msg;
	if (!pending.isEmpty()) {
	queueRequest(byteBufferMsg);
	return;
	}

	writeWithPendingDetection(byteBufferMsg, false);
	}
	}

	private void writeWithPendingDetection(Buffer msg, boolean wasPending) throws IOException {
	try {
	asyncIn.writeBuffer(msg).addListener(future -> {
	if (future.isWritten()) {
	if (wasPending) {
	pending.remove();
	}

	try {
	writePendingIfAny();
	} catch (IOException e) {
	log.error("Failed ({}) to re-write pending: {}", e.getClass().getSimpleName(), e.getMessage());
	}
	} else {
	Throwable t = future.getException();
	log.warn("Failed to write message", t);
	}
	});
	} catch (final WritePendingException e) {
	if (!wasPending) {
	queueRequest(msg);
	}
	}
	}

	private synchronized void writePendingIfAny() throws IOException {
	if (pending.peek() == null) {
	return;
	}

	Buffer msg = pending.peek();
	writeWithPendingDetection(msg, true);
	}

	private void queueRequest(Buffer msg) {
	msg.rpos(0);
	pending.add(msg);
	}
	}
	}