@@ -36,9 +36,6 @@ type Conn struct {
36
36
closer io.Closer
37
37
client bool
38
38
39
- // read limit for a message in bytes.
40
- msgReadLimit int64
41
-
42
39
closeOnce sync.Once
43
40
closeErr error
44
41
closed chan struct {}
@@ -50,10 +47,13 @@ type Conn struct {
50
47
writeFrameLock chan struct {}
51
48
writeHeaderBuf []byte
52
49
writeHeader * header
50
+ // read limit for a message in bytes.
51
+ msgReadLimit int64
53
52
54
53
// messageWriter state.
55
54
writeMsgOpcode opcode
56
55
writeMsgCtx context.Context
56
+ readMsgLeft int64
57
57
58
58
// Used to ensure the previous reader is read till EOF before allowing
59
59
// a new one.
@@ -371,20 +371,19 @@ func (c *Conn) reader(ctx context.Context) (MessageType, io.Reader, error) {
371
371
c .readMsgHeader = h
372
372
c .readFrameEOF = false
373
373
c .readMaskPos = 0
374
+ c .readMsgLeft = c .msgReadLimit
374
375
375
376
r := & messageReader {
376
- c : c ,
377
- left : c .msgReadLimit ,
377
+ c : c ,
378
378
}
379
379
c .previousReader = r
380
380
return MessageType (h .opcode ), r , nil
381
381
}
382
382
383
383
// messageReader enables reading a data frame from the WebSocket connection.
384
384
type messageReader struct {
385
- c * Conn
386
- left int64
387
- eof bool
385
+ c * Conn
386
+ eof bool
388
387
}
389
388
390
389
// Read reads as many bytes as possible into p.
@@ -406,14 +405,14 @@ func (r *messageReader) read(p []byte) (int, error) {
406
405
return 0 , xerrors .Errorf ("cannot use EOFed reader" )
407
406
}
408
407
409
- if r .left <= 0 {
408
+ if r .c . readMsgLeft <= 0 {
410
409
err := xerrors .Errorf ("read limited at %v bytes" , r .c .msgReadLimit )
411
410
r .c .Close (StatusMessageTooBig , err .Error ())
412
411
return 0 , err
413
412
}
414
413
415
- if int64 (len (p )) > r .left {
416
- p = p [:r .left ]
414
+ if int64 (len (p )) > r .c . readMsgLeft {
415
+ p = p [:r .c . readMsgLeft ]
417
416
}
418
417
419
418
if r .c .readFrameEOF {
@@ -441,7 +440,7 @@ func (r *messageReader) read(p []byte) (int, error) {
441
440
n , err := r .c .readFramePayload (r .c .readMsgCtx , p )
442
441
443
442
h .payloadLength -= int64 (n )
444
- r .left -= int64 (n )
443
+ r .c . readMsgLeft -= int64 (n )
445
444
if h .masked {
446
445
r .c .readMaskPos = fastXOR (h .maskKey , r .c .readMaskPos , p )
447
446
}
0 commit comments