package org.bouncycastle.crypto.tls.test;

import java.security.SecureRandom;

import junit.framework.TestCase;

import org.bouncycastle.crypto.tls.DTLSClientProtocol;
import org.bouncycastle.crypto.tls.DTLSServerProtocol;
import org.bouncycastle.crypto.tls.DTLSTransport;
import org.bouncycastle.crypto.tls.DatagramTransport;

public class DTLSProtocolTest extends TestCase {

    public void testClientServer() throws Exception {

        SecureRandom secureRandom = new SecureRandom();

        DTLSClientProtocol clientProtocol = new DTLSClientProtocol(secureRandom);
        DTLSServerProtocol serverProtocol = new DTLSServerProtocol(secureRandom);

        MockDatagramAssociation network = new MockDatagramAssociation(1500);

        ServerThread serverThread = new ServerThread(serverProtocol, network.getServer());
        serverThread.start();

        DatagramTransport clientTransport = network.getClient();

        clientTransport = new LoggingDatagramTransport(clientTransport, System.out);

        MockDTLSClient client = new MockDTLSClient();
        DTLSTransport dtlsClient = clientProtocol.connect(client, clientTransport);

        // byte[] data = new byte[64];
        // secureRandom.nextBytes(data);
        //
        // OutputStream output = clientProtocol.getOutputStream();
        // output.write(data);
        // output.close();
        //
        // byte[] echo = Streams.readAll(clientProtocol.getInputStream());

        dtlsClient.close();

        serverThread.join();

        // assertTrue(Arrays.areEqual(data, echo));
    }

    static class ServerThread extends Thread {
        private final DTLSServerProtocol serverProtocol;
        private final DatagramTransport serverTransport;

        ServerThread(DTLSServerProtocol serverProtocol, DatagramTransport serverTransport) {
            this.serverProtocol = serverProtocol;
            this.serverTransport = serverTransport;
        }

        public void run() {
            try {
                MockDTLSServer server = new MockDTLSServer();
                DTLSTransport dtlsServer = serverProtocol.accept(server, serverTransport);
                // Streams.pipeAll(serverProtocol.getInputStream(),
                // serverProtocol.getOutputStream());
//                byte[] buf = new byte[dtlsServer.getReceiveLimit()];
//                dtlsServer.receive(buf, 0, buf.length, 60000);
                dtlsServer.close();
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    }
}

@@ -933,32 +933,131 @@
 getClient();%0D%0A%0D%0A
+        clientTransport = new UnreliableDatagramTransport(clientTransport, secureRandom, 0, 0);%0D%0A%0D%0A
         clientTr
@@ -1638,11 +1638,15 @@
 ead.
-joi
+shutdow
 n();
@@ -1874,16 +1874,70 @@
 nsport;%0D
+%0A        private volatile boolean isShutdown = false;%0D
 %0A%0D%0A     
@@ -2365,156 +2365,134 @@
     
-// Streams.pipeAll(serverProtocol.getInputStream(),%0D%0A                // serverProtocol.getOutputStream());%0D%0A//                byte%5B%5D buf = new byte%5B
+byte%5B%5D buf = new byte%5BdtlsServer.getReceiveLimit()%5D;%0D%0A                while (!isShutdown) %7B%0D%0A                    int length = 
 dtls
@@ -2502,65 +2502,128 @@
 ver.
-getR
+r
 eceive
-Limit()%5D;%0D%0A//                dtlsServer.receive
+(buf, 0, buf.length, 1000);%0D%0A                    if (length %3E= 0) %7B%0D%0A                        serverTransport.send
 (buf
@@ -2631,27 +2631,58 @@
  0, 
-buf.
 length
-, 60000);
+);%0D%0A                    %7D%0D%0A                %7D
 %0D%0A  
@@ -2822,23 +2822,204 @@
 %0D%0A        %7D%0D
+%0A%0D%0A        void shutdown() throws InterruptedException %7B%0D%0A            if (!isShutdown) %7B%0D%0A                isShutdown = true;%0D%0A                this.join();%0D%0A            %7D%0D%0A        %7D%0D
 %0A    %7D%0D%0A%7D%0D%0A

Add proper shutdown and unreliable layer for testing