Skip to content

Commit 18ddfcf

Browse files
TooTallNateTooTallNate
committed
Fixed glitch with WebSocketClient where CPU usage would spike at 100%. Closes TooTallNate#3.
1 parent 70e46c2 commit 18ddfcf

File tree

1 file changed

+106
-85
lines changed

1 file changed

+106
-85
lines changed

‎src/net/tootallnate/websocket/WebSocketClient.java‎

Lines changed: 106 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import java.net.InetSocketAddress;
55
import java.net.URI;
66
import java.nio.ByteBuffer;
7+
import java.nio.channels.ClosedChannelException;
78
import java.nio.channels.SelectionKey;
89
import java.nio.channels.Selector;
910
import java.nio.channels.SocketChannel;
@@ -38,17 +39,27 @@ public abstract class WebSocketClient implements Runnable, WebSocketListener {
3839
*/
3940
private WebSocket conn;
4041
/**
41-
* @author The Websocket mode this client is in.
42+
* The SocketChannel instance this client uses.
43+
*/
44+
private SocketChannel client;
45+
/**
46+
* The 'Selector' used to get event keys from the underlying socket.
47+
*/
48+
private Selector selector;
49+
/**
50+
* The Websocket mode this client is in.
4251
*/
4352
private Draft draft;
4453
/**
4554
* Number 1 used in handshake
4655
*/
4756
private int number1=0;
48-
/** Number 2 used in handshake
57+
/**
58+
* Number 2 used in handshake
4959
*/
5060
private int number2=0;
51-
/** Key3 used in handshake
61+
/**
62+
* Key3 used in handshake
5263
*/
5364
private byte[] key3=null;
5465
public static enum Draft{
@@ -106,26 +117,33 @@ public void send(String text) throws IOException {
106117

107118
// Runnable IMPLEMENTATION /////////////////////////////////////////////////
108119
public void run() {
109-
try {
110-
int port = uri.getPort();
111-
if (port == -1) {
112-
port = WebSocket.DEFAULT_PORT;
113-
}
120+
int port = uri.getPort();
121+
if (port == -1) {
122+
port = WebSocket.DEFAULT_PORT;
123+
}
114124

115-
// The WebSocket constructor expects a SocketChannel that is
116-
// non-blocking, and has a Selector attached to it.
117-
SocketChannel client = SocketChannel.open();
125+
// The WebSocket constructor expects a SocketChannel that is
126+
// non-blocking, and has a Selector attached to it.
127+
try {
128+
client = SocketChannel.open();
118129
client.configureBlocking(false);
119130
client.connect(new InetSocketAddress(uri.getHost(), port));
120131

121-
Selector selector = Selector.open();
132+
selector = Selector.open();
122133

123134
this.conn = new WebSocket(client, new LinkedBlockingQueue<ByteBuffer>(), this);
124-
client.register(selector, client.validOps());
135+
// At first, we're only interested in the 'CONNECT' keys.
136+
client.register(selector, SelectionKey.OP_CONNECT);
125137

126-
// Continuous loop that is only supposed to end when close is called
127-
while (selector.select(500) > 0) {
138+
} catch (IOException ex) {
139+
ex.printStackTrace();
140+
return;
141+
}
128142

143+
// Continuous loop that is only supposed to end when "close" is called.
144+
while (true) {
145+
try {
146+
selector.select();
129147
Set<SelectionKey> keys = selector.selectedKeys();
130148
Iterator<SelectionKey> i = keys.iterator();
131149

@@ -135,14 +153,16 @@ public void run() {
135153

136154
// When 'conn' has connected to the host
137155
if (key.isConnectable()) {
138-
156+
139157
// Ensure connection is finished
140158
if (client.isConnectionPending()) {
141159
client.finishConnect();
142160
}
143161

162+
// Now that we're connected, re-register for only 'READ' keys.
163+
client.register(selector, SelectionKey.OP_READ);
164+
144165
// Now send WebSocket client-side handshake
145-
Random r=new Random();
146166
String path = "/" + uri.getPath();
147167
String host = uri.getHost() + (port != WebSocket.DEFAULT_PORT ? ":" + port : "");
148168
String origin = null; // TODO: Make 'origin' configurable
@@ -151,17 +171,16 @@ public void run() {
151171
"Connection: Upgrade\r\n" +
152172
"Host: " + host + "\r\n" +
153173
"Origin: " + origin + "\r\n";
154-
if(this.draft==Draft.DRAFT76)
155-
{
156-
request+="Sec-WebSocket-Key1: " + this.generateKey() + "\r\n";
157-
request+="Sec-WebSocket-Key2: " + this.generateKey() + "\r\n";
158-
this.key3=new byte[8];
159-
r.nextBytes(this.key3);
160-
}
161-
//extraHeaders.toString() +
162-
request+="\r\n";
174+
if (this.draft == Draft.DRAFT76) {
175+
request += "Sec-WebSocket-Key1: " + this.generateKey() + "\r\n";
176+
request += "Sec-WebSocket-Key2: " + this.generateKey() + "\r\n";
177+
this.key3 = new byte[8];
178+
(new Random()).nextBytes(this.key3);
179+
}
180+
//extraHeaders.toString() +
181+
request+="\r\n";
163182
conn.socketChannel().write(ByteBuffer.wrap(request.getBytes(WebSocket.UTF8_CHARSET)));
164-
if(this.key3 !=null){
183+
if (this.key3 != null) {
165184
conn.socketChannel().write(ByteBuffer.wrap(this.key3));
166185
}
167186
}
@@ -171,44 +190,46 @@ public void run() {
171190
conn.handleRead();
172191
}
173192
}
193+
} catch (IOException ex) {
194+
ex.printStackTrace();
195+
} catch (NoSuchAlgorithmException ex) {
196+
ex.printStackTrace();
174197
}
175-
} catch (IOException ex) {
176-
ex.printStackTrace();
177-
} catch (NoSuchAlgorithmException e) {
178-
// TODO Auto-generated catch block
179-
e.printStackTrace();
180198
}
199+
200+
//System.err.println("WebSocketClient thread ended!");
181201
}
182-
private String generateKey(){
183-
Random r=new Random();
184-
long maxNumber=4294967295L;
185-
long spaces=r.nextInt(12)+1;
186-
int max=new Long(maxNumber/spaces).intValue();
187-
max=Math.abs(max);
188-
int number=r.nextInt(max)+1;
189-
if(this.number1==0){
190-
this.number1=number;
202+
203+
private String generateKey() {
204+
Random r = new Random();
205+
long maxNumber = 4294967295L;
206+
long spaces = r.nextInt(12) + 1;
207+
int max = new Long(maxNumber / spaces).intValue();
208+
max = Math.abs(max);
209+
int number = r.nextInt(max) + 1;
210+
if (this.number1 == 0) {
211+
this.number1 = number;
191212
}
192-
else{
193-
this.number2=number;
194-
}
195-
long product=number*spaces;
196-
String key=Long.toString(product);
197-
int numChars=r.nextInt(12);
198-
for (int i=0;i<numChars;i++){
199-
int position=r.nextInt(key.length());
200-
position=Math.abs(position);
201-
char randChar=(char)(r.nextInt(95)+33);
213+
else {
214+
this.number2 = number;
215+
}
216+
long product = number * spaces;
217+
String key = Long.toString(product);
218+
int numChars = r.nextInt(12);
219+
for (int i=0; i < numChars; i++){
220+
int position = r.nextInt(key.length());
221+
position = Math.abs(position);
222+
char randChar = (char)(r.nextInt(95) + 33);
202223
//exclude numbers here
203-
if(randChar >= 48 && randChar <=57){
204-
randChar-=15;
224+
if(randChar >= 48 && randChar <= 57){
225+
randChar -= 15;
205226
}
206-
key=new StringBuilder(key).insert(position, randChar).toString();
227+
key = new StringBuilder(key).insert(position, randChar).toString();
207228
}
208-
for (int i=0;i<spaces;i++){
209-
int position=r.nextInt(key.length()-1)+1;
210-
position=Math.abs(position);
211-
key=new StringBuilder(key).insert(position,"\u0020").toString();
229+
for (int i = 0; i < spaces; i++){
230+
int position = r.nextInt(key.length() - 1) + 1;
231+
position = Math.abs(position);
232+
key = new StringBuilder(key).insert(position,"\u0020").toString();
212233
}
213234
return key;
214235
}
@@ -228,36 +249,36 @@ private String generateKey(){
228249
public boolean onHandshakeRecieved(WebSocket conn, String handshake,byte[] reply) throws IOException, NoSuchAlgorithmException {
229250
// TODO: Do some parsing of the returned handshake, and close connection
230251
// (return false) if we recieved anything unexpected.
231-
if(this.draft==Draft.DRAFT76){
232-
if(reply==null){
233-
return false;
234-
}
235-
byte[] challenge=new byte[]{
236-
(byte)( this.number1 >> 24 ),
237-
(byte)( (this.number1 << 8) >> 24 ),
238-
(byte)( (this.number1 << 16) >> 24 ),
239-
(byte)( (this.number1 << 24) >> 24 ),
240-
(byte)( this.number2 >> 24 ),
241-
(byte)( (this.number2 << 8) >> 24 ),
242-
(byte)( (this.number2 << 16) >> 24 ),
243-
(byte)( (this.number2 << 24) >> 24 ),
244-
this.key3[0],
245-
this.key3[1],
246-
this.key3[2],
247-
this.key3[3],
248-
this.key3[4],
249-
this.key3[5],
250-
this.key3[6],
251-
this.key3[7]
252-
};
253-
MessageDigest md5=MessageDigest.getInstance("MD5");
254-
byte[] expected=md5.digest(challenge);
255-
for(int i=0;i<reply.length;i++){
256-
if(expected[i]!=reply[i]){
252+
if(this.draft == Draft.DRAFT76) {
253+
if (reply == null){
254+
return false;
255+
}
256+
byte[] challenge = new byte[] {
257+
(byte)( this.number1 >> 24 ),
258+
(byte)( (this.number1 << 8) >> 24 ),
259+
(byte)( (this.number1 << 16) >> 24 ),
260+
(byte)( (this.number1 << 24) >> 24 ),
261+
(byte)( this.number2 >> 24 ),
262+
(byte)( (this.number2 << 8) >> 24 ),
263+
(byte)( (this.number2 << 16) >> 24 ),
264+
(byte)( (this.number2 << 24) >> 24 ),
265+
this.key3[0],
266+
this.key3[1],
267+
this.key3[2],
268+
this.key3[3],
269+
this.key3[4],
270+
this.key3[5],
271+
this.key3[6],
272+
this.key3[7]
273+
};
274+
MessageDigest md5 = MessageDigest.getInstance("MD5");
275+
byte[] expected = md5.digest(challenge);
276+
for(int i = 0; i < reply.length; i++){
277+
if (expected[i] != reply[i]) {
257278
return false;
258279
}
259280
}
260-
}
281+
}
261282
return true;
262283
}
263284

0 commit comments

Comments
 (0)