1   /*
2    *  Licensed to the Apache Software Foundation (ASF) under one
3    *  or more contributor license agreements.  See the NOTICE file
4    *  distributed with this work for additional information
5    *  regarding copyright ownership.  The ASF licenses this file
6    *  to you under the Apache License, Version 2.0 (the
7    *  "License"); you may not use this file except in compliance
8    *  with the License.  You may obtain a copy of the License at
9    *
10   *    http://www.apache.org/licenses/LICENSE-2.0
11   *
12   *  Unless required by applicable law or agreed to in writing,
13   *  software distributed under the License is distributed on an
14   *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15   *  KIND, either express or implied.  See the License for the
16   *  specific language governing permissions and limitations
17   *  under the License.
18   *
19   */
20  package org.apache.mina.filter.stream;
21  
22  import java.io.ByteArrayInputStream;
23  import java.io.IOException;
24  import java.io.InputStream;
25  import java.net.InetSocketAddress;
26  import java.net.SocketAddress;
27  import java.security.MessageDigest;
28  import java.util.LinkedList;
29  import java.util.Queue;
30  import java.util.Random;
31  import java.util.concurrent.CountDownLatch;
32  import java.util.concurrent.TimeUnit;
33  
34  import junit.framework.TestCase;
35  
36  import org.apache.mina.common.DefaultWriteRequest;
37  import org.apache.mina.common.DummySession;
38  import org.apache.mina.common.IdleStatus;
39  import org.apache.mina.common.IoBuffer;
40  import org.apache.mina.common.IoFutureListener;
41  import org.apache.mina.common.IoHandlerAdapter;
42  import org.apache.mina.common.IoSession;
43  import org.apache.mina.common.WriteFuture;
44  import org.apache.mina.common.WriteRequest;
45  import org.apache.mina.common.IoFilter.NextFilter;
46  import org.apache.mina.transport.socket.nio.NioSocketAcceptor;
47  import org.apache.mina.transport.socket.nio.NioSocketConnector;
48  import org.apache.mina.util.AvailablePortFinder;
49  import org.easymock.AbstractMatcher;
50  import org.easymock.MockControl;
51  
52  /**
53   * Tests {@link StreamWriteFilter}.
54   *
55   * @author The Apache MINA Project (dev@mina.apache.org)
56   * @version $Rev: 600461 $, $Date: 2007-12-03 02:55:52 -0700 (Mon, 03 Dec 2007) $
57   */
58  public class StreamWriteFilterTest extends TestCase {
59      private MockControl mockNextFilter;
60  
61      private IoSession session;
62  
63      private NextFilter nextFilter;
64      private static final byte[] BUF = new byte[0];
65  
66      @Override
67      protected void setUp() throws Exception {
68          super.setUp();
69  
70          /*
71           * Create the mocks.
72           */
73          session = new DummySession();
74          mockNextFilter = MockControl.createControl(NextFilter.class);
75          nextFilter = (NextFilter) mockNextFilter.getMock();
76      }
77  
78      @Override
79      protected void tearDown() throws Exception {
80          assertFalse(session.containsAttribute(StreamWriteFilter.CURRENT_STREAM));
81          assertFalse(session
82                  .containsAttribute(StreamWriteFilter.CURRENT_WRITE_REQUEST));
83          assertFalse(session
84                  .containsAttribute(StreamWriteFilter.WRITE_REQUEST_QUEUE));
85          super.tearDown();
86      }
87  
88      public void testWriteEmptyStream() throws Exception {
89          StreamWriteFilter filter = new StreamWriteFilter();
90  
91          InputStream stream = new ByteArrayInputStream(BUF);
92          WriteRequest writeRequest = new DefaultWriteRequest(stream,
93                  new DummyWriteFuture());
94  
95          /*
96           * Record expectations
97           */
98          nextFilter.messageSent(session, writeRequest);
99  
100         /*
101          * Replay.
102          */
103         mockNextFilter.replay();
104 
105         filter.filterWrite(nextFilter, session, writeRequest);
106 
107         /*
108          * Verify.
109          */
110         mockNextFilter.verify();
111 
112         assertTrue(writeRequest.getFuture().isWritten());
113     }
114 
115     /**
116      * Tests that the filter just passes objects which aren't InputStreams
117      * through to the next filter.
118      *
119      * @throws Exception when something goes wrong
120      */
121     public void testWriteNonStreamMessage() throws Exception {
122         StreamWriteFilter filter = new StreamWriteFilter();
123 
124         Object message = new Object();
125         WriteRequest writeRequest = new DefaultWriteRequest(message,
126                 new DummyWriteFuture());
127 
128         /*
129          * Record expectations
130          */
131         nextFilter.filterWrite(session, writeRequest);
132         nextFilter.messageSent(session, writeRequest);
133 
134         /*
135          * Replay.
136          */
137         mockNextFilter.replay();
138 
139         filter.filterWrite(nextFilter, session, writeRequest);
140         filter.messageSent(nextFilter, session, writeRequest);
141 
142         /*
143          * Verify.
144          */
145         mockNextFilter.verify();
146     }
147 
148     /**
149      * Tests when the contents of the stream fits into one write buffer.
150      *
151      * @throws Exception when something goes wrong
152      */
153     public void testWriteSingleBufferStream() throws Exception {
154         StreamWriteFilter filter = new StreamWriteFilter();
155 
156         byte[] data = new byte[] { 1, 2, 3, 4 };
157 
158         InputStream stream = new ByteArrayInputStream(data);
159         WriteRequest writeRequest = new DefaultWriteRequest(stream,
160                 new DummyWriteFuture());
161 
162         /*
163          * Record expectations
164          */
165         nextFilter.filterWrite(session, new DefaultWriteRequest(IoBuffer
166                 .wrap(data)));
167         mockNextFilter.setMatcher(new WriteRequestMatcher());
168         nextFilter.messageSent(session, writeRequest);
169 
170         /*
171          * Replay.
172          */
173         mockNextFilter.replay();
174 
175         filter.filterWrite(nextFilter, session, writeRequest);
176         filter.messageSent(nextFilter, session, writeRequest);
177 
178         /*
179          * Verify.
180          */
181         mockNextFilter.verify();
182 
183         assertTrue(writeRequest.getFuture().isWritten());
184     }
185 
186     /**
187      * Tests when the contents of the stream doesn't fit into one write buffer.
188      *
189      * @throws Exception when something goes wrong
190      */
191     public void testWriteSeveralBuffersStream() throws Exception {
192         StreamWriteFilter filter = new StreamWriteFilter();
193         filter.setWriteBufferSize(4);
194 
195         byte[] data = new byte[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };
196         byte[] chunk1 = new byte[] { 1, 2, 3, 4 };
197         byte[] chunk2 = new byte[] { 5, 6, 7, 8 };
198         byte[] chunk3 = new byte[] { 9, 10 };
199 
200         InputStream stream = new ByteArrayInputStream(data);
201         WriteRequest writeRequest = new DefaultWriteRequest(stream,
202                 new DummyWriteFuture());
203 
204         WriteRequest chunk1Request = new DefaultWriteRequest(IoBuffer
205                 .wrap(chunk1));
206         WriteRequest chunk2Request = new DefaultWriteRequest(IoBuffer
207                 .wrap(chunk2));
208         WriteRequest chunk3Request = new DefaultWriteRequest(IoBuffer
209                 .wrap(chunk3));
210 
211         /*
212          * Record expectations
213          */
214         nextFilter.filterWrite(session, chunk1Request);
215         mockNextFilter.setMatcher(new WriteRequestMatcher());
216         nextFilter.filterWrite(session, chunk2Request);
217         nextFilter.filterWrite(session, chunk3Request);
218         nextFilter.messageSent(session, writeRequest);
219 
220         /*
221          * Replay.
222          */
223         mockNextFilter.replay();
224 
225         filter.filterWrite(nextFilter, session, writeRequest);
226         filter.messageSent(nextFilter, session, chunk1Request);
227         filter.messageSent(nextFilter, session, chunk2Request);
228         filter.messageSent(nextFilter, session, chunk3Request);
229 
230         /*
231          * Verify.
232          */
233         mockNextFilter.verify();
234 
235         assertTrue(writeRequest.getFuture().isWritten());
236     }
237 
238     public void testWriteWhileWriteInProgress() throws Exception {
239         StreamWriteFilter filter = new StreamWriteFilter();
240 
241         Queue<WriteRequest> queue = new LinkedList<WriteRequest>();
242         InputStream stream = new ByteArrayInputStream(new byte[5]);
243 
244         /*
245          * Make up the situation.
246          */
247         session.setAttribute(StreamWriteFilter.CURRENT_STREAM, stream);
248         session.setAttribute(StreamWriteFilter.WRITE_REQUEST_QUEUE, queue);
249 
250         /*
251          * Replay.  (We recorded *nothing* because nothing should occur.)
252          */
253         mockNextFilter.replay();
254 
255         WriteRequest wr = new DefaultWriteRequest(new Object(),
256                 new DummyWriteFuture());
257         filter.filterWrite(nextFilter, session, wr);
258         assertEquals(1, queue.size());
259         assertSame(wr, queue.poll());
260 
261         /*
262          * Verify.
263          */
264         mockNextFilter.verify();
265 
266         session.removeAttribute(StreamWriteFilter.CURRENT_STREAM);
267         session.removeAttribute(StreamWriteFilter.WRITE_REQUEST_QUEUE);
268     }
269 
270     public void testWritesWriteRequestQueueWhenFinished() throws Exception {
271         StreamWriteFilter filter = new StreamWriteFilter();
272 
273         WriteRequest wrs[] = new WriteRequest[] {
274                 new DefaultWriteRequest(new Object(), new DummyWriteFuture()),
275                 new DefaultWriteRequest(new Object(), new DummyWriteFuture()),
276                 new DefaultWriteRequest(new Object(), new DummyWriteFuture()) };
277         Queue<WriteRequest> queue = new LinkedList<WriteRequest>();
278         queue.add(wrs[0]);
279         queue.add(wrs[1]);
280         queue.add(wrs[2]);
281         InputStream stream = new ByteArrayInputStream(BUF);
282 
283         /*
284          * Make up the situation.
285          */
286         session.setAttribute(StreamWriteFilter.CURRENT_STREAM, stream);
287         session.setAttribute(StreamWriteFilter.CURRENT_WRITE_REQUEST,
288                 new DefaultWriteRequest(stream));
289         session.setAttribute(StreamWriteFilter.WRITE_REQUEST_QUEUE, queue);
290 
291         /*
292          * Record expectations
293          */
294         nextFilter.filterWrite(session, wrs[0]);
295         nextFilter.filterWrite(session, wrs[1]);
296         nextFilter.filterWrite(session, wrs[2]);
297         nextFilter.messageSent(session, new DefaultWriteRequest(stream));
298         mockNextFilter.setMatcher(new WriteRequestMatcher());
299 
300         /*
301          * Replay.
302          */
303         mockNextFilter.replay();
304 
305         filter.messageSent(nextFilter, session, new DefaultWriteRequest(
306                 new Object()));
307         assertEquals(0, queue.size());
308 
309         /*
310          * Verify.
311          */
312         mockNextFilter.verify();
313     }
314 
315     /**
316      * Tests that {@link StreamWriteFilter#setWriteBufferSize(int)} checks the
317      * specified size.
318      */
319     public void testSetWriteBufferSize() {
320         StreamWriteFilter filter = new StreamWriteFilter();
321 
322         try {
323             filter.setWriteBufferSize(0);
324             fail("0 writeBuferSize specified. IllegalArgumentException expected.");
325         } catch (IllegalArgumentException iae) {
326         }
327 
328         try {
329             filter.setWriteBufferSize(-100);
330             fail("Negative writeBuferSize specified. IllegalArgumentException expected.");
331         } catch (IllegalArgumentException iae) {
332         }
333 
334         filter.setWriteBufferSize(1);
335         assertEquals(1, filter.getWriteBufferSize());
336         filter.setWriteBufferSize(1024);
337         assertEquals(1024, filter.getWriteBufferSize());
338     }
339 
340     public void testWriteUsingSocketTransport() throws Exception {
341         NioSocketAcceptor acceptor = new NioSocketAcceptor();
342         acceptor.setReuseAddress(true);
343         SocketAddress address = new InetSocketAddress("localhost",
344                 AvailablePortFinder.getNextAvailable());
345 
346         NioSocketConnector connector = new NioSocketConnector();
347 
348         FixedRandomInputStream stream = new FixedRandomInputStream(
349                 4 * 1024 * 1024);
350 
351         SenderHandler sender = new SenderHandler(stream);
352         ReceiverHandler receiver = new ReceiverHandler(stream.size);
353 
354         acceptor.setHandler(sender);
355         connector.setHandler(receiver);
356         
357         acceptor.bind(address);
358         connector.connect(address);
359         sender.latch.await();
360         receiver.latch.await();
361 
362         acceptor.dispose();
363 
364         assertEquals(stream.bytesRead, receiver.bytesRead);
365         assertEquals(stream.size, receiver.bytesRead);
366         byte[] expectedMd5 = stream.digest.digest();
367         byte[] actualMd5 = receiver.digest.digest();
368         assertEquals(expectedMd5.length, actualMd5.length);
369         for (int i = 0; i < expectedMd5.length; i++) {
370             assertEquals(expectedMd5[i], actualMd5[i]);
371         }
372     }
373 
374     private static class FixedRandomInputStream extends InputStream {
375         long size;
376 
377         long bytesRead = 0;
378 
379         Random random = new Random();
380 
381         MessageDigest digest;
382 
383         private FixedRandomInputStream(long size) throws Exception {
384             this.size = size;
385             digest = MessageDigest.getInstance("MD5");
386         }
387 
388         @Override
389         public int read() throws IOException {
390             if (isAllWritten()) {
391                 return -1;
392             }
393             bytesRead++;
394             byte b = (byte) random.nextInt(255);
395             digest.update(b);
396             return b;
397         }
398 
399         public long getBytesRead() {
400             return bytesRead;
401         }
402 
403         public long getSize() {
404             return size;
405         }
406 
407         public boolean isAllWritten() {
408             return bytesRead >= size;
409         }
410     }
411 
412     private static class SenderHandler extends IoHandlerAdapter {
413         final CountDownLatch latch = new CountDownLatch(1);
414 
415         InputStream inputStream;
416 
417         StreamWriteFilter streamWriteFilter = new StreamWriteFilter();
418 
419         private SenderHandler(InputStream inputStream) {
420             this.inputStream = inputStream;
421         }
422 
423         @Override
424         public void sessionCreated(IoSession session) throws Exception {
425             super.sessionCreated(session);
426             session.getFilterChain().addLast("codec", streamWriteFilter);
427         }
428 
429         @Override
430         public void sessionOpened(IoSession session) throws Exception {
431             session.write(inputStream);
432         }
433 
434         @Override
435         public void exceptionCaught(IoSession session, Throwable cause)
436                 throws Exception {
437             latch.countDown();
438         }
439 
440         @Override
441         public void sessionClosed(IoSession session) throws Exception {
442             latch.countDown();
443         }
444 
445         @Override
446         public void sessionIdle(IoSession session, IdleStatus status)
447                 throws Exception {
448             latch.countDown();
449         }
450 
451         @Override
452         public void messageSent(IoSession session, Object message)
453                 throws Exception {
454             if (message == inputStream) {
455                 latch.countDown();
456             }
457         }
458     }
459 
460     private static class ReceiverHandler extends IoHandlerAdapter {
461         final CountDownLatch latch = new CountDownLatch(1);
462 
463         long bytesRead = 0;
464 
465         long size = 0;
466 
467         MessageDigest digest;
468 
469         private ReceiverHandler(long size) throws Exception {
470             this.size = size;
471             digest = MessageDigest.getInstance("MD5");
472         }
473 
474         @Override
475         public void sessionCreated(IoSession session) throws Exception {
476             super.sessionCreated(session);
477 
478             session.getConfig().setIdleTime(IdleStatus.READER_IDLE, 5);
479         }
480 
481         @Override
482         public void sessionIdle(IoSession session, IdleStatus status)
483                 throws Exception {
484             session.close();
485         }
486 
487         @Override
488         public void exceptionCaught(IoSession session, Throwable cause)
489                 throws Exception {
490             latch.countDown();
491         }
492 
493         @Override
494         public void sessionClosed(IoSession session) throws Exception {
495             latch.countDown();
496         }
497 
498         @Override
499         public void messageReceived(IoSession session, Object message)
500                 throws Exception {
501             IoBuffer buf = (IoBuffer) message;
502             while (buf.hasRemaining()) {
503                 digest.update(buf.get());
504                 bytesRead++;
505             }
506             if (bytesRead >= size) {
507                 session.close();
508             }
509         }
510     }
511 
512     public static class WriteRequestMatcher extends AbstractMatcher {
513         @Override
514         protected boolean argumentMatches(Object expected, Object actual) {
515             if (expected instanceof WriteRequest
516                     && actual instanceof WriteRequest) {
517                 WriteRequest w1 = (WriteRequest) expected;
518                 WriteRequest w2 = (WriteRequest) actual;
519 
520                 return w1.getMessage().equals(w2.getMessage())
521                         && w1.getFuture().isWritten() == w2.getFuture()
522                         .isWritten();
523             }
524             return super.argumentMatches(expected, actual);
525         }
526     }
527 
528     private static class DummyWriteFuture implements WriteFuture {
529         private boolean written;
530 
531         public boolean isWritten() {
532             return written;
533         }
534 
535         public void setWritten() {
536             this.written = true;
537         }
538 
539         public IoSession getSession() {
540             return null;
541         }
542 
543         public Object getLock() {
544             return this;
545         }
546 
547         public void join() {
548         }
549 
550         public boolean join(long timeoutInMillis) {
551             return true;
552         }
553 
554         public boolean isReady() {
555             return true;
556         }
557 
558         public WriteFuture addListener(IoFutureListener<?> listener) {
559             return this;
560         }
561 
562         public WriteFuture removeListener(IoFutureListener<?> listener) {
563             return this;
564         }
565 
566         public WriteFuture await() throws InterruptedException {
567             return this;
568         }
569 
570         public boolean await(long timeout, TimeUnit unit)
571                 throws InterruptedException {
572             return true;
573         }
574 
575         public boolean await(long timeoutMillis) throws InterruptedException {
576             return true;
577         }
578 
579         public WriteFuture awaitUninterruptibly() {
580             return this;
581         }
582 
583         public boolean awaitUninterruptibly(long timeout, TimeUnit unit) {
584             return true;
585         }
586 
587         public boolean awaitUninterruptibly(long timeoutMillis) {
588             return true;
589         }
590 
591         public Throwable getException() {
592             return null;
593         }
594 
595         public void setException(Throwable cause) {
596             throw new IllegalStateException();
597         }
598     }
599 }