bughz/glider
1
0
Fork 0
mirror of https://github.com/nadoo/glider.git synced 2025-02-24 01:45:39 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

ws: fixed a bug in frame reader

Browse Source
This commit is contained in:
nadoo 2018-07-23 01:41:18 +08:00
parent cdabe41e34
commit 68d8fb8b60
1 changed files with 52 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
proxy/ws/frame.go
View File

@ -30,13 +30,16 @@ import (
) )
const ( const (
finalBit byte = 1 << 7
defaultFrameSize = 4096 defaultFrameSize = 4096
maxFrameHeaderSize = 2 + 8 + 4 // Fixed header + length + mask maxFrameHeaderSize = 2 + 8 + 4 // Fixed header + length + mask
maskKeyLen = 4
)
const (
finalBit byte = 1 << 7
maskBit byte = 1 << 7 maskBit byte = 1 << 7
opCodeBinary byte = 2 opCodeBinary byte = 2
opClose byte = 8 opClose byte = 8
maskKeyLen = 4
) )
type frameWriter struct { type frameWriter struct {
@ -85,18 +88,21 @@ func (w *frameWriter) ReadFrom(r io.Reader) (n int64, err error) {
binary.BigEndian.PutUint64(buf[2:2+lengthFieldLen], uint64(nr)) binary.BigEndian.PutUint64(buf[2:2+lengthFieldLen], uint64(nr))
} }
// header and length
_, ew := w.Writer.Write(buf[:2+lengthFieldLen]) _, ew := w.Writer.Write(buf[:2+lengthFieldLen])
if ew != nil { if ew != nil {
err = ew err = ew
break break
} }
// maskkey
_, ew = w.Writer.Write(w.maskKey) _, ew = w.Writer.Write(w.maskKey)
if ew != nil { if ew != nil {
err = ew err = ew
break break
} }
// payload
payloadBuf = payloadBuf[:nr] payloadBuf = payloadBuf[:nr]
for i := range payloadBuf { for i := range payloadBuf {
payloadBuf[i] = payloadBuf[i] ^ w.maskKey[i%4] payloadBuf[i] = payloadBuf[i] ^ w.maskKey[i%4]
@ -124,7 +130,7 @@ func (w *frameWriter) ReadFrom(r io.Reader) (n int64, err error) {
type frameReader struct { type frameReader struct {
io.Reader io.Reader
buf []byte buf []byte
leftover []byte leftBytes int64
} }
// FrameReader returns a chunked reader // FrameReader returns a chunked reader
@ -136,12 +142,7 @@ func FrameReader(r io.Reader) io.Reader {
} }
func (r *frameReader) Read(b []byte) (int, error) { func (r *frameReader) Read(b []byte) (int, error) {
if len(r.leftover) > 0 { if r.leftBytes == 0 {
n := copy(b, r.leftover)
r.leftover = r.leftover[n:]
return n, nil
}
// get msg header // get msg header
_, err := io.ReadFull(r.Reader, r.buf[:2]) _, err := io.ReadFull(r.Reader, r.buf[:2])
if err != nil { if err != nil {
@ -151,26 +152,39 @@ func (r *frameReader) Read(b []byte) (int, error) {
// final := r.buf[0]&finalBit != 0 // final := r.buf[0]&finalBit != 0
// frameType := int(r.buf[0] & 0xf) // frameType := int(r.buf[0] & 0xf)
// mask := r.buf[1]&maskBit != 0 // mask := r.buf[1]&maskBit != 0
len := int64(r.buf[1] & 0x7f) r.leftBytes = int64(r.buf[1] & 0x7f)
switch len { switch r.leftBytes {
case 126: case 126:
io.ReadFull(r.Reader, r.buf[:2]) io.ReadFull(r.Reader, r.buf[:2])
len = int64(binary.BigEndian.Uint16(r.buf[0:])) r.leftBytes = int64(binary.BigEndian.Uint16(r.buf[0:]))
case 127: case 127:
io.ReadFull(r.Reader, r.buf[:8]) io.ReadFull(r.Reader, r.buf[:8])
len = int64(binary.BigEndian.Uint64(r.buf[0:])) r.leftBytes = int64(binary.BigEndian.Uint64(r.buf[0:]))
}
}
var n, m int
var err error
if r.leftBytes > int64(len(r.buf)) {
if len(b) < len(r.buf) {
n, err = r.Reader.Read(r.buf[:len(b)])
} else {
n, err = r.Reader.Read(r.buf)
}
} else {
if int64(len(b)) < r.leftBytes {
n, err = io.ReadFull(r.Reader, r.buf[:len(b)])
} else {
n, err = io.ReadFull(r.Reader, r.buf[:r.leftBytes])
}
} }
// get payload
_, err = io.ReadFull(r.Reader, r.buf[:len])
if err != nil { if err != nil {
return 0, err return 0, err
} }
m := copy(b, r.buf[:len]) m = copy(b, r.buf[:n])
if m < int(len) { r.leftBytes -= int64(m)
r.leftover = r.buf[m:len]
}
return m, err return m, err
} }