package org.apache.mina.filter.ssl;
   
   import static org.junit.Assert.fail;
   
   import java.io.IOException;
   import java.nio.ByteBuffer;
   import java.security.GeneralSecurityException;
   import java.security.KeyStore;
   import java.security.Security;
  import java.util.Deque;
  import java.util.concurrent.BlockingDeque;
  import java.util.concurrent.LinkedBlockingDeque;
  
  import javax.net.ssl.KeyManagerFactory;
  import javax.net.ssl.SSLContext;
  import javax.net.ssl.SSLEngine;
  import javax.net.ssl.SSLEngineResult;
  import javax.net.ssl.SSLException;
  import javax.net.ssl.SSLEngineResult.HandshakeStatus;
  import javax.net.ssl.SSLEngineResult.Status;
  import javax.net.ssl.TrustManagerFactory;
  
  import org.apache.mina.core.buffer.IoBuffer;
  import org.junit.Ignore;
  import org.junit.Test;
  
  public class SslEngineTest
  {
      private BlockingDeque<ByteBuffer> clientQueue = new LinkedBlockingDeque<>(); 
      private BlockingDeque<ByteBuffer> serverQueue = new LinkedBlockingDeque<>(); 
  
      private class Handshaker implements Runnable {
          private SSLEngine sslEngine;
          private ByteBuffer workBuffer;
          private ByteBuffer emptyBuffer= ByteBuffer.allocate(0);
          
          private void push(Deque<ByteBuffer> queue, ByteBuffer buffer) {
              ByteBuffer result = ByteBuffer.allocate(buffer.capacity());
              result.put(buffer);
              queue.addFirst(result);
          }
          
          public void run()
          {
              HandshakeStatus handshakeStatus = sslEngine.getHandshakeStatus();
              SSLEngineResult result;
  
              try
              {
                  while (handshakeStatus != HandshakeStatus.FINISHED) {
                      switch (handshakeStatus)
                      {
                          case NEED_TASK:
                              break;
                              
                          case NEED_UNWRAP:
                              // The SSLEngine waits for some input.
                              // We may have received too few data (TCP fragmentation)
                              // 
                              ByteBuffer data = serverQueue.takeLast();
                              result = sslEngine.unwrap(data, workBuffer);
                              
                              while (result.getStatus()  == SSLEngineResult.Status.BUFFER_UNDERFLOW) {
                                  // We need more data, until then, wait.
                                  //ByteBuffer data = serverQueue.takeLast();
                                  result = sslEngine.unwrap(data, workBuffer);
                              }
                              
                              handshakeStatus = sslEngine.getHandshakeStatus();
                              break;
  
                          case NEED_WRAP:
                          case NOT_HANDSHAKING:
                              result = sslEngine.wrap(emptyBuffer, workBuffer);
      
                              workBuffer.flip();
                              
                              if (workBuffer.hasRemaining()) {
                                  push(clientQueue, workBuffer);
                                  workBuffer.clear();
                              }
                              
                              handshakeStatus = result.getHandshakeStatus();
                              
                              break;
      
                          case FINISHED:
                          
                      }
                  }
              }
              catch ( SSLException e )
              {
                  // TODO Auto-generated catch block
                  e.printStackTrace();
              }
              catch ( InterruptedException e )
              {
                  // TODO Auto-generated catch block
                 e.printStackTrace();
             }
         }
         
         public Handshaker(SSLEngine sslEngine) {
             this.sslEngine = sslEngine;
             int packetBufferSize = sslEngine.getSession().getPacketBufferSize();
             workBuffer = ByteBuffer.allocate(packetBufferSize);
         }
     }
     
     /** A JVM independant KEY_MANAGER_FACTORY algorithm */
     private static final String KEY_MANAGER_FACTORY_ALGORITHM;
 
     static {
         String algorithm = Security.getProperty("ssl.KeyManagerFactory.algorithm");
         if (algorithm == null) {
             algorithm = KeyManagerFactory.getDefaultAlgorithm();
         }
 
         KEY_MANAGER_FACTORY_ALGORITHM = algorithm;
     }
 
     /** App data buffer for the client SSLEngine*/
     private IoBuffer inNetBufferClient;
 
     /** Net data buffer for the client SSLEngine */
     private IoBuffer outNetBufferClient;
 
     /** App data buffer for the server SSLEngine */
     private IoBuffer inNetBufferServer;
 
     /** Net data buffer for the server SSLEngine */
     private IoBuffer outNetBufferServer;
     
     private final IoBuffer emptyBuffer = IoBuffer.allocate(0);
 
 
     private static SSLContext createSSLContext() throws IOException, GeneralSecurityException {
         char[] passphrase = "password".toCharArray();
 
         SSLContext ctx = SSLContext.getInstance("TLS");
         KeyManagerFactory kmf = KeyManagerFactory.getInstance(KEY_MANAGER_FACTORY_ALGORITHM);
         TrustManagerFactory tmf = TrustManagerFactory.getInstance(KEY_MANAGER_FACTORY_ALGORITHM);
 
         KeyStore ks = KeyStore.getInstance("JKS");
         KeyStore ts = KeyStore.getInstance("JKS");
 
         ks.load(SslTest.class.getResourceAsStream("keystore.sslTest"), passphrase);
         ts.load(SslTest.class.getResourceAsStream("truststore.sslTest"), passphrase);
 
         kmf.init(ks, passphrase);
         tmf.init(ts);
         ctx.init(kmf.getKeyManagers(), tmf.getTrustManagers(), null);
 
         return ctx;
     }
 
     
     /**
      * Decrypt the incoming buffer and move the decrypted data to an
      * application buffer.
      */
     private SSLEngineResult unwrap(SSLEngine sslEngine, IoBuffer inBuffer, IoBuffer outBuffer) throws SSLException {
         // We first have to create the application buffer if it does not exist
         if (outBuffer == null) {
             outBuffer = IoBuffer.allocate(inBuffer.remaining());
         } else {
             // We already have one, just add the new data into it
             outBuffer.expand(inBuffer.remaining());
         }
 
         SSLEngineResult res;
         Status status;
         HandshakeStatus localHandshakeStatus;
 
         do {
             // Decode the incoming data
             res = sslEngine.unwrap(inBuffer.buf(), outBuffer.buf());
             status = res.getStatus();
 
             // We can be processing the Handshake
             localHandshakeStatus = res.getHandshakeStatus();
 
             if (status == SSLEngineResult.Status.BUFFER_OVERFLOW) {
                 // We have to grow the target buffer, it's too small.
                 // Then we can call the unwrap method again
                 int newCapacity = sslEngine.getSession().getApplicationBufferSize();
                 
                 if (inBuffer.remaining() >= newCapacity) {
                     // The buffer is already larger than the max buffer size suggested by the SSL engine.
                     // Raising it any more will not make sense and it will end up in an endless loop. Throwing an error is safer
                     throw new SSLException("SSL buffer overflow");
                 }
 
                 inBuffer.expand(newCapacity);
                 continue;
             }
         } while (((status == SSLEngineResult.Status.OK) || (status == SSLEngineResult.Status.BUFFER_OVERFLOW))
                 && ((localHandshakeStatus == SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING) || 
                         (localHandshakeStatus == SSLEngineResult.HandshakeStatus.NEED_UNWRAP)));
 
         return res;
     }
 
     
     private SSLEngineResult.Status unwrapHandshake(SSLEngine sslEngine, IoBuffer appBuffer, IoBuffer netBuffer) throws SSLException {
         // Prepare the net data for reading.
         if ((appBuffer == null) || !appBuffer.hasRemaining()) {
             // Need more data.
             return SSLEngineResult.Status.BUFFER_UNDERFLOW;
         }
 
         SSLEngineResult res = unwrap(sslEngine, appBuffer, netBuffer);
         HandshakeStatus handshakeStatus = res.getHandshakeStatus();
 
         //checkStatus(res);
 
         // If handshake finished, no data was produced, and the status is still
         // ok, try to unwrap more
         if ((handshakeStatus == SSLEngineResult.HandshakeStatus.FINISHED)
                 && (res.getStatus() == SSLEngineResult.Status.OK)
                 && appBuffer.hasRemaining()) {
             res = unwrap(sslEngine, appBuffer, netBuffer);
 
             // prepare to be written again
             if (appBuffer.hasRemaining()) {
                 appBuffer.compact();
             } else {
                 appBuffer.free();
                 appBuffer = null;
             }
         } else {
             // prepare to be written again
             if (appBuffer.hasRemaining()) {
                 appBuffer.compact();
             } else {
                 appBuffer.free();
                 appBuffer = null;
             }
         }
 
         return res.getStatus();
     }
 
     
     /* no qualifier */boolean isInboundDone(SSLEngine sslEngine) {
         return sslEngine == null || sslEngine.isInboundDone();
     }
 
     
     /* no qualifier */boolean isOutboundDone(SSLEngine sslEngine) {
         return sslEngine == null || sslEngine.isOutboundDone();
     }
 
     
     /**
      * Perform any handshaking processing.
      */
     /* no qualifier */HandshakeStatus handshake(SSLEngine sslEngine, IoBuffer appBuffer, IoBuffer netBuffer ) throws SSLException {
         SSLEngineResult result;
         HandshakeStatus handshakeStatus = sslEngine.getHandshakeStatus();
 
         for (;;) {
             switch (handshakeStatus) {
             case FINISHED:
                 //handshakeComplete = true;
                 return handshakeStatus;
 
             case NEED_TASK:
                 //handshakeStatus = doTasks();
                 break;
 
             case NEED_UNWRAP:
                 // we need more data read
                 SSLEngineResult.Status status = unwrapHandshake(sslEngine, appBuffer, netBuffer);
                 handshakeStatus = sslEngine.getHandshakeStatus();
 
                 return handshakeStatus;
 
             case NEED_WRAP:
                 result = sslEngine.wrap(emptyBuffer.buf(), netBuffer.buf());
 
                 while ( result.getStatus() == SSLEngineResult.Status.BUFFER_OVERFLOW ) {
                     netBuffer.capacity(netBuffer.capacity() << 1);
                     netBuffer.limit(netBuffer.capacity());
 
                     result = sslEngine.wrap(emptyBuffer.buf(), netBuffer.buf());
                 }
 
                 netBuffer.flip();
                 return result.getHandshakeStatus();
 
             case NOT_HANDSHAKING:
                 result = sslEngine.wrap(emptyBuffer.buf(), netBuffer.buf());
 
                 while ( result.getStatus() == SSLEngineResult.Status.BUFFER_OVERFLOW ) {
                     netBuffer.capacity(netBuffer.capacity() << 1);
                     netBuffer.limit(netBuffer.capacity());
 
                     result = sslEngine.wrap(emptyBuffer.buf(), netBuffer.buf());
                 }
 
                 netBuffer.flip();
                 handshakeStatus = result.getHandshakeStatus();
                 return handshakeStatus;
 
             default:
                 throw new IllegalStateException("error");
             }
         }
     }
 
     
     /**
      * Do all the outstanding handshake tasks in the current Thread.
      */
     private SSLEngineResult.HandshakeStatus doTasks(SSLEngine sslEngine) {
         /*
          * We could run this in a separate thread, but I don't see the need for
          * this when used from SSLFilter. Use thread filters in MINA instead?
          */
         Runnable runnable;
         while ((runnable = sslEngine.getDelegatedTask()) != null) {
             //Thread thread = new Thread(runnable);
             //thread.start();
             runnable.run();
         }
         return sslEngine.getHandshakeStatus();
     }
     
     
     private HandshakeStatus handshake(SSLEngine sslEngine, HandshakeStatus expected, 
         IoBuffer inBuffer, IoBuffer outBuffer, boolean dumpBuffer) throws SSLException {
         HandshakeStatus handshakeStatus = handshake(sslEngine, inBuffer, outBuffer);
 
         if ( handshakeStatus != expected) {
             fail();
         }
         
         if (dumpBuffer) {
             System.out.println("Message:" + outBuffer);
         }
         
         return handshakeStatus;
     }
 
     
     @Test
     @Ignore
     public void testSSL() throws Exception {
         // Initialise the client SSLEngine
         SSLContext sslContextClient = createSSLContext();
         SSLEngine sslEngineClient = sslContextClient.createSSLEngine();
         int packetBufferSize = sslEngineClient.getSession().getPacketBufferSize();
         inNetBufferClient = IoBuffer.allocate(packetBufferSize).setAutoExpand(true);
         outNetBufferClient = IoBuffer.allocate(packetBufferSize).setAutoExpand(true);
         
         sslEngineClient.setUseClientMode(true);
 
         // Initialise the Server SSLEngine
         SSLContext sslContextServer = createSSLContext();
         SSLEngine sslEngineServer = sslContextServer.createSSLEngine();
         packetBufferSize = sslEngineServer.getSession().getPacketBufferSize();
         inNetBufferServer = IoBuffer.allocate(packetBufferSize).setAutoExpand(true);
         outNetBufferServer = IoBuffer.allocate(packetBufferSize).setAutoExpand(true);
         
         sslEngineServer.setUseClientMode(false);
         
         Handshaker handshakerClient = new Handshaker( sslEngineClient );
         Handshaker handshakerServer = new Handshaker( sslEngineServer );
         
         handshakerServer.run();
 
         HandshakeStatus handshakeStatusClient = sslEngineClient.getHandshakeStatus();
         HandshakeStatus handshakeStatusServer = sslEngineServer.getHandshakeStatus();
 
     // <<< Server
         // Start the server
         handshakeStatusServer = handshake(sslEngineServer, HandshakeStatus.NEED_UNWRAP, 
             null, outNetBufferServer, false);
         
     // >>> Client
         // Now start the client, which will generate a CLIENT_HELLO,
         // stored into the outNetBufferClient
         handshakeStatusClient = handshake(sslEngineClient, HandshakeStatus.NEED_UNWRAP, 
             null, outNetBufferClient, true);
         
     // <<< Server
         // Process the CLIENT_HELLO on the server
         handshakeStatusServer = handshake(sslEngineServer, HandshakeStatus.NEED_TASK, 
             outNetBufferClient, outNetBufferServer, false);
 
         // Process the tasks on the server, prepare the SERVER_HELLO message
         handshakeStatusServer = doTasks(sslEngineServer);
         
         // We should be ready to generate the SERVER_HELLO message
         if ( handshakeStatusServer != HandshakeStatus.NEED_WRAP) {
             fail();
         }
         
         // Get the SERVER_HELLO message, with all the associated messages
         // ([Certificate], [ServerKeyExchange], [CertificateRequest], ServerHelloDone)
         outNetBufferServer.clear();
         handshakeStatusServer = handshake(sslEngineServer, HandshakeStatus.NEED_UNWRAP, 
             null, outNetBufferServer, true);
         
     // >>> Client
         // Process the SERVER_HELLO message on the client
         handshakeStatusClient = handshake(sslEngineClient, HandshakeStatus.NEED_TASK, 
             outNetBufferServer, inNetBufferClient, false);
     
         // Prepare the client response
         handshakeStatusClient = doTasks(sslEngineClient);
     
         // We should get back the Client messages ([Certificate],
         // ClientKeyExchange, [CertificateVerify])
         if ( handshakeStatusClient != HandshakeStatus.NEED_WRAP) {
             fail();
         }
     
         // Generate the [Certificate], ClientKeyExchange, [CertificateVerify] messages
         outNetBufferClient.clear();
         handshakeStatusClient = handshake(sslEngineClient, HandshakeStatus.NEED_WRAP, 
             null, outNetBufferClient, true);
         
     // <<< Server
         // Process the CLIENT_KEY_EXCHANGE on the server
         outNetBufferServer.clear();
         handshakeStatusServer = handshake(sslEngineServer, HandshakeStatus.NEED_TASK, 
             outNetBufferClient, outNetBufferServer, false);
 
         // Do the controls
         handshakeStatusServer = doTasks(sslEngineServer);
         
         // The server is waiting for more
         if ( handshakeStatusServer != HandshakeStatus.NEED_UNWRAP) {
             fail();
         }
 
     // >>> Client
         // The CHANGE_CIPHER_SPEC message generation
         outNetBufferClient.clear();
         handshakeStatusClient = handshake(sslEngineClient, HandshakeStatus.NEED_WRAP, 
             null, outNetBufferClient, true);
 
     // <<< Server
         // Process the CHANGE_CIPHER_SPEC on the server
         outNetBufferServer.clear();
         handshakeStatusServer = handshake(sslEngineServer, HandshakeStatus.NEED_UNWRAP, 
             outNetBufferClient, outNetBufferServer, false);
 
     // >>> Client
         // Generate the FINISHED message on thee client
         outNetBufferClient.clear();
         handshakeStatusClient = handshake(sslEngineClient, HandshakeStatus.NEED_UNWRAP, 
             null, outNetBufferClient, true);
 
     // <<< Server
         // Process the client FINISHED message
         outNetBufferServer.clear();
         handshakeStatusServer = handshake(sslEngineServer, HandshakeStatus.NEED_WRAP, 
             outNetBufferClient, outNetBufferServer, false);
 
         // Generate the CHANGE_CIPHER_SPEC message on the server
         handshakeStatusServer = handshake(sslEngineServer, HandshakeStatus.NEED_WRAP, 
             null, outNetBufferServer, true);
         
     // >>> Client
         // Process the server CHANGE_SCIPHER_SPEC message on the client
         outNetBufferClient.clear();
         handshakeStatusClient = handshake(sslEngineClient, HandshakeStatus.NEED_UNWRAP, 
             outNetBufferServer, outNetBufferClient, false);
 
     // <<< Server
         // Generate the server FINISHED message
         outNetBufferServer.clear();
         handshakeStatusServer = handshake(sslEngineServer, HandshakeStatus.FINISHED, 
             null, outNetBufferServer, true);
 
     // >>> Client
         // Process the server FINISHED message on the client
         outNetBufferClient.clear();
         handshakeStatusClient = handshake(sslEngineClient, HandshakeStatus.NOT_HANDSHAKING, 
             outNetBufferServer, outNetBufferClient, false);
     }
 }