Skip to content

Commit 219f54e

Browse files
marci4marci4
authored
Merge pull request TooTallNate#761 from marci4/MemoryManagement
Implements Memory Management
2 parents 64b7574 + 8741403 commit 219f54e

File tree

9 files changed

+440
-47
lines changed

9 files changed

+440
-47
lines changed

‎src/main/example/simplelogger.properties‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
org.slf4j.simpleLogger.logFile=System.out
2-
org.slf4j.simpleLogger.defaultLogLevel=error
2+
org.slf4j.simpleLogger.defaultLogLevel=trace
33
org.slf4j.simpleLogger.showDateTime=true
44
org.slf4j.simpleLogger.dateTimeFormat=yyyy-MM-dd HH:mm:ss.SSS
55
org.slf4j.simpleLogger.showThreadName=false

‎src/main/java/org/java_websocket/SSLSocketChannel.java‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -265,8 +265,15 @@ private boolean doHandshake() throws IOException {
265265
peerNetData.clear();
266266

267267
handshakeStatus = engine.getHandshakeStatus();
268-
while( handshakeStatus != SSLEngineResult.HandshakeStatus.FINISHED && handshakeStatus != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING ) {
268+
boolean handshakeComplete = false;
269+
while( !handshakeComplete) {
269270
switch(handshakeStatus) {
271+
case FINISHED:
272+
handshakeComplete = !this.peerNetData.hasRemaining();
273+
if (handshakeComplete)
274+
return true;
275+
socketChannel.write(this.peerNetData);
276+
break;
270277
case NEED_UNWRAP:
271278
if( socketChannel.read( peerNetData ) < 0 ) {
272279
if( engine.isInboundDone() && engine.isOutboundDone() ) {
@@ -316,6 +323,8 @@ private boolean doHandshake() throws IOException {
316323
}
317324
break;
318325
case NEED_WRAP:
326+
327+
319328
myNetData.clear();
320329
try {
321330
result = engine.wrap( myAppData, myNetData );
@@ -363,8 +372,7 @@ private boolean doHandshake() throws IOException {
363372
}
364373
handshakeStatus = engine.getHandshakeStatus();
365374
break;
366-
case FINISHED:
367-
break;
375+
368376
case NOT_HANDSHAKING:
369377
break;
370378
default:

‎src/main/java/org/java_websocket/WebSocketImpl.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,7 @@
2828
import org.java_websocket.drafts.Draft;
2929
import org.java_websocket.drafts.Draft_6455;
3030
import org.java_websocket.enums.*;
31-
import org.java_websocket.exceptions.IncompleteHandshakeException;
32-
import org.java_websocket.exceptions.InvalidDataException;
33-
import org.java_websocket.exceptions.InvalidHandshakeException;
34-
import org.java_websocket.exceptions.WebsocketNotConnectedException;
31+
import org.java_websocket.exceptions.*;
3532
import org.java_websocket.framing.CloseFrame;
3633
import org.java_websocket.framing.Framedata;
3734
import org.java_websocket.framing.PingFrame;
@@ -369,6 +366,12 @@ private void decodeFrames( ByteBuffer socketBuffer ) {
369366
log.trace( "matched frame: {}" , f );
370367
draft.processFrame( this, f );
371368
}
369+
} catch ( LimitExedeedException e ) {
370+
if (e.getLimit() == Integer.MAX_VALUE) {
371+
log.error("Closing due to invalid size of frame", e);
372+
wsl.onWebsocketError(this, e);
373+
}
374+
close(e);
372375
} catch ( InvalidDataException e ) {
373376
log.error("Closing due to invalid data in frame", e);
374377
wsl.onWebsocketError( this, e );

‎src/main/java/org/java_websocket/drafts/Draft_6455.java‎

Lines changed: 124 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ public class Draft_6455 extends Draft {
8686
/**
8787
* Attribute for the payload of the current continuous frame
8888
*/
89-
private List<ByteBuffer> byteBufferList;
89+
private final List<ByteBuffer> byteBufferList;
9090

9191
/**
9292
* Attribute for the current incomplete frame
@@ -98,6 +98,13 @@ public class Draft_6455 extends Draft {
9898
*/
9999
private final Random reuseableRandom = new Random();
100100

101+
/**
102+
* Attribute for the maximum allowed size of a frame
103+
*
104+
* @since 1.4.0
105+
*/
106+
private int maxFrameSize;
107+
101108
/**
102109
* Constructor for the websocket protocol specified by RFC 6455 with default extensions
103110
* @since 1.3.5
@@ -135,7 +142,32 @@ public Draft_6455( List<IExtension> inputExtensions ) {
135142
* @since 1.3.7
136143
*/
137144
public Draft_6455( List<IExtension> inputExtensions , List<IProtocol> inputProtocols ) {
138-
if (inputExtensions == null || inputProtocols == null) {
145+
this(inputExtensions, inputProtocols, Integer.MAX_VALUE);
146+
}
147+
148+
/**
149+
* Constructor for the websocket protocol specified by RFC 6455 with custom extensions and protocols
150+
*
151+
* @param inputExtensions the extensions which should be used for this draft
152+
* @param inputMaxFrameSize the maximum allowed size of a frame (the real payload size, decoded frames can be bigger)
153+
*
154+
* @since 1.4.0
155+
*/
156+
public Draft_6455( List<IExtension> inputExtensions , int inputMaxFrameSize) {
157+
this(inputExtensions, Collections.<IProtocol>singletonList( new Protocol( "" )), inputMaxFrameSize);
158+
}
159+
160+
/**
161+
* Constructor for the websocket protocol specified by RFC 6455 with custom extensions and protocols
162+
*
163+
* @param inputExtensions the extensions which should be used for this draft
164+
* @param inputProtocols the protocols which should be used for this draft
165+
* @param inputMaxFrameSize the maximum allowed size of a frame (the real payload size, decoded frames can be bigger)
166+
*
167+
* @since 1.4.0
168+
*/
169+
public Draft_6455( List<IExtension> inputExtensions , List<IProtocol> inputProtocols, int inputMaxFrameSize ) {
170+
if (inputExtensions == null || inputProtocols == null || inputMaxFrameSize < 1) {
139171
throw new IllegalArgumentException();
140172
}
141173
knownExtensions = new ArrayList<IExtension>( inputExtensions.size());
@@ -153,6 +185,7 @@ public Draft_6455( List<IExtension> inputExtensions , List<IProtocol> inputProto
153185
knownExtensions.add( this.knownExtensions.size(), extension );
154186
}
155187
knownProtocols.addAll( inputProtocols );
188+
maxFrameSize = inputMaxFrameSize;
156189
}
157190

158191
@Override
@@ -263,6 +296,17 @@ public IProtocol getProtocol() {
263296
return protocol;
264297
}
265298

299+
300+
/**
301+
* Getter for the maximum allowed payload size which is used by this draft
302+
*
303+
* @return the size, which is allowed for the payload
304+
* @since 1.4.0
305+
*/
306+
public int getMaxFrameSize() {
307+
return maxFrameSize;
308+
}
309+
266310
/**
267311
* Getter for all available protocols for this draft
268312
* @return the protocols which are enabled for this draft
@@ -337,7 +381,7 @@ public Draft copyInstance() {
337381
for( IProtocol protocol : getKnownProtocols() ) {
338382
newProtocols.add( protocol.copyInstance() );
339383
}
340-
return new Draft_6455( newExtensions, newProtocols );
384+
return new Draft_6455( newExtensions, newProtocols, maxFrameSize );
341385
}
342386

343387
@Override
@@ -440,14 +484,18 @@ public Framedata translateSingleFrame( ByteBuffer buffer ) throws IncompleteExce
440484
bytes[i] = buffer.get( /*1 + i*/ );
441485
}
442486
long length = new BigInteger( bytes ).longValue();
443-
if( length > Integer.MAX_VALUE ) {
487+
if( length > Integer.MAX_VALUE) {
444488
log.trace( "Limit exedeed: Payloadsize is to big..." );
445489
throw new LimitExedeedException( "Payloadsize is to big..." );
446490
} else {
447491
payloadlength = ( int ) length;
448492
}
449493
}
450494
}
495+
if( payloadlength > maxFrameSize) {
496+
log.trace( "Payload limit reached. Allowed: {0} Current: {1}" , maxFrameSize, payloadlength);
497+
throw new LimitExedeedException( "Payload limit reached.", maxFrameSize );
498+
}
451499

452500
// int maskskeystart = foff + realpacketsize;
453501
realpacketsize += ( MASK ? 4 : 0 );
@@ -682,13 +730,15 @@ public void processFrame( WebSocketImpl webSocketImpl, Framedata frame ) throws
682730
throw new InvalidDataException( CloseFrame.PROTOCOL_ERROR, "Previous continuous frame sequence not completed." );
683731
}
684732
current_continuous_frame = frame;
685-
byteBufferList.add( frame.getPayloadData() );
733+
addToBufferList(frame.getPayloadData());
734+
checkBufferLimit();
686735
} else if( frame.isFin() ) {
687736
if( current_continuous_frame == null ) {
688737
log.trace( "Protocol error: Previous continuous frame sequence not completed." );
689738
throw new InvalidDataException( CloseFrame.PROTOCOL_ERROR, "Continuous frame sequence was not started." );
690739
}
691-
byteBufferList.add( frame.getPayloadData() );
740+
addToBufferList(frame.getPayloadData());
741+
checkBufferLimit();
692742
if( current_continuous_frame.getOpcode() == Opcode.TEXT ) {
693743
((FramedataImpl1) current_continuous_frame).setPayload( getPayloadFromByteBufferList() );
694744
((FramedataImpl1) current_continuous_frame ).isValid();
@@ -709,7 +759,7 @@ public void processFrame( WebSocketImpl webSocketImpl, Framedata frame ) throws
709759
}
710760
}
711761
current_continuous_frame = null;
712-
byteBufferList.clear();
762+
clearBufferList();
713763
} else if( current_continuous_frame == null ) {
714764
log.error( "Protocol error: Continuous frame sequence was not started." );
715765
throw new InvalidDataException( CloseFrame.PROTOCOL_ERROR, "Continuous frame sequence was not started." );
@@ -723,7 +773,7 @@ public void processFrame( WebSocketImpl webSocketImpl, Framedata frame ) throws
723773
}
724774
//Checking if the current continuous frame contains a correct payload with the other frames combined
725775
if( curop == Opcode.CONTINUOUS && current_continuous_frame != null ) {
726-
byteBufferList.add( frame.getPayloadData() );
776+
addToBufferList(frame.getPayloadData());
727777
}
728778
} else if( current_continuous_frame != null ) {
729779
log.error( "Protocol error: Continuous frame sequence not completed." );
@@ -748,6 +798,38 @@ public void processFrame( WebSocketImpl webSocketImpl, Framedata frame ) throws
748798
}
749799
}
750800

801+
/**
802+
* Clear the current bytebuffer list
803+
*/
804+
private void clearBufferList() {
805+
synchronized (byteBufferList) {
806+
byteBufferList.clear();
807+
}
808+
}
809+
810+
/**
811+
* Add a payload to the current bytebuffer list
812+
* @param payloadData the new payload
813+
*/
814+
private void addToBufferList(ByteBuffer payloadData) {
815+
synchronized (byteBufferList) {
816+
byteBufferList.add(payloadData);
817+
}
818+
}
819+
820+
/**
821+
* Check the current size of the buffer and throw an exception if the size is bigger than the max allowed frame size
822+
* @throws LimitExedeedException if the current size is bigger than the allowed size
823+
*/
824+
private void checkBufferLimit() throws LimitExedeedException {
825+
long totalSize = getByteBufferListSize();
826+
if( totalSize > maxFrameSize ) {
827+
clearBufferList();
828+
log.trace("Payload limit reached. Allowed: {0} Current: {1}", maxFrameSize, totalSize);
829+
throw new LimitExedeedException(maxFrameSize);
830+
}
831+
}
832+
751833
@Override
752834
public CloseHandshakeType getCloseHandshakeType() {
753835
return CloseHandshakeType.TWOWAY;
@@ -760,24 +842,27 @@ public String toString() {
760842
result += " extension: " + getExtension().toString();
761843
if ( getProtocol() != null )
762844
result += " protocol: " + getProtocol().toString();
845+
result += " max frame size: " + this.maxFrameSize;
763846
return result;
764847
}
765848

766849
@Override
767-
public boolean equals( Object o ) {
768-
if( this == o ) return true;
769-
if( o == null || getClass() != o.getClass() ) return false;
850+
public boolean equals(Object o) {
851+
if (this == o) return true;
852+
if (o == null || getClass() != o.getClass()) return false;
770853

771-
Draft_6455 that = ( Draft_6455 ) o;
854+
Draft_6455 that = (Draft_6455) o;
772855

773-
if( extension != null ? !extension.equals( that.extension ) : that.extension != null ) return false;
774-
return protocol != null ? protocol.equals( that.protocol ) : that.protocol == null;
856+
if (maxFrameSize != that.getMaxFrameSize()) return false;
857+
if (extension != null ? !extension.equals(that.getExtension()) : that.getExtension() != null) return false;
858+
return protocol != null ? protocol.equals(that.getProtocol()) : that.getProtocol() == null;
775859
}
776860

777861
@Override
778862
public int hashCode() {
779863
int result = extension != null ? extension.hashCode() : 0;
780-
result = 31 * result + ( protocol != null ? protocol.hashCode() : 0 );
864+
result = 31 * result + (protocol != null ? protocol.hashCode() : 0);
865+
result = 31 * result + (int) (maxFrameSize ^ (maxFrameSize >>> 32));
781866
return result;
782867
}
783868

@@ -788,18 +873,32 @@ public int hashCode() {
788873
*/
789874
private ByteBuffer getPayloadFromByteBufferList() throws LimitExedeedException {
790875
long totalSize = 0;
791-
for (ByteBuffer buffer : byteBufferList) {
792-
totalSize +=buffer.limit();
793-
}
794-
if (totalSize > Integer.MAX_VALUE) {
795-
log.trace( "Payloadsize is to big...");
796-
throw new LimitExedeedException( "Payloadsize is to big..." );
797-
}
798-
ByteBuffer resultingByteBuffer = ByteBuffer.allocate( (int) totalSize );
799-
for (ByteBuffer buffer : byteBufferList) {
800-
resultingByteBuffer.put( buffer );
876+
ByteBuffer resultingByteBuffer;
877+
synchronized (byteBufferList) {
878+
for (ByteBuffer buffer : byteBufferList) {
879+
totalSize += buffer.limit();
880+
}
881+
checkBufferLimit();
882+
resultingByteBuffer = ByteBuffer.allocate( (int) totalSize );
883+
for (ByteBuffer buffer : byteBufferList) {
884+
resultingByteBuffer.put( buffer );
885+
}
801886
}
802887
resultingByteBuffer.flip();
803888
return resultingByteBuffer;
804889
}
890+
891+
/**
892+
* Get the current size of the resulting bytebuffer in the bytebuffer list
893+
* @return the size as long (to not get an integer overflow)
894+
*/
895+
private long getByteBufferListSize() {
896+
long totalSize = 0;
897+
synchronized (byteBufferList) {
898+
for (ByteBuffer buffer : byteBufferList) {
899+
totalSize += buffer.limit();
900+
}
901+
}
902+
return totalSize;
903+
}
805904
}

‎src/main/java/org/java_websocket/exceptions/LimitExedeedException.java‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,13 +37,38 @@ public class LimitExedeedException extends InvalidDataException {
3737
*/
3838
private static final long serialVersionUID = 6908339749836826785L;
3939

40+
/**
41+
* A closer indication about the limit
42+
*/
43+
private final int limit;
44+
4045
/**
4146
* constructor for a LimitExedeedException
4247
* <p>
4348
* calling InvalidDataException with closecode TOOBIG
4449
*/
4550
public LimitExedeedException() {
51+
this(Integer.MAX_VALUE);
52+
}
53+
54+
/**
55+
* constructor for a LimitExedeedException
56+
* <p>
57+
* calling InvalidDataException with closecode TOOBIG
58+
*/
59+
public LimitExedeedException(int limit) {
4660
super( CloseFrame.TOOBIG);
61+
this.limit = limit;
62+
}
63+
64+
/**
65+
* constructor for a LimitExedeedException
66+
* <p>
67+
* calling InvalidDataException with closecode TOOBIG
68+
*/
69+
public LimitExedeedException(String s, int limit) {
70+
super( CloseFrame.TOOBIG, s);
71+
this.limit = limit;
4772
}
4873

4974
/**
@@ -54,7 +79,14 @@ public LimitExedeedException() {
5479
* @param s the detail message.
5580
*/
5681
public LimitExedeedException(String s) {
57-
super( CloseFrame.TOOBIG, s);
82+
this(s, Integer.MAX_VALUE);
5883
}
5984

85+
/**
86+
* Get the limit which was hit so this exception was caused
87+
* @return the limit as int
88+
*/
89+
public int getLimit() {
90+
return limit;
91+
}
6092
}

0 commit comments

Comments
 (0)