Skip to content

Commit 2840565

Browse files
DavidiusdadiDavidiusdadi
committed
fixed problems with wrapped channels in general and improved SSLSocketChannel2
1 parent 8bfa6f9 commit 2840565

File tree

5 files changed

+79
-53
lines changed

5 files changed

+79
-53
lines changed

‎example/SSLServerExample.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
import java.nio.channels.SelectionKey;
77
import java.security.KeyStore;
88
import java.util.List;
9+
import java.util.concurrent.ExecutorService;
10+
import java.util.concurrent.Executors;
911

1012
import javax.net.ssl.KeyManagerFactory;
1113
import javax.net.ssl.SSLContext;
@@ -59,6 +61,8 @@ public static void main( String[] args ) throws Exception {
5961

6062
class SSLWebSocketServerFactory implements WebSocketServer.WebSocketServerFactory {
6163
private SSLContext sslcontext;
64+
private ExecutorService exec = Executors.newSingleThreadScheduledExecutor();
65+
6266
SSLWebSocketServerFactory( SSLContext sslContext ) {
6367
this.sslcontext = sslContext;
6468
}
@@ -67,7 +71,7 @@ class SSLWebSocketServerFactory implements WebSocketServer.WebSocketServerFactor
6771
public ByteChannel wrapChannel( SelectionKey c ) throws IOException {
6872
SSLEngine e = sslcontext.createSSLEngine();
6973
e.setUseClientMode( false );
70-
return new SSLSocketChannel2( c, e );
74+
return new SSLSocketChannel2( c, e, exec );
7175
}
7276

7377
@Override

‎src/org/java_websocket/SSLSocketChannel2.java‎

Lines changed: 45 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import java.nio.channels.SelectableChannel;
1414
import java.nio.channels.SelectionKey;
1515
import java.nio.channels.SocketChannel;
16+
import java.util.concurrent.ExecutorService;
1617

1718
import javax.net.ssl.SSLEngine;
1819
import javax.net.ssl.SSLEngineResult;
@@ -26,6 +27,8 @@
2627
public class SSLSocketChannel2 implements ByteChannel, WrappedByteChannel {
2728
private static ByteBuffer emptybuffer = ByteBuffer.allocate( 0 );
2829

30+
private ExecutorService exec;
31+
2932
/** raw payload incomming */
3033
private ByteBuffer inData;
3134
/** encrypted data outgoing */
@@ -39,12 +42,13 @@ public class SSLSocketChannel2 implements ByteChannel, WrappedByteChannel {
3942
private SSLEngineResult res;
4043
private SSLEngine sslEngine;
4144

42-
public SSLSocketChannel2( SelectionKey key , SSLEngine sslEngine ) throws IOException {
45+
public SSLSocketChannel2( SelectionKey key , SSLEngine sslEngine , ExecutorService exec ) throws IOException {
4346
this.sc = (SocketChannel) key.channel();
4447
this.key = key;
4548
this.sslEngine = sslEngine;
49+
this.exec = exec;
4650

47-
key.interestOps( key.interestOps() | SelectionKey.OP_WRITE );
51+
this.key.interestOps( key.interestOps() | SelectionKey.OP_WRITE );
4852

4953
sslEngine.setEnableSessionCreation( true );
5054
SSLSession session = sslEngine.getSession();
@@ -55,21 +59,18 @@ public SSLSocketChannel2( SelectionKey key , SSLEngine sslEngine ) throws IOExce
5559
inCrypt.flip();
5660
outCrypt.flip();
5761

58-
wrap( emptybuffer );
62+
sc.write( wrap( emptybuffer ) );
5963

6064
processHandshake();
6165
}
6266

6367
private boolean processHandshake() throws IOException {
6468
if( res.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_UNWRAP ) {
65-
if( !inCrypt.hasRemaining() )
66-
inCrypt.clear();
69+
inCrypt.compact();
6770
sc.read( inCrypt );
6871
inCrypt.flip();
6972
unwrap();
7073
if( res.getHandshakeStatus() != SSLEngineResult.HandshakeStatus.FINISHED ) {
71-
// if( !outData.hasRemaining() )
72-
// outData.clear();
7374
sc.write( wrap( emptybuffer ) );
7475
}
7576
} else if( res.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_WRAP ) {
@@ -82,8 +83,7 @@ private boolean processHandshake() throws IOException {
8283
}
8384

8485
private synchronized ByteBuffer wrap( ByteBuffer b ) throws SSLException {
85-
if( !outCrypt.hasRemaining() )
86-
outCrypt.clear();
86+
outCrypt.compact();
8787
res = sslEngine.wrap( b, outCrypt );
8888
outCrypt.flip();
8989
return outCrypt;
@@ -94,7 +94,6 @@ private synchronized ByteBuffer unwrap() throws SSLException {
9494
while ( inCrypt.hasRemaining() ) {
9595
res = sslEngine.unwrap( inCrypt, inData );
9696
if( res.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_TASK ) {
97-
// Task
9897
Runnable task;
9998
while ( ( task = sslEngine.getDelegatedTask() ) != null ) {
10099
task.run();
@@ -104,19 +103,18 @@ private synchronized ByteBuffer unwrap() throws SSLException {
104103
} else if( res.getStatus() == SSLEngineResult.Status.BUFFER_UNDERFLOW ) {
105104
break;
106105
} else if( res.getStatus() == SSLEngineResult.Status.BUFFER_OVERFLOW ) {
107-
assert ( false );
106+
throw new RuntimeException( "destenation buffer to small" );
108107
}
109108
}
110109
inData.flip();
111110
return inData;
112111
}
113112

114-
private void createBuffers( SSLSession session ) {
113+
protected void createBuffers( SSLSession session ) {
115114
int appBufferMax = session.getApplicationBufferSize();
116115
int netBufferMax = session.getPacketBufferSize();
117116

118-
inData = ByteBuffer.allocate( 65536 );
119-
117+
inData = ByteBuffer.allocate( appBufferMax );
120118
outCrypt = ByteBuffer.allocate( netBufferMax );
121119
inCrypt = ByteBuffer.allocate( netBufferMax );
122120
}
@@ -138,47 +136,24 @@ public int read( ByteBuffer dst ) throws IOException {
138136
int amount = 0, limit;
139137
// test if there was a buffer overflow in dst
140138
if( inData.hasRemaining() ) {
141-
limit = Math.min( inData.remaining(), dst.remaining() );
142-
for( int i = 0 ; i < limit ; i++ ) {
143-
dst.put( inData.get() );
144-
amount++;
145-
}
146-
return amount;
139+
return transfereTo( inData, dst );
147140
}
148141
// test if some bytes left from last read (e.g. BUFFER_UNDERFLOW)
149142
if( inCrypt.hasRemaining() ) {
150143
unwrap();
151-
inData.flip();
152-
limit = Math.min( inData.limit(), dst.remaining() );
153-
for( int i = 0 ; i < limit ; i++ ) {
154-
dst.put( inData.get() );
155-
amount++;
156-
}
157-
if( res.getStatus() != SSLEngineResult.Status.BUFFER_UNDERFLOW ) {
158-
inCrypt.clear();
159-
inCrypt.flip();
160-
return amount;
161-
}
144+
amount += transfereTo( inData, dst );
162145
}
163146
if( !inCrypt.hasRemaining() )
164147
inCrypt.clear();
165148
else
166149
inCrypt.compact();
167150

168151
if( sc.read( inCrypt ) == -1 ) {
169-
inCrypt.clear();
170-
inCrypt.flip();
171152
return -1;
172153
}
173154
inCrypt.flip();
174155
unwrap();
175-
// write in dst
176-
// inData.flip();
177-
limit = Math.min( inData.limit(), dst.remaining() );
178-
for( int i = 0 ; i < limit ; i++ ) {
179-
dst.put( inData.get() );
180-
amount++;
181-
}
156+
amount += transfereTo( inData, dst );
182157
return amount;
183158

184159
}
@@ -190,12 +165,11 @@ public boolean isConnected() {
190165
public void close() throws IOException {
191166
sslEngine.closeOutbound();
192167
sslEngine.getSession().invalidate();
193-
outCrypt.compact();
194168
int wr = sc.write( wrap( emptybuffer ) );
195169
sc.close();
196170
}
197171

198-
private boolean isHandShakeComplete(){
172+
private boolean isHandShakeComplete() {
199173
HandshakeStatus status = res.getHandshakeStatus();
200174
return status == SSLEngineResult.HandshakeStatus.FINISHED || status == SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING;
201175
}
@@ -231,7 +205,35 @@ public boolean isNeedWrite() {
231205
}
232206

233207
@Override
234-
public void write() throws IOException {
208+
public void writeMore() throws IOException {
235209
write( emptybuffer );
236210
}
211+
212+
@Override
213+
public boolean isNeedRead() {
214+
return inData.hasRemaining();
215+
}
216+
217+
@Override
218+
public int readMore( ByteBuffer dst ) {
219+
return transfereTo( inData, dst );
220+
}
221+
222+
private int transfereTo( ByteBuffer from, ByteBuffer to ) {
223+
int fremain = from.remaining();
224+
int toremain = to.remaining();
225+
if( fremain > toremain ) {
226+
// FIXME there should be a more efficient transfer method
227+
int limit = Math.min( fremain, toremain );
228+
for( int i = 0 ; i < limit ; i++ ) {
229+
to.put( from.get() );
230+
}
231+
return limit;
232+
} else {
233+
to.put( from );
234+
return fremain;
235+
}
236+
237+
}
238+
237239
}

‎src/org/java_websocket/SocketChannelIOHelper.java‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,14 @@
44
import java.nio.ByteBuffer;
55
import java.nio.channels.ByteChannel;
66

7-
import org.java_websocket.drafts.Draft;
8-
97
public class SocketChannelIOHelper {
108

11-
private static ByteBuffer emptybuffer = ByteBuffer.allocate( 0 );
12-
139
public static boolean read( final ByteBuffer buf, WebSocketImpl ws, ByteChannel channel ) throws IOException {
1410
buf.clear();
1511
int read = channel.read( buf );
1612
buf.flip();
1713

1814
if( read == -1 ) {
19-
Draft draft = ws.getDraft();
2015
ws.eot( null );
2116
return false;
2217
}
@@ -31,7 +26,7 @@ public static boolean batch( WebSocketImpl ws, ByteChannel sockchannel ) throws
3126
if( sockchannel instanceof WrappedByteChannel ) {
3227
WrappedByteChannel c = (WrappedByteChannel) sockchannel;
3328
if( c.isNeedWrite() ) {
34-
c.write();
29+
c.writeMore();
3530
return !c.isNeedWrite();
3631
}
3732
return true;
@@ -53,7 +48,7 @@ public static boolean batch( WebSocketImpl ws, ByteChannel sockchannel ) throws
5348
sockchannel.close();
5449
}
5550
}
56-
return true;
51+
return sockchannel instanceof WrappedByteChannel == true ? !( (WrappedByteChannel) sockchannel ).isNeedWrite() : true;
5752
}
5853

5954
}

‎src/org/java_websocket/WrappedByteChannel.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
package org.java_websocket;
22

33
import java.io.IOException;
4+
import java.nio.ByteBuffer;
45
import java.nio.channels.ByteChannel;
56

67
public interface WrappedByteChannel extends ByteChannel {
78
public boolean isNeedWrite();
8-
public void write() throws IOException;
9+
public void writeMore() throws IOException;
10+
11+
public boolean isNeedRead();
12+
public int readMore( ByteBuffer dst );
913
}

‎src/org/java_websocket/server/WebSocketServer.java‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import java.util.Collections;
1717
import java.util.HashSet;
1818
import java.util.Iterator;
19+
import java.util.LinkedList;
1920
import java.util.List;
2021
import java.util.Set;
2122
import java.util.concurrent.BlockingQueue;
@@ -27,6 +28,7 @@
2728
import org.java_websocket.WebSocketAdapter;
2829
import org.java_websocket.WebSocketFactory;
2930
import org.java_websocket.WebSocketImpl;
31+
import org.java_websocket.WrappedByteChannel;
3032
import org.java_websocket.drafts.Draft;
3133
import org.java_websocket.framing.CloseFrame;
3234
import org.java_websocket.handshake.ClientHandshake;
@@ -70,6 +72,7 @@ public abstract class WebSocketServer extends WebSocketAdapter implements Runnab
7072
private List<WebSocketWorker> decoders;
7173

7274
private BlockingQueue<WebSocketImpl> oqueue;
75+
private List<WebSocketImpl> iqueue;
7376
private BlockingQueue<ByteBuffer> buffers;
7477
private int queueinvokes = 0;
7578
private AtomicInteger queuesize = new AtomicInteger( 0 );
@@ -136,6 +139,7 @@ public WebSocketServer( InetSocketAddress address , int decodercount , List<Draf
136139
setAddress( address );
137140

138141
oqueue = new LinkedBlockingQueue<WebSocketImpl>();
142+
iqueue = new LinkedList<WebSocketImpl>();
139143

140144
decoders = new ArrayList<WebSocketWorker>( decodercount );
141145
buffers = new LinkedBlockingQueue<ByteBuffer>();
@@ -278,6 +282,11 @@ public void run() {
278282
conn.inQueue.put( buf );
279283
queue( conn );
280284
i.remove();
285+
if( conn.channel instanceof WrappedByteChannel ) {
286+
if( ( (WrappedByteChannel) conn.channel ).isNeedRead() ) {
287+
iqueue.add( conn );
288+
}
289+
}
281290
} else {
282291
pushBuffer( buf );
283292
}
@@ -297,14 +306,26 @@ public void run() {
297306
}
298307
}
299308
}
309+
while ( !iqueue.isEmpty() ) {
310+
conn = iqueue.remove( 0 );
311+
WrappedByteChannel c = ( (WrappedByteChannel) conn.channel );
312+
ByteBuffer buf = takeBuffer();
313+
buf.clear();
314+
c.readMore( buf );
315+
buf.flip();
316+
conn.inQueue.put( buf );
317+
queue( conn );
318+
if(c.isNeedRead())
319+
iqueue.add( conn );
320+
}
300321
} catch ( CancelledKeyException e ) {
301322
// an other thread may cancel the key
302323
} catch ( IOException ex ) {
303324
if( key != null )
304325
key.cancel();
305326
handleIOException( conn, ex );
306327
} catch ( InterruptedException e ) {
307-
return;
328+
return;// FIXME controlled shutdown
308329
}
309330
}
310331
} catch ( RuntimeException e ) {

0 commit comments

Comments
 (0)