reader.go raw
1 package ws
2
3 import (
4 "encoding/binary"
5 "io"
6 "unicode/utf8"
7 )
8
9 // Reader parses client frames without writing anything. Control frames are
10 // returned to the caller so that the socket's writer domain remains the only
11 // writer on the connection.
12 type Reader struct {
13 br *bufioReader
14 MaxMessage int64
15 fragOp byte
16 frag []byte
17 inFrag bool
18 }
19
20 // bufioReader is the small interface Reader needs; it avoids importing bufio
21 // in this file's public signature while still accepting *bufio.Reader.
22 type bufioReader struct {
23 ReadFull func([]byte) (int, error)
24 }
25
26 // NewReader wraps a buffered reader.
27 func NewReader(br interface{ Read([]byte) (int, error) }, max int64) *Reader {
28 if max <= 0 {
29 max = 1 << 20
30 }
31 return &Reader{
32 br: &bufioReader{ReadFull: func(p []byte) (int, error) { return io.ReadFull(br, p) }},
33 MaxMessage: max,
34 }
35 }
36
37 // Read returns the next frame. Data messages are reassembled across
38 // continuations; ping, pong and close frames are returned immediately.
39 func (r *Reader) Read() (op byte, payload []byte, err error) {
40 for {
41 hdr := make([]byte, 2)
42 if _, err := r.br.ReadFull(hdr); err != nil {
43 return 0, nil, err
44 }
45 fin := hdr[0]&0x80 != 0
46 op := hdr[0] & 0x0f
47 if hdr[0]&0x70 != 0 {
48 return 0, nil, ErrProtocol
49 }
50 if hdr[1]&0x80 == 0 {
51 return 0, nil, ErrProtocol
52 }
53 length := int64(hdr[1] & 0x7f)
54 switch length {
55 case 126:
56 ext := make([]byte, 2)
57 if _, err := r.br.ReadFull(ext); err != nil {
58 return 0, nil, err
59 }
60 length = int64(binary.BigEndian.Uint16(ext))
61 case 127:
62 ext := make([]byte, 8)
63 if _, err := r.br.ReadFull(ext); err != nil {
64 return 0, nil, err
65 }
66 v := binary.BigEndian.Uint64(ext)
67 if v > 1<<62 {
68 return 0, nil, ErrTooLarge
69 }
70 length = int64(v)
71 }
72 control := op&0x08 != 0
73 if control && (!fin || length > 125) {
74 return 0, nil, ErrProtocol
75 }
76 if length > r.MaxMessage || int64(len(r.frag))+length > r.MaxMessage {
77 return 0, nil, ErrTooLarge
78 }
79 mask := make([]byte, 4)
80 if _, err := r.br.ReadFull(mask); err != nil {
81 return 0, nil, err
82 }
83 data := make([]byte, length)
84 if _, err := r.br.ReadFull(data); err != nil {
85 return 0, nil, err
86 }
87 for i := range data {
88 data[i] ^= mask[i&3]
89 }
90
91 switch op {
92 case OpPing, OpPong, OpClose:
93 return op, data, nil
94 case OpText, OpBinary:
95 if r.inFrag {
96 return 0, nil, ErrProtocol
97 }
98 if fin {
99 if op == OpText && !utf8.Valid(data) {
100 return 0, nil, ErrProtocol
101 }
102 return op, data, nil
103 }
104 r.fragOp = op
105 r.frag = append(r.frag[:0], data...)
106 r.inFrag = true
107 case OpContinuation:
108 if !r.inFrag {
109 return 0, nil, ErrProtocol
110 }
111 r.frag = append(r.frag, data...)
112 if fin {
113 if r.fragOp == OpText && !utf8.Valid(r.frag) {
114 return 0, nil, ErrProtocol
115 }
116 out := append([]byte(nil), r.frag...)
117 r.inFrag = false
118 r.frag = r.frag[:0]
119 return r.fragOp, out, nil
120 }
121 default:
122 return 0, nil, ErrProtocol
123 }
124 }
125 }
126
127 // AppendFrame appends one unmasked server frame to dst.
128 func AppendFrame(dst []byte, op byte, payload []byte) []byte {
129 dst = append(dst, 0x80|op)
130 switch {
131 case len(payload) < 126:
132 dst = append(dst, byte(len(payload)))
133 case len(payload) <= 0xffff:
134 dst = append(dst, 126, byte(len(payload)>>8), byte(len(payload)))
135 default:
136 dst = append(dst, 127)
137 var ext [8]byte
138 binary.BigEndian.PutUint64(ext[:], uint64(len(payload)))
139 dst = append(dst, ext[:]...)
140 }
141 return append(dst, payload...)
142 }
143
144 // AppendClose appends a close frame with a status code and reason.
145 func AppendClose(dst []byte, code uint16, reason string) []byte {
146 payload := make([]byte, 2+len(reason))
147 binary.BigEndian.PutUint16(payload[:2], code)
148 copy(payload[2:], reason)
149 return AppendFrame(dst, OpClose, payload)
150 }
151