conn.go raw
1 // Package ws implements the server side of RFC 6455 WebSockets.
2 //
3 // Only what a relay needs is here: the opening handshake, masked client
4 // frames, fragmentation, ping/pong, close and a maximum message size. Server
5 // frames are never masked. Reads and writes may run concurrently (one reader
6 // goroutine, one writer goroutine); Write is internally serialized.
7 package ws
8
9 import (
10 "bufio"
11 "crypto/sha1"
12 "encoding/base64"
13 "encoding/binary"
14 "errors"
15 "io"
16 "net"
17 "strings"
18 "sync"
19 "time"
20 "unicode/utf8"
21 )
22
23 // Frame opcodes.
24 const (
25 OpContinuation = 0x0
26 OpText = 0x1
27 OpBinary = 0x2
28 OpClose = 0x8
29 OpPing = 0x9
30 OpPong = 0xa
31 )
32
33 const acceptGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
34
35 var (
36 // ErrProtocol is returned for a framing violation.
37 ErrProtocol = errors.New("ws: protocol error")
38 // ErrTooLarge is returned when a message exceeds MaxMessage.
39 ErrTooLarge = errors.New("ws: message too large")
40 // ErrClosed is returned after the peer closed the connection.
41 ErrClosed = errors.New("ws: closed")
42 )
43
44 // Conn is a server-side WebSocket connection.
45 type Conn struct {
46 conn net.Conn
47 br *bufio.Reader
48 bw *bufio.Writer
49
50 wmu sync.Mutex
51 dead bool
52
53 // MaxMessage bounds one (possibly fragmented) message. Zero means 1 MiB.
54 MaxMessage int64
55 }
56
57 // Upgrade validates the handshake headers and writes the 101 response. The
58 // caller must have checked that the request is a GET.
59 func Upgrade(conn net.Conn, brw *bufio.ReadWriter, header map[string]string) (*Conn, error) {
60 if !headerContains(header["upgrade"], "websocket") {
61 return nil, errors.New("ws: missing Upgrade: websocket")
62 }
63 if !headerContains(header["connection"], "upgrade") {
64 return nil, errors.New("ws: missing Connection: Upgrade")
65 }
66 if header["sec-websocket-version"] != "13" {
67 return nil, errors.New("ws: unsupported websocket version")
68 }
69 key := header["sec-websocket-key"]
70 if key == "" {
71 return nil, errors.New("ws: missing Sec-WebSocket-Key")
72 }
73 sum := sha1.Sum([]byte(key + acceptGUID))
74 accept := base64.StdEncoding.EncodeToString(sum[:])
75
76 if _, err := brw.WriteString("HTTP/1.1 101 Switching Protocols\r\n" +
77 "Upgrade: websocket\r\n" +
78 "Connection: Upgrade\r\n" +
79 "Sec-WebSocket-Accept: " + accept + "\r\n\r\n"); err != nil {
80 return nil, err
81 }
82 if err := brw.Flush(); err != nil {
83 return nil, err
84 }
85 return &Conn{conn: conn, br: brw.Reader, bw: brw.Writer, MaxMessage: 1 << 20}, nil
86 }
87
88 func headerContains(value, token string) bool {
89 for _, part := range strings.Split(value, ",") {
90 if strings.EqualFold(strings.TrimSpace(part), token) {
91 return true
92 }
93 }
94 return false
95 }
96
97 // SetReadDeadline sets the deadline for the next read.
98 func (c *Conn) SetReadDeadline(t time.Time) error { return c.conn.SetReadDeadline(t) }
99
100 // SetWriteDeadline sets the deadline for the next write.
101 func (c *Conn) SetWriteDeadline(t time.Time) error { return c.conn.SetWriteDeadline(t) }
102
103 // RemoteAddr returns the peer address.
104 func (c *Conn) RemoteAddr() string { return c.conn.RemoteAddr().String() }
105
106 // ReadMessage returns the next complete data message. Text messages are
107 // validated as UTF-8. Control frames are handled internally.
108 func (c *Conn) ReadMessage() (opcode byte, payload []byte, err error) {
109 if c.MaxMessage <= 0 {
110 c.MaxMessage = 1 << 20
111 }
112 var fragOp byte
113 var message []byte
114 inFrag := false
115
116 for {
117 hdr := make([]byte, 2)
118 if _, err := io.ReadFull(c.br, hdr); err != nil {
119 return 0, nil, err
120 }
121 fin := hdr[0]&0x80 != 0
122 rsv := hdr[0] & 0x70
123 op := hdr[0] & 0x0f
124 masked := hdr[1]&0x80 != 0
125 length := int64(hdr[1] & 0x7f)
126 if rsv != 0 {
127 c.fail(1002, "reserved bits set")
128 return 0, nil, ErrProtocol
129 }
130 if !masked {
131 c.fail(1002, "client frames must be masked")
132 return 0, nil, ErrProtocol
133 }
134 switch length {
135 case 126:
136 ext := make([]byte, 2)
137 if _, err := io.ReadFull(c.br, ext); err != nil {
138 return 0, nil, err
139 }
140 length = int64(binary.BigEndian.Uint16(ext))
141 case 127:
142 ext := make([]byte, 8)
143 if _, err := io.ReadFull(c.br, ext); err != nil {
144 return 0, nil, err
145 }
146 v := binary.BigEndian.Uint64(ext)
147 if v > 1<<62 {
148 c.fail(1009, "frame too large")
149 return 0, nil, ErrTooLarge
150 }
151 length = int64(v)
152 }
153
154 control := op&0x08 != 0
155 if control {
156 if !fin {
157 c.fail(1002, "fragmented control frame")
158 return 0, nil, ErrProtocol
159 }
160 if length > 125 {
161 c.fail(1002, "oversized control frame")
162 return 0, nil, ErrProtocol
163 }
164 }
165 if length > c.MaxMessage || int64(len(message))+length > c.MaxMessage {
166 c.fail(1009, "message too large")
167 return 0, nil, ErrTooLarge
168 }
169
170 mask := make([]byte, 4)
171 if _, err := io.ReadFull(c.br, mask); err != nil {
172 return 0, nil, err
173 }
174 data := make([]byte, length)
175 if _, err := io.ReadFull(c.br, data); err != nil {
176 return 0, nil, err
177 }
178 for i := range data {
179 data[i] ^= mask[i&3]
180 }
181
182 switch op {
183 case OpPing:
184 if err := c.WriteFrame(OpPong, data); err != nil {
185 return 0, nil, err
186 }
187 continue
188 case OpPong:
189 continue
190 case OpClose:
191 code := uint16(1000)
192 if len(data) >= 2 {
193 code = binary.BigEndian.Uint16(data[:2])
194 }
195 _ = c.Close(code, "")
196 return 0, nil, ErrClosed
197 case OpText, OpBinary:
198 if inFrag {
199 c.fail(1002, "new data frame inside a fragmented message")
200 return 0, nil, ErrProtocol
201 }
202 if fin {
203 if op == OpText && !utf8.Valid(data) {
204 c.fail(1007, "invalid UTF-8")
205 return 0, nil, ErrProtocol
206 }
207 return op, data, nil
208 }
209 fragOp = op
210 message = append(message[:0], data...)
211 inFrag = true
212 case OpContinuation:
213 if !inFrag {
214 c.fail(1002, "continuation without a message")
215 return 0, nil, ErrProtocol
216 }
217 message = append(message, data...)
218 if fin {
219 if fragOp == OpText && !utf8.Valid(message) {
220 c.fail(1007, "invalid UTF-8")
221 return 0, nil, ErrProtocol
222 }
223 return fragOp, message, nil
224 }
225 default:
226 c.fail(1002, "unknown opcode")
227 return 0, nil, ErrProtocol
228 }
229 }
230 }
231
232 // WriteText sends a text message.
233 func (c *Conn) WriteText(data []byte) error { return c.WriteFrame(OpText, data) }
234
235 // WriteFrame sends one unfragmented frame.
236 func (c *Conn) WriteFrame(op byte, payload []byte) error {
237 c.wmu.Lock()
238 defer c.wmu.Unlock()
239 if c.dead {
240 return ErrClosed
241 }
242 var hdr [10]byte
243 hdr[0] = 0x80 | op
244 n := 2
245 switch {
246 case len(payload) < 126:
247 hdr[1] = byte(len(payload))
248 case len(payload) <= 0xffff:
249 hdr[1] = 126
250 binary.BigEndian.PutUint16(hdr[2:4], uint16(len(payload)))
251 n = 4
252 default:
253 hdr[1] = 127
254 binary.BigEndian.PutUint64(hdr[2:10], uint64(len(payload)))
255 n = 10
256 }
257 if _, err := c.bw.Write(hdr[:n]); err != nil {
258 return err
259 }
260 if len(payload) > 0 {
261 if _, err := c.bw.Write(payload); err != nil {
262 return err
263 }
264 }
265 return c.bw.Flush()
266 }
267
268 // Close sends a close frame. It is not an error to call it more than once.
269 func (c *Conn) Close(code uint16, reason string) error {
270 c.wmu.Lock()
271 defer c.wmu.Unlock()
272 if c.dead {
273 return nil
274 }
275 c.dead = true
276 payload := make([]byte, 2+len(reason))
277 binary.BigEndian.PutUint16(payload[:2], code)
278 copy(payload[2:], reason)
279 var hdr [4]byte
280 hdr[0] = 0x80 | OpClose
281 hdr[1] = byte(len(payload))
282 binary.BigEndian.PutUint16(hdr[2:4], 0)
283 if _, err := c.bw.Write(hdr[:2]); err != nil {
284 return err
285 }
286 if _, err := c.bw.Write(payload); err != nil {
287 return err
288 }
289 return c.bw.Flush()
290 }
291
292 func (c *Conn) fail(code uint16, reason string) {
293 _ = c.Close(code, reason)
294 }
295
296 // CloseSocket closes the underlying TCP connection.
297 func (c *Conn) CloseSocket() error { return c.conn.Close() }
298
299 // HandshakeResponse validates the opening handshake headers and returns the
300 // complete 101 response. Reader workers use it so that only their paired
301 // writer domain writes to the socket.
302 func HandshakeResponse(header map[string]string) ([]byte, error) {
303 if !headerContains(header["upgrade"], "websocket") {
304 return nil, errors.New("ws: missing Upgrade: websocket")
305 }
306 if !headerContains(header["connection"], "upgrade") {
307 return nil, errors.New("ws: missing Connection: Upgrade")
308 }
309 if header["sec-websocket-version"] != "13" {
310 return nil, errors.New("ws: unsupported websocket version")
311 }
312 key := header["sec-websocket-key"]
313 if key == "" {
314 return nil, errors.New("ws: missing Sec-WebSocket-Key")
315 }
316 sum := sha1.Sum([]byte(key + acceptGUID))
317 accept := base64.StdEncoding.EncodeToString(sum[:])
318 out := []byte("HTTP/1.1 101 Switching Protocols\r\n" +
319 "Upgrade: websocket\r\n" +
320 "Connection: Upgrade\r\n" +
321 "Sec-WebSocket-Accept: " + accept + "\r\n\r\n")
322 return out, nil
323 }
324