SecureNioChannel


1   package org.apache.tomcat.util.net;
2   
3   import java.io.IOException  ;
4   import java.nio.ByteBuffer  ;
5   import java.nio.channels.SelectionKey  ;
6   import java.nio.channels.SocketChannel  ;
7   import javax.net.ssl.SSLEngine;
8   import javax.net.ssl.SSLEngineResult;
9   import javax.net.ssl.SSLEngineResult.HandshakeStatus;
10  import javax.net.ssl.SSLEngineResult.Status;
11  import java.nio.channels.Selector  ;
12  
13  /**
14   * 
15   * Implementation of a secure socket channel
16   * @author Filip Hanik
17   * @version 1.0
18   */
19  
20  public class SecureNioChannel extends NioChannel  {
21      
22      protected ByteBuffer   netInBuffer;
23      protected ByteBuffer   netOutBuffer;
24      
25      protected SSLEngine sslEngine;
26      
27      protected boolean initHandshakeComplete = false;
28      protected HandshakeStatus initHandshakeStatus; //gets set by begin handshake
29      
30      protected boolean closed = false;
31      protected boolean closing = false;
32      
33      protected NioSelectorPool pool;
34      
35      public SecureNioChannel(SocketChannel   channel, SSLEngine engine, 
36                              ApplicationBufferHandler bufHandler, NioSelectorPool pool) throws IOException   {
37          super(channel,bufHandler);
38          this.sslEngine = engine;
39          int appBufSize = sslEngine.getSession().getApplicationBufferSize();
40          int netBufSize = sslEngine.getSession().getPacketBufferSize();
41          //allocate network buffers - TODO, add in optional direct non-direct buffers
42          if ( netInBuffer == null ) netInBuffer = ByteBuffer.allocateDirect(netBufSize);
43          if ( netOutBuffer == null ) netOutBuffer = ByteBuffer.allocateDirect(netBufSize);
44          
45          //selector pool for blocking operations
46          this.pool = pool;
47          
48          //ensure that the application has a large enough read/write buffers
49          //by doing this, we should not encounter any buffer overflow errors
50          bufHandler.expand(bufHandler.getReadBuffer(), appBufSize);
51          bufHandler.expand(bufHandler.getWriteBuffer(), appBufSize);
52          reset();
53      }
54      
55      public void reset(SSLEngine engine) throws IOException   {
56          this.sslEngine = engine;
57          reset();
58      }
59      public void reset() throws IOException   {
60          super.reset();
61          netOutBuffer.position(0);
62          netOutBuffer.limit(0);
63          netInBuffer.position(0);
64          netInBuffer.limit(0);
65          initHandshakeComplete = false;
66          closed = false;
67          closing = false;
68          //initiate handshake
69          sslEngine.beginHandshake();
70          initHandshakeStatus = sslEngine.getHandshakeStatus();
71      }
72      
73      public int getBufferSize() {
74          int size = super.getBufferSize();
75          size += netInBuffer!=null?netInBuffer.capacity():0;
76          size += netOutBuffer!=null?netOutBuffer.capacity():0;
77          return size;
78      }
79  
80      
81  //===========================================================================================    
82  //                  NIO SSL METHODS
83  //===========================================================================================
84      /**
85       * returns true if the network buffer has 
86       * been flushed out and is empty
87       * @return boolean
88       */
89      public boolean flush(Selector   s, long timeout) throws IOException   {
90          pool.write(netOutBuffer,this,s,timeout);
91          return !netOutBuffer.hasRemaining();
92      }
93      
94      /**
95       * Flushes the buffer to the network, non blocking
96       * @param buf ByteBuffer
97       * @return boolean true if the buffer has been emptied out, false otherwise
98       * @throws IOException
99       */
100     protected boolean flush(ByteBuffer   buf) throws IOException   {
101         int remaining = buf.remaining();
102         if ( remaining > 0 ) {
103             int written = sc.write(buf);
104             return written >= remaining;
105         }else {
106             return true;
107         }
108     }
109     
110     /**
111      * Performs SSL handshake, non blocking, but performs NEED_TASK on the same thread.<br>
112      * Hence, you should never call this method using your Acceptor thread, as you would slow down
113      * your system significantly.<br>
114      * The return for this operation is 0 if the handshake is complete and a positive value if it is not complete.
115      * In the event of a positive value coming back, reregister the selection key for the return values interestOps.
116      * @param read boolean - true if the underlying channel is readable
117      * @param write boolean - true if the underlying channel is writable
118      * @return int - 0 if hand shake is complete, otherwise it returns a SelectionKey interestOps value
119      * @throws IOException
120      */
121     public int handshake(boolean read, boolean write) throws IOException   {
122         if ( initHandshakeComplete ) return 0; //we have done our initial handshake
123         
124         if (!flush(netOutBuffer)) return SelectionKey.OP_WRITE; //we still have data to write
125         
126         SSLEngineResult handshake = null;
127         
128         while (!initHandshakeComplete) {
129             switch ( initHandshakeStatus ) {
130                 case NOT_HANDSHAKING: {
131                     //should never happen
132                     throw new IOException  ("NOT_HANDSHAKING during handshake");
133                 }
134                 case FINISHED: {
135                     //we are complete if we have delivered the last package
136                     initHandshakeComplete = !netOutBuffer.hasRemaining();
137                     //return 0 if we are complete, otherwise we still have data to write
138                     return initHandshakeComplete?0:SelectionKey.OP_WRITE; 
139                 }
140                 case NEED_WRAP: {
141                     //perform the wrap function
142                     handshake = handshakeWrap(write);
143                     if ( handshake.getStatus() == Status.OK ){
144                         if (initHandshakeStatus == HandshakeStatus.NEED_TASK) 
145                             initHandshakeStatus = tasks();
146                     } else {
147                         //wrap should always work with our buffers
148                         throw new IOException  ("Unexpected status:" + handshake.getStatus() + " during handshake WRAP.");
149                     }
150                     if ( initHandshakeStatus != HandshakeStatus.NEED_UNWRAP || (!flush(netOutBuffer)) ) {
151                         //should actually return OP_READ if we have NEED_UNWRAP
152                         return SelectionKey.OP_WRITE;
153                     }
154                     //fall down to NEED_UNWRAP on the same call, will result in a 
155                     //BUFFER_UNDERFLOW if it needs data
156                 }
157                 case NEED_UNWRAP: {
158                     //perform the unwrap function
159                     handshake = handshakeUnwrap(read);
160                     if ( handshake.getStatus() == Status.OK ) {
161                         if (initHandshakeStatus == HandshakeStatus.NEED_TASK) 
162                             initHandshakeStatus = tasks();
163                     } else if ( handshake.getStatus() == Status.BUFFER_UNDERFLOW ){
164                         //read more data, reregister for OP_READ
165                         return SelectionKey.OP_READ;
166                     } else {
167                         throw new IOException  ("Invalid handshake status:"+initHandshakeStatus+" during handshake UNWRAP.");
168                     }//switch
169                     break;
170                 }
171                 case NEED_TASK: {
172                     initHandshakeStatus = tasks();
173                     break;
174                 }
175                 default: throw new IllegalStateException  ("Invalid handshake status:"+initHandshakeStatus);
176             }//switch
177         }//while      
178         //return 0 if we are complete, otherwise reregister for any activity that 
179         //would cause this method to be called again.
180         return initHandshakeComplete?0:(SelectionKey.OP_WRITE|SelectionKey.OP_READ);
181     }
182     
183     /**
184      * Executes all the tasks needed on the same thread.
185      * @return HandshakeStatus
186      */
187     protected SSLEngineResult.HandshakeStatus tasks() {
188         Runnable   r = null;
189         while ( (r = sslEngine.getDelegatedTask()) != null) {
190             r.run();
191         }
192         return sslEngine.getHandshakeStatus();
193     }
194 
195     /**
196      * Performs the WRAP function
197      * @param doWrite boolean
198      * @return SSLEngineResult
199      * @throws IOException
200      */
201     protected SSLEngineResult handshakeWrap(boolean doWrite) throws IOException   {
202         //this should never be called with a network buffer that contains data
203         //so we can clear it here.
204         netOutBuffer.clear();
205         //perform the wrap
206         SSLEngineResult result = sslEngine.wrap(bufHandler.getWriteBuffer(), netOutBuffer);
207         //prepare the results to be written
208         netOutBuffer.flip();
209         //set the status
210         initHandshakeStatus = result.getHandshakeStatus();
211         //optimization, if we do have a writable channel, write it now
212         if ( doWrite ) flush(netOutBuffer);
213         return result;
214     }
215     
216     /**
217      * Perform handshake unwrap
218      * @param doread boolean
219      * @return SSLEngineResult
220      * @throws IOException
221      */
222     protected SSLEngineResult handshakeUnwrap(boolean doread) throws IOException   {
223         
224         if (netInBuffer.position() == netInBuffer.limit()) {
225             //clear the buffer if we have emptied it out on data
226             netInBuffer.clear();
227         }
228         if ( doread )  {
229             //if we have data to read, read it
230             int read = sc.read(netInBuffer);
231             if (read == -1) throw new IOException  ("EOF encountered during handshake.");
232         }        
233         SSLEngineResult result;
234         boolean cont = false;
235         //loop while we can perform pure SSLEngine data
236         do {
237             //prepare the buffer with the incoming data
238             netInBuffer.flip();
239             //call unwrap
240             result = sslEngine.unwrap(netInBuffer, bufHandler.getReadBuffer());
241             //compact the buffer, this is an optional method, wonder what would happen if we didn't
242             netInBuffer.compact();
243             //read in the status
244             initHandshakeStatus = result.getHandshakeStatus();
245             if ( result.getStatus() == SSLEngineResult.Status.OK &&
246                  result.getHandshakeStatus() == HandshakeStatus.NEED_TASK ) {
247                 //execute tasks if we need to
248                 initHandshakeStatus = tasks();
249             }
250             //perform another unwrap?
251             cont = result.getStatus() == SSLEngineResult.Status.OK &&
252                    initHandshakeStatus == HandshakeStatus.NEED_UNWRAP;
253         }while ( cont );
254         return result;
255     }
256     
257     /**
258      * Sends a SSL close message, will not physically close the connection here.<br>
259      * To close the connection, you could do something like
260      * <pre><code>
261      *   close();
262      *   while (isOpen() && !myTimeoutFunction()) Thread.sleep(25);
263      *   if ( isOpen() ) close(true); //forces a close if you timed out
264      * </code></pre>
265      * @throws IOException if an I/O error occurs
266      * @throws IOException if there is data on the outgoing network buffer and we are unable to flush it
267      * @todo Implement this java.io.Closeable method
268      */
269     public void close() throws IOException   {
270         if (closing) return;
271         closing = true;
272         sslEngine.closeOutbound();
273 
274         if (!flush(netOutBuffer)) {
275             throw new IOException  ("Remaining data in the network buffer, can't send SSL close message, force a close with close(true) instead");
276         }
277         //prep the buffer for the close message
278         netOutBuffer.clear();
279         //perform the close, since we called sslEngine.closeOutbound
280         SSLEngineResult handshake = sslEngine.wrap(getEmptyBuf(), netOutBuffer);
281         //we should be in a close state
282         if (handshake.getStatus() != SSLEngineResult.Status.CLOSED) {
283             throw new IOException  ("Invalid close state, will not send network data.");
284         }
285         //prepare the buffer for writing
286         netOutBuffer.flip();
287         //if there is data to be written
288         flush(netOutBuffer);
289 
290         //is the channel closed?
291         closed = (!netOutBuffer.hasRemaining() && (handshake.getHandshakeStatus() != HandshakeStatus.NEED_WRAP));
292     }
293 
294     /**
295      * Force a close, can throw an IOException
296      * @param force boolean
297      * @throws IOException
298      */
299     public void close(boolean force) throws IOException   {
300         try {
301             close();
302         }finally {
303             if ( force || closed ) {
304                 closed = true;
305                 sc.socket().close();
306                 sc.close();
307             }
308         }
309     }
310 
311     /**
312      * Reads a sequence of bytes from this channel into the given buffer.
313      *
314      * @param dst The buffer into which bytes are to be transferred
315      * @return The number of bytes read, possibly zero, or <tt>-1</tt> if the channel has reached end-of-stream
316      * @throws IOException If some other I/O error occurs
317      * @throws IllegalArgumentException if the destination buffer is different than bufHandler.getReadBuffer()
318      * @todo Implement this java.nio.channels.ReadableByteChannel method
319      */
320     public int read(ByteBuffer   dst) throws IOException   {
321         //if we want to take advantage of the expand function, make sure we only use the ApplicationBufferHandler's buffers
322         if ( dst != bufHandler.getReadBuffer() ) throw new IllegalArgumentException  ("You can only read using the application read buffer provided by the handler.");
323         //are we in the middle of closing or closed?
324         if ( closing || closed) return -1;
325         //did we finish our handshake?
326         if (!initHandshakeComplete) throw new IllegalStateException  ("Handshake incomplete, you must complete handshake before reading data.");
327 
328         //read from the network
329         int netread = sc.read(netInBuffer);
330         //did we reach EOF? if so send EOF up one layer.
331         if (netread == -1) return -1;
332         
333         //the data read
334         int read = 0;
335         //the SSL engine result
336         SSLEngineResult unwrap;
337         do {
338             //prepare the buffer
339             netInBuffer.flip();
340             //unwrap the data
341             unwrap = sslEngine.unwrap(netInBuffer, dst);
342             //compact the buffer
343             netInBuffer.compact();
344             
345             if ( unwrap.getStatus()==Status.OK || unwrap.getStatus()==Status.BUFFER_UNDERFLOW ) {
346                 //we did receive some data, add it to our total
347                 read += unwrap.bytesProduced();
348                 //perform any tasks if needed
349                 if (unwrap.getHandshakeStatus() == HandshakeStatus.NEED_TASK) tasks();
350                 //if we need more network data, then bail out for now.
351                 if ( unwrap.getStatus() == Status.BUFFER_UNDERFLOW ) break;
352             }else {
353                 //here we should trap BUFFER_OVERFLOW and call expand on the buffer
354                 //for now, throw an exception, as we initialized the buffers
355                 //in the constructor
356                 throw new IOException  ("Unable to unwrap data, invalid status: " + unwrap.getStatus());
357             }
358         } while ( (netInBuffer.position() != 0)); //continue to unwrapping as long as the input buffer has stuff
359         return (read);
360     }
361 
362     /**
363      * Writes a sequence of bytes to this channel from the given buffer.
364      *
365      * @param src The buffer from which bytes are to be retrieved
366      * @return The number of bytes written, possibly zero
367      * @throws IOException If some other I/O error occurs
368      * @todo Implement this java.nio.channels.WritableByteChannel method
369      */
370     public int write(ByteBuffer   src) throws IOException   {
371         //make sure we can handle expand, and that we only use on buffer
372         if ( src != bufHandler.getWriteBuffer() ) throw new IllegalArgumentException  ("You can only write using the application write buffer provided by the handler.");
373         //are we closing or closed?
374         if ( closing || closed) throw new IOException  ("Channel is in closing state.");
375         
376         //the number of bytes written
377         int written = 0;
378         
379         if (!flush(netOutBuffer)) {
380             //we haven't emptied out the buffer yet
381             return written;
382         }
383 
384         /*
385          * The data buffer is empty, we can reuse the entire buffer.
386          */
387         netOutBuffer.clear();
388 
389         SSLEngineResult result = sslEngine.wrap(src, netOutBuffer);
390         written = result.bytesConsumed();
391         netOutBuffer.flip();
392 
393         if (result.getStatus() == Status.OK) {
394             if (result.getHandshakeStatus() == HandshakeStatus.NEED_TASK) tasks();
395         } else {
396             throw new IOException  ("Unable to wrap data, invalid engine state: " +result.getStatus());
397         }        
398 
399         //force a flush
400         flush(netOutBuffer);
401 
402         return written;
403     }
404     
405     /**
406      * Callback interface to be able to expand buffers
407      * when buffer overflow exceptions happen
408      */
409     public static interface ApplicationBufferHandler {
410         public ByteBuffer   expand(ByteBuffer   buffer, int remaining);
411         public ByteBuffer   getReadBuffer();
412         public ByteBuffer   getWriteBuffer();
413     }
414 
415     public ApplicationBufferHandler getBufHandler() {
416         return bufHandler;
417     }
418 
419     public boolean isInitHandshakeComplete() {
420         return initHandshakeComplete;
421     }
422 
423     public boolean isClosing() {
424         return closing;
425     }
426 
427     public SSLEngine getSslEngine() {
428         return sslEngine;
429     }
430 
431     public ByteBuffer   getEmptyBuf() {
432         return emptyBuf;
433     }
434 
435     public void setBufHandler(ApplicationBufferHandler bufHandler) {
436         this.bufHandler = bufHandler;
437     }
438     
439     public SocketChannel   getIOChannel() {
440         return sc;
441     }
442 
443 }
A to Z: JavaDoc & Examples Daily Java News & Articles Open Source Projects Open Source Codes Free Computer Books Remove Frame
Popular Tags