Skip to content

Commit 86b3c05

Browse files
DavidiusdadiDavidiusdadi
committed
made WebSocketClient work with WrappedByteChannels as well
1 parent c59f9c0 commit 86b3c05

3 files changed

Lines changed: 45 additions & 18 deletions

File tree

‎src/org/java_websocket/SocketChannelIOHelper.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,18 @@ public static boolean read( final ByteBuffer buf, WebSocketImpl ws, ByteChannel
1818
return read != 0;
1919
}
2020

21+
public static boolean readMore( final ByteBuffer buf, WebSocketImpl ws, WrappedByteChannel channel ) throws IOException {
22+
buf.clear();
23+
int read = channel.readMore( buf );
24+
buf.flip();
25+
26+
if( read == -1 ) {
27+
ws.eot( null );
28+
return false;
29+
}
30+
return channel.isNeedRead();
31+
}
32+
2133
/** Returns whether the whole outQueue has been flushed */
2234
public static boolean batch( WebSocketImpl ws, ByteChannel sockchannel ) throws IOException {
2335
ByteBuffer buffer = ws.outQueue.peek();

‎src/org/java_websocket/client/WebSocketClient.java‎

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.java_websocket.WebSocketAdapter;
2525
import org.java_websocket.WebSocketFactory;
2626
import org.java_websocket.WebSocketImpl;
27+
import org.java_websocket.WrappedByteChannel;
2728
import org.java_websocket.drafts.Draft;
2829
import org.java_websocket.drafts.Draft_10;
2930
import org.java_websocket.exceptions.InvalidHandshakeException;
@@ -69,7 +70,7 @@ public abstract class WebSocketClient extends WebSocketAdapter implements Runnab
6970

7071
private final Lock closelock = new ReentrantLock();
7172

72-
private Map<String, String> headers;
73+
private Map<String,String> headers;
7374

7475
WebSocketFactory wf = new WebSocketFactory() {
7576
@Override
@@ -98,10 +99,10 @@ public WebSocketClient( URI serverURI ) {
9899
* must call <var>connect</var> first to initiate the socket connection.
99100
*/
100101
public WebSocketClient( URI serverUri , Draft draft ) {
101-
this(serverUri, draft, null);
102+
this( serverUri, draft, null );
102103
}
103104

104-
public WebSocketClient( URI serverUri , Draft draft, Map<String, String> headers) {
105+
public WebSocketClient( URI serverUri , Draft draft , Map<String,String> headers ) {
105106
if( serverUri == null ) {
106107
throw new IllegalArgumentException();
107108
}
@@ -192,7 +193,7 @@ public void run() {
192193
interruptableRun();
193194

194195
try {
195-
if (selector != null) // if the initialization in <code>tryToConnect</code> fails, it could be null
196+
if( selector != null ) // if the initialization in <code>tryToConnect</code> fails, it could be null
196197
selector.close();
197198
} catch ( IOException e ) {
198199
onError( e );
@@ -234,7 +235,6 @@ private final void interruptableRun() {
234235
}
235236
SelectionKey key = null;
236237
selector.select();
237-
SocketChannelIOHelper.batch( conn, wrappedchannel );
238238
Set<SelectionKey> keys = selector.selectedKeys();
239239
Iterator<SelectionKey> i = keys.iterator();
240240
while ( i.hasNext() ) {
@@ -256,12 +256,28 @@ private final void interruptableRun() {
256256
conn.close( e ); // http error
257257
}
258258
}
259+
if( key.isWritable() ) {
260+
if( SocketChannelIOHelper.batch( conn, wrappedchannel ) ) {
261+
if( key.isValid() )
262+
key.interestOps( SelectionKey.OP_READ );
263+
} else {
264+
key.interestOps( SelectionKey.OP_READ | SelectionKey.OP_WRITE );
265+
}
266+
}
267+
}
268+
if( wrappedchannel instanceof WrappedByteChannel ) {
269+
WrappedByteChannel w = (WrappedByteChannel) wrappedchannel;
270+
if( w.isNeedRead() ) {
271+
while ( SocketChannelIOHelper.read( buff, conn, w ) ) {
272+
conn.decode( buff );
273+
}
274+
}
259275
}
260276
}
261277
} catch ( IOException e ) {
262278
onError( e );
263279
conn.close( CloseFrame.ABNORMAL_CLOSE );
264-
} catch ( RuntimeException e ) {
280+
} catch ( RuntimeException e ) {
265281
// this catch case covers internal errors only and indicates a bug in this websocket implementation
266282
onError( e );
267283
conn.eot( e );
@@ -278,7 +294,7 @@ private void finishConnect( SelectionKey key ) throws IOException , InvalidHands
278294
channel.finishConnect();
279295
}
280296
// Now that we're connected, re-register for only 'READ' keys.
281-
key.interestOps( SelectionKey.OP_READ );
297+
key.interestOps( SelectionKey.OP_READ | SelectionKey.OP_WRITE );
282298

283299
wrappedchannel = wf.wrapChannel( key );
284300
sendHandshake();
@@ -300,9 +316,9 @@ private void sendHandshake() throws IOException , InvalidHandshakeException , In
300316
HandshakeImpl1Client handshake = new HandshakeImpl1Client();
301317
handshake.setResourceDescriptor( path );
302318
handshake.put( "Host", host );
303-
if (headers != null) {
304-
for (Map.Entry<String, String> kv : headers.entrySet()) {
305-
handshake.put(kv.getKey(), kv.getValue());
319+
if( headers != null ) {
320+
for( Map.Entry<String,String> kv : headers.entrySet() ) {
321+
handshake.put( kv.getKey(), kv.getValue() );
306322
}
307323
}
308324
conn.startHandshake( handshake );
@@ -316,9 +332,9 @@ private void sendHandshake() throws IOException , InvalidHandshakeException , In
316332
* @return Returns '0 = CONNECTING', '1 = OPEN', '2 = CLOSING' or '3 = CLOSED'
317333
*/
318334
public int getReadyState() {
319-
if (conn == null) {
320-
return WebSocket.READY_STATE_CONNECTING;
321-
}
335+
if( conn == null ) {
336+
return WebSocket.READY_STATE_CONNECTING;
337+
}
322338
return conn.getReadyState();
323339
}
324340

@@ -371,6 +387,7 @@ public final void onWebsocketError( WebSocket conn, Exception ex ) {
371387

372388
@Override
373389
public final void onWriteDemand( WebSocket conn ) {
390+
channel.keyFor( selector ).interestOps( SelectionKey.OP_READ | SelectionKey.OP_WRITE );
374391
selector.wakeup();
375392
}
376393

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -310,13 +310,11 @@ public void run() {
310310
conn = iqueue.remove( 0 );
311311
WrappedByteChannel c = ( (WrappedByteChannel) conn.channel );
312312
ByteBuffer buf = takeBuffer();
313-
buf.clear();
314-
c.readMore( buf );
315-
buf.flip();
313+
if(SocketChannelIOHelper.readMore( buf, conn, c ))
314+
iqueue.add( conn );
316315
conn.inQueue.put( buf );
317316
queue( conn );
318-
if(c.isNeedRead())
319-
iqueue.add( conn );
317+
320318
}
321319
} catch ( CancelledKeyException e ) {
322320
// an other thread may cancel the key

0 commit comments

Comments
 (0)