1// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package websocket
6
7import (
8	"bufio"
9	"encoding/binary"
10	"errors"
11	"io"
12	"io/ioutil"
13	"math/rand"
14	"net"
15	"strconv"
16	"sync"
17	"time"
18	"unicode/utf8"
19)
20
21const (
22	// Frame header byte 0 bits from Section 5.2 of RFC 6455
23	finalBit = 1 << 7
24	rsv1Bit  = 1 << 6
25	rsv2Bit  = 1 << 5
26	rsv3Bit  = 1 << 4
27
28	// Frame header byte 1 bits from Section 5.2 of RFC 6455
29	maskBit = 1 << 7
30
31	maxFrameHeaderSize         = 2 + 8 + 4 // Fixed header + length + mask
32	maxControlFramePayloadSize = 125
33
34	writeWait = time.Second
35
36	defaultReadBufferSize  = 4096
37	defaultWriteBufferSize = 4096
38
39	continuationFrame = 0
40	noFrame           = -1
41)
42
43// Close codes defined in RFC 6455, section 11.7.
44const (
45	CloseNormalClosure           = 1000
46	CloseGoingAway               = 1001
47	CloseProtocolError           = 1002
48	CloseUnsupportedData         = 1003
49	CloseNoStatusReceived        = 1005
50	CloseAbnormalClosure         = 1006
51	CloseInvalidFramePayloadData = 1007
52	ClosePolicyViolation         = 1008
53	CloseMessageTooBig           = 1009
54	CloseMandatoryExtension      = 1010
55	CloseInternalServerErr       = 1011
56	CloseServiceRestart          = 1012
57	CloseTryAgainLater           = 1013
58	CloseTLSHandshake            = 1015
59)
60
61// The message types are defined in RFC 6455, section 11.8.
62const (
63	// TextMessage denotes a text data message. The text message payload is
64	// interpreted as UTF-8 encoded text data.
65	TextMessage = 1
66
67	// BinaryMessage denotes a binary data message.
68	BinaryMessage = 2
69
70	// CloseMessage denotes a close control message. The optional message
71	// payload contains a numeric code and text. Use the FormatCloseMessage
72	// function to format a close message payload.
73	CloseMessage = 8
74
75	// PingMessage denotes a ping control message. The optional message payload
76	// is UTF-8 encoded text.
77	PingMessage = 9
78
79	// PongMessage denotes a pong control message. The optional message payload
80	// is UTF-8 encoded text.
81	PongMessage = 10
82)
83
84// ErrCloseSent is returned when the application writes a message to the
85// connection after sending a close message.
86var ErrCloseSent = errors.New("websocket: close sent")
87
88// ErrReadLimit is returned when reading a message that is larger than the
89// read limit set for the connection.
90var ErrReadLimit = errors.New("websocket: read limit exceeded")
91
92// netError satisfies the net Error interface.
93type netError struct {
94	msg       string
95	temporary bool
96	timeout   bool
97}
98
99func (e *netError) Error() string   { return e.msg }
100func (e *netError) Temporary() bool { return e.temporary }
101func (e *netError) Timeout() bool   { return e.timeout }
102
103// CloseError represents a close message.
104type CloseError struct {
105	// Code is defined in RFC 6455, section 11.7.
106	Code int
107
108	// Text is the optional text payload.
109	Text string
110}
111
112func (e *CloseError) Error() string {
113	s := []byte("websocket: close ")
114	s = strconv.AppendInt(s, int64(e.Code), 10)
115	switch e.Code {
116	case CloseNormalClosure:
117		s = append(s, " (normal)"...)
118	case CloseGoingAway:
119		s = append(s, " (going away)"...)
120	case CloseProtocolError:
121		s = append(s, " (protocol error)"...)
122	case CloseUnsupportedData:
123		s = append(s, " (unsupported data)"...)
124	case CloseNoStatusReceived:
125		s = append(s, " (no status)"...)
126	case CloseAbnormalClosure:
127		s = append(s, " (abnormal closure)"...)
128	case CloseInvalidFramePayloadData:
129		s = append(s, " (invalid payload data)"...)
130	case ClosePolicyViolation:
131		s = append(s, " (policy violation)"...)
132	case CloseMessageTooBig:
133		s = append(s, " (message too big)"...)
134	case CloseMandatoryExtension:
135		s = append(s, " (mandatory extension missing)"...)
136	case CloseInternalServerErr:
137		s = append(s, " (internal server error)"...)
138	case CloseTLSHandshake:
139		s = append(s, " (TLS handshake error)"...)
140	}
141	if e.Text != "" {
142		s = append(s, ": "...)
143		s = append(s, e.Text...)
144	}
145	return string(s)
146}
147
148// IsCloseError returns boolean indicating whether the error is a *CloseError
149// with one of the specified codes.
150func IsCloseError(err error, codes ...int) bool {
151	if e, ok := err.(*CloseError); ok {
152		for _, code := range codes {
153			if e.Code == code {
154				return true
155			}
156		}
157	}
158	return false
159}
160
161// IsUnexpectedCloseError returns boolean indicating whether the error is a
162// *CloseError with a code not in the list of expected codes.
163func IsUnexpectedCloseError(err error, expectedCodes ...int) bool {
164	if e, ok := err.(*CloseError); ok {
165		for _, code := range expectedCodes {
166			if e.Code == code {
167				return false
168			}
169		}
170		return true
171	}
172	return false
173}
174
175var (
176	errWriteTimeout        = &netError{msg: "websocket: write timeout", timeout: true, temporary: true}
177	errUnexpectedEOF       = &CloseError{Code: CloseAbnormalClosure, Text: io.ErrUnexpectedEOF.Error()}
178	errBadWriteOpCode      = errors.New("websocket: bad write message type")
179	errWriteClosed         = errors.New("websocket: write closed")
180	errInvalidControlFrame = errors.New("websocket: invalid control frame")
181)
182
183func newMaskKey() [4]byte {
184	n := rand.Uint32()
185	return [4]byte{byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24)}
186}
187
188func hideTempErr(err error) error {
189	if e, ok := err.(net.Error); ok && e.Temporary() {
190		err = &netError{msg: e.Error(), timeout: e.Timeout()}
191	}
192	return err
193}
194
195func isControl(frameType int) bool {
196	return frameType == CloseMessage || frameType == PingMessage || frameType == PongMessage
197}
198
199func isData(frameType int) bool {
200	return frameType == TextMessage || frameType == BinaryMessage
201}
202
203var validReceivedCloseCodes = map[int]bool{
204	// see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
205
206	CloseNormalClosure:           true,
207	CloseGoingAway:               true,
208	CloseProtocolError:           true,
209	CloseUnsupportedData:         true,
210	CloseNoStatusReceived:        false,
211	CloseAbnormalClosure:         false,
212	CloseInvalidFramePayloadData: true,
213	ClosePolicyViolation:         true,
214	CloseMessageTooBig:           true,
215	CloseMandatoryExtension:      true,
216	CloseInternalServerErr:       true,
217	CloseServiceRestart:          true,
218	CloseTryAgainLater:           true,
219	CloseTLSHandshake:            false,
220}
221
222func isValidReceivedCloseCode(code int) bool {
223	return validReceivedCloseCodes[code] || (code >= 3000 && code <= 4999)
224}
225
226// The Conn type represents a WebSocket connection.
227type Conn struct {
228	conn        net.Conn
229	isServer    bool
230	subprotocol string
231
232	// Write fields
233	mu            chan bool // used as mutex to protect write to conn
234	writeBuf      []byte    // frame is constructed in this buffer.
235	writeDeadline time.Time
236	writer        io.WriteCloser // the current writer returned to the application
237	isWriting     bool           // for best-effort concurrent write detection
238
239	writeErrMu sync.Mutex
240	writeErr   error
241
242	enableWriteCompression bool
243	compressionLevel       int
244	newCompressionWriter   func(io.WriteCloser, int) io.WriteCloser
245
246	// Read fields
247	reader        io.ReadCloser // the current reader returned to the application
248	readErr       error
249	br            *bufio.Reader
250	readRemaining int64 // bytes remaining in current frame.
251	readFinal     bool  // true the current message has more frames.
252	readLength    int64 // Message size.
253	readLimit     int64 // Maximum message size.
254	readMaskPos   int
255	readMaskKey   [4]byte
256	handlePong    func(string) error
257	handlePing    func(string) error
258	handleClose   func(int, string) error
259	readErrCount  int
260	messageReader *messageReader // the current low-level reader
261
262	readDecompress         bool // whether last read frame had RSV1 set
263	newDecompressionReader func(io.Reader) io.ReadCloser
264}
265
266func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int) *Conn {
267	return newConnBRW(conn, isServer, readBufferSize, writeBufferSize, nil)
268}
269
270type writeHook struct {
271	p []byte
272}
273
274func (wh *writeHook) Write(p []byte) (int, error) {
275	wh.p = p
276	return len(p), nil
277}
278
279func newConnBRW(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int, brw *bufio.ReadWriter) *Conn {
280	mu := make(chan bool, 1)
281	mu <- true
282
283	var br *bufio.Reader
284	if readBufferSize == 0 && brw != nil && brw.Reader != nil {
285		// Reuse the supplied bufio.Reader if the buffer has a useful size.
286		// This code assumes that peek on a reader returns
287		// bufio.Reader.buf[:0].
288		brw.Reader.Reset(conn)
289		if p, err := brw.Reader.Peek(0); err == nil && cap(p) >= 256 {
290			br = brw.Reader
291		}
292	}
293	if br == nil {
294		if readBufferSize == 0 {
295			readBufferSize = defaultReadBufferSize
296		}
297		if readBufferSize < maxControlFramePayloadSize {
298			readBufferSize = maxControlFramePayloadSize
299		}
300		br = bufio.NewReaderSize(conn, readBufferSize)
301	}
302
303	var writeBuf []byte
304	if writeBufferSize == 0 && brw != nil && brw.Writer != nil {
305		// Use the bufio.Writer's buffer if the buffer has a useful size. This
306		// code assumes that bufio.Writer.buf[:1] is passed to the
307		// bufio.Writer's underlying writer.
308		var wh writeHook
309		brw.Writer.Reset(&wh)
310		brw.Writer.WriteByte(0)
311		brw.Flush()
312		if cap(wh.p) >= maxFrameHeaderSize+256 {
313			writeBuf = wh.p[:cap(wh.p)]
314		}
315	}
316
317	if writeBuf == nil {
318		if writeBufferSize == 0 {
319			writeBufferSize = defaultWriteBufferSize
320		}
321		writeBuf = make([]byte, writeBufferSize+maxFrameHeaderSize)
322	}
323
324	c := &Conn{
325		isServer:               isServer,
326		br:                     br,
327		conn:                   conn,
328		mu:                     mu,
329		readFinal:              true,
330		writeBuf:               writeBuf,
331		enableWriteCompression: true,
332		compressionLevel:       defaultCompressionLevel,
333	}
334	c.SetCloseHandler(nil)
335	c.SetPingHandler(nil)
336	c.SetPongHandler(nil)
337	return c
338}
339
340// Subprotocol returns the negotiated protocol for the connection.
341func (c *Conn) Subprotocol() string {
342	return c.subprotocol
343}
344
345// Close closes the underlying network connection without sending or waiting
346// for a close message.
347func (c *Conn) Close() error {
348	return c.conn.Close()
349}
350
351// LocalAddr returns the local network address.
352func (c *Conn) LocalAddr() net.Addr {
353	return c.conn.LocalAddr()
354}
355
356// RemoteAddr returns the remote network address.
357func (c *Conn) RemoteAddr() net.Addr {
358	return c.conn.RemoteAddr()
359}
360
361// Write methods
362
363func (c *Conn) writeFatal(err error) error {
364	err = hideTempErr(err)
365	c.writeErrMu.Lock()
366	if c.writeErr == nil {
367		c.writeErr = err
368	}
369	c.writeErrMu.Unlock()
370	return err
371}
372
373func (c *Conn) write(frameType int, deadline time.Time, buf0, buf1 []byte) error {
374	<-c.mu
375	defer func() { c.mu <- true }()
376
377	c.writeErrMu.Lock()
378	err := c.writeErr
379	c.writeErrMu.Unlock()
380	if err != nil {
381		return err
382	}
383
384	c.conn.SetWriteDeadline(deadline)
385	if len(buf1) == 0 {
386		_, err = c.conn.Write(buf0)
387	} else {
388		err = c.writeBufs(buf0, buf1)
389	}
390	if err != nil {
391		return c.writeFatal(err)
392	}
393	if frameType == CloseMessage {
394		c.writeFatal(ErrCloseSent)
395	}
396	return nil
397}
398
399// WriteControl writes a control message with the given deadline. The allowed
400// message types are CloseMessage, PingMessage and PongMessage.
401func (c *Conn) WriteControl(messageType int, data []byte, deadline time.Time) error {
402	if !isControl(messageType) {
403		return errBadWriteOpCode
404	}
405	if len(data) > maxControlFramePayloadSize {
406		return errInvalidControlFrame
407	}
408
409	b0 := byte(messageType) | finalBit
410	b1 := byte(len(data))
411	if !c.isServer {
412		b1 |= maskBit
413	}
414
415	buf := make([]byte, 0, maxFrameHeaderSize+maxControlFramePayloadSize)
416	buf = append(buf, b0, b1)
417
418	if c.isServer {
419		buf = append(buf, data...)
420	} else {
421		key := newMaskKey()
422		buf = append(buf, key[:]...)
423		buf = append(buf, data...)
424		maskBytes(key, 0, buf[6:])
425	}
426
427	d := time.Hour * 1000
428	if !deadline.IsZero() {
429		d = deadline.Sub(time.Now())
430		if d < 0 {
431			return errWriteTimeout
432		}
433	}
434
435	timer := time.NewTimer(d)
436	select {
437	case <-c.mu:
438		timer.Stop()
439	case <-timer.C:
440		return errWriteTimeout
441	}
442	defer func() { c.mu <- true }()
443
444	c.writeErrMu.Lock()
445	err := c.writeErr
446	c.writeErrMu.Unlock()
447	if err != nil {
448		return err
449	}
450
451	c.conn.SetWriteDeadline(deadline)
452	_, err = c.conn.Write(buf)
453	if err != nil {
454		return c.writeFatal(err)
455	}
456	if messageType == CloseMessage {
457		c.writeFatal(ErrCloseSent)
458	}
459	return err
460}
461
462func (c *Conn) prepWrite(messageType int) error {
463	// Close previous writer if not already closed by the application. It's
464	// probably better to return an error in this situation, but we cannot
465	// change this without breaking existing applications.
466	if c.writer != nil {
467		c.writer.Close()
468		c.writer = nil
469	}
470
471	if !isControl(messageType) && !isData(messageType) {
472		return errBadWriteOpCode
473	}
474
475	c.writeErrMu.Lock()
476	err := c.writeErr
477	c.writeErrMu.Unlock()
478	return err
479}
480
481// NextWriter returns a writer for the next message to send. The writer's Close
482// method flushes the complete message to the network.
483//
484// There can be at most one open writer on a connection. NextWriter closes the
485// previous writer if the application has not already done so.
486//
487// All message types (TextMessage, BinaryMessage, CloseMessage, PingMessage and
488// PongMessage) are supported.
489func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error) {
490	if err := c.prepWrite(messageType); err != nil {
491		return nil, err
492	}
493
494	mw := &messageWriter{
495		c:         c,
496		frameType: messageType,
497		pos:       maxFrameHeaderSize,
498	}
499	c.writer = mw
500	if c.newCompressionWriter != nil && c.enableWriteCompression && isData(messageType) {
501		w := c.newCompressionWriter(c.writer, c.compressionLevel)
502		mw.compress = true
503		c.writer = w
504	}
505	return c.writer, nil
506}
507
508type messageWriter struct {
509	c         *Conn
510	compress  bool // whether next call to flushFrame should set RSV1
511	pos       int  // end of data in writeBuf.
512	frameType int  // type of the current frame.
513	err       error
514}
515
516func (w *messageWriter) fatal(err error) error {
517	if w.err != nil {
518		w.err = err
519		w.c.writer = nil
520	}
521	return err
522}
523
524// flushFrame writes buffered data and extra as a frame to the network. The
525// final argument indicates that this is the last frame in the message.
526func (w *messageWriter) flushFrame(final bool, extra []byte) error {
527	c := w.c
528	length := w.pos - maxFrameHeaderSize + len(extra)
529
530	// Check for invalid control frames.
531	if isControl(w.frameType) &&
532		(!final || length > maxControlFramePayloadSize) {
533		return w.fatal(errInvalidControlFrame)
534	}
535
536	b0 := byte(w.frameType)
537	if final {
538		b0 |= finalBit
539	}
540	if w.compress {
541		b0 |= rsv1Bit
542	}
543	w.compress = false
544
545	b1 := byte(0)
546	if !c.isServer {
547		b1 |= maskBit
548	}
549
550	// Assume that the frame starts at beginning of c.writeBuf.
551	framePos := 0
552	if c.isServer {
553		// Adjust up if mask not included in the header.
554		framePos = 4
555	}
556
557	switch {
558	case length >= 65536:
559		c.writeBuf[framePos] = b0
560		c.writeBuf[framePos+1] = b1 | 127
561		binary.BigEndian.PutUint64(c.writeBuf[framePos+2:], uint64(length))
562	case length > 125:
563		framePos += 6
564		c.writeBuf[framePos] = b0
565		c.writeBuf[framePos+1] = b1 | 126
566		binary.BigEndian.PutUint16(c.writeBuf[framePos+2:], uint16(length))
567	default:
568		framePos += 8
569		c.writeBuf[framePos] = b0
570		c.writeBuf[framePos+1] = b1 | byte(length)
571	}
572
573	if !c.isServer {
574		key := newMaskKey()
575		copy(c.writeBuf[maxFrameHeaderSize-4:], key[:])
576		maskBytes(key, 0, c.writeBuf[maxFrameHeaderSize:w.pos])
577		if len(extra) > 0 {
578			return c.writeFatal(errors.New("websocket: internal error, extra used in client mode"))
579		}
580	}
581
582	// Write the buffers to the connection with best-effort detection of
583	// concurrent writes. See the concurrency section in the package
584	// documentation for more info.
585
586	if c.isWriting {
587		panic("concurrent write to websocket connection")
588	}
589	c.isWriting = true
590
591	err := c.write(w.frameType, c.writeDeadline, c.writeBuf[framePos:w.pos], extra)
592
593	if !c.isWriting {
594		panic("concurrent write to websocket connection")
595	}
596	c.isWriting = false
597
598	if err != nil {
599		return w.fatal(err)
600	}
601
602	if final {
603		c.writer = nil
604		return nil
605	}
606
607	// Setup for next frame.
608	w.pos = maxFrameHeaderSize
609	w.frameType = continuationFrame
610	return nil
611}
612
613func (w *messageWriter) ncopy(max int) (int, error) {
614	n := len(w.c.writeBuf) - w.pos
615	if n <= 0 {
616		if err := w.flushFrame(false, nil); err != nil {
617			return 0, err
618		}
619		n = len(w.c.writeBuf) - w.pos
620	}
621	if n > max {
622		n = max
623	}
624	return n, nil
625}
626
627func (w *messageWriter) Write(p []byte) (int, error) {
628	if w.err != nil {
629		return 0, w.err
630	}
631
632	if len(p) > 2*len(w.c.writeBuf) && w.c.isServer {
633		// Don't buffer large messages.
634		err := w.flushFrame(false, p)
635		if err != nil {
636			return 0, err
637		}
638		return len(p), nil
639	}
640
641	nn := len(p)
642	for len(p) > 0 {
643		n, err := w.ncopy(len(p))
644		if err != nil {
645			return 0, err
646		}
647		copy(w.c.writeBuf[w.pos:], p[:n])
648		w.pos += n
649		p = p[n:]
650	}
651	return nn, nil
652}
653
654func (w *messageWriter) WriteString(p string) (int, error) {
655	if w.err != nil {
656		return 0, w.err
657	}
658
659	nn := len(p)
660	for len(p) > 0 {
661		n, err := w.ncopy(len(p))
662		if err != nil {
663			return 0, err
664		}
665		copy(w.c.writeBuf[w.pos:], p[:n])
666		w.pos += n
667		p = p[n:]
668	}
669	return nn, nil
670}
671
672func (w *messageWriter) ReadFrom(r io.Reader) (nn int64, err error) {
673	if w.err != nil {
674		return 0, w.err
675	}
676	for {
677		if w.pos == len(w.c.writeBuf) {
678			err = w.flushFrame(false, nil)
679			if err != nil {
680				break
681			}
682		}
683		var n int
684		n, err = r.Read(w.c.writeBuf[w.pos:])
685		w.pos += n
686		nn += int64(n)
687		if err != nil {
688			if err == io.EOF {
689				err = nil
690			}
691			break
692		}
693	}
694	return nn, err
695}
696
697func (w *messageWriter) Close() error {
698	if w.err != nil {
699		return w.err
700	}
701	if err := w.flushFrame(true, nil); err != nil {
702		return err
703	}
704	w.err = errWriteClosed
705	return nil
706}
707
708// WritePreparedMessage writes prepared message into connection.
709func (c *Conn) WritePreparedMessage(pm *PreparedMessage) error {
710	frameType, frameData, err := pm.frame(prepareKey{
711		isServer:         c.isServer,
712		compress:         c.newCompressionWriter != nil && c.enableWriteCompression && isData(pm.messageType),
713		compressionLevel: c.compressionLevel,
714	})
715	if err != nil {
716		return err
717	}
718	if c.isWriting {
719		panic("concurrent write to websocket connection")
720	}
721	c.isWriting = true
722	err = c.write(frameType, c.writeDeadline, frameData, nil)
723	if !c.isWriting {
724		panic("concurrent write to websocket connection")
725	}
726	c.isWriting = false
727	return err
728}
729
730// WriteMessage is a helper method for getting a writer using NextWriter,
731// writing the message and closing the writer.
732func (c *Conn) WriteMessage(messageType int, data []byte) error {
733
734	if c.isServer && (c.newCompressionWriter == nil || !c.enableWriteCompression) {
735		// Fast path with no allocations and single frame.
736
737		if err := c.prepWrite(messageType); err != nil {
738			return err
739		}
740		mw := messageWriter{c: c, frameType: messageType, pos: maxFrameHeaderSize}
741		n := copy(c.writeBuf[mw.pos:], data)
742		mw.pos += n
743		data = data[n:]
744		return mw.flushFrame(true, data)
745	}
746
747	w, err := c.NextWriter(messageType)
748	if err != nil {
749		return err
750	}
751	if _, err = w.Write(data); err != nil {
752		return err
753	}
754	return w.Close()
755}
756
757// SetWriteDeadline sets the write deadline on the underlying network
758// connection. After a write has timed out, the websocket state is corrupt and
759// all future writes will return an error. A zero value for t means writes will
760// not time out.
761func (c *Conn) SetWriteDeadline(t time.Time) error {
762	c.writeDeadline = t
763	return nil
764}
765
766// Read methods
767
768func (c *Conn) advanceFrame() (int, error) {
769	// 1. Skip remainder of previous frame.
770
771	if c.readRemaining > 0 {
772		if _, err := io.CopyN(ioutil.Discard, c.br, c.readRemaining); err != nil {
773			return noFrame, err
774		}
775	}
776
777	// 2. Read and parse first two bytes of frame header.
778
779	p, err := c.read(2)
780	if err != nil {
781		return noFrame, err
782	}
783
784	final := p[0]&finalBit != 0
785	frameType := int(p[0] & 0xf)
786	mask := p[1]&maskBit != 0
787	c.readRemaining = int64(p[1] & 0x7f)
788
789	c.readDecompress = false
790	if c.newDecompressionReader != nil && (p[0]&rsv1Bit) != 0 {
791		c.readDecompress = true
792		p[0] &^= rsv1Bit
793	}
794
795	if rsv := p[0] & (rsv1Bit | rsv2Bit | rsv3Bit); rsv != 0 {
796		return noFrame, c.handleProtocolError("unexpected reserved bits 0x" + strconv.FormatInt(int64(rsv), 16))
797	}
798
799	switch frameType {
800	case CloseMessage, PingMessage, PongMessage:
801		if c.readRemaining > maxControlFramePayloadSize {
802			return noFrame, c.handleProtocolError("control frame length > 125")
803		}
804		if !final {
805			return noFrame, c.handleProtocolError("control frame not final")
806		}
807	case TextMessage, BinaryMessage:
808		if !c.readFinal {
809			return noFrame, c.handleProtocolError("message start before final message frame")
810		}
811		c.readFinal = final
812	case continuationFrame:
813		if c.readFinal {
814			return noFrame, c.handleProtocolError("continuation after final message frame")
815		}
816		c.readFinal = final
817	default:
818		return noFrame, c.handleProtocolError("unknown opcode " + strconv.Itoa(frameType))
819	}
820
821	// 3. Read and parse frame length.
822
823	switch c.readRemaining {
824	case 126:
825		p, err := c.read(2)
826		if err != nil {
827			return noFrame, err
828		}
829		c.readRemaining = int64(binary.BigEndian.Uint16(p))
830	case 127:
831		p, err := c.read(8)
832		if err != nil {
833			return noFrame, err
834		}
835		c.readRemaining = int64(binary.BigEndian.Uint64(p))
836	}
837
838	// 4. Handle frame masking.
839
840	if mask != c.isServer {
841		return noFrame, c.handleProtocolError("incorrect mask flag")
842	}
843
844	if mask {
845		c.readMaskPos = 0
846		p, err := c.read(len(c.readMaskKey))
847		if err != nil {
848			return noFrame, err
849		}
850		copy(c.readMaskKey[:], p)
851	}
852
853	// 5. For text and binary messages, enforce read limit and return.
854
855	if frameType == continuationFrame || frameType == TextMessage || frameType == BinaryMessage {
856
857		c.readLength += c.readRemaining
858		if c.readLimit > 0 && c.readLength > c.readLimit {
859			c.WriteControl(CloseMessage, FormatCloseMessage(CloseMessageTooBig, ""), time.Now().Add(writeWait))
860			return noFrame, ErrReadLimit
861		}
862
863		return frameType, nil
864	}
865
866	// 6. Read control frame payload.
867
868	var payload []byte
869	if c.readRemaining > 0 {
870		payload, err = c.read(int(c.readRemaining))
871		c.readRemaining = 0
872		if err != nil {
873			return noFrame, err
874		}
875		if c.isServer {
876			maskBytes(c.readMaskKey, 0, payload)
877		}
878	}
879
880	// 7. Process control frame payload.
881
882	switch frameType {
883	case PongMessage:
884		if err := c.handlePong(string(payload)); err != nil {
885			return noFrame, err
886		}
887	case PingMessage:
888		if err := c.handlePing(string(payload)); err != nil {
889			return noFrame, err
890		}
891	case CloseMessage:
892		closeCode := CloseNoStatusReceived
893		closeText := ""
894		if len(payload) >= 2 {
895			closeCode = int(binary.BigEndian.Uint16(payload))
896			if !isValidReceivedCloseCode(closeCode) {
897				return noFrame, c.handleProtocolError("invalid close code")
898			}
899			closeText = string(payload[2:])
900			if !utf8.ValidString(closeText) {
901				return noFrame, c.handleProtocolError("invalid utf8 payload in close frame")
902			}
903		}
904		if err := c.handleClose(closeCode, closeText); err != nil {
905			return noFrame, err
906		}
907		return noFrame, &CloseError{Code: closeCode, Text: closeText}
908	}
909
910	return frameType, nil
911}
912
913func (c *Conn) handleProtocolError(message string) error {
914	c.WriteControl(CloseMessage, FormatCloseMessage(CloseProtocolError, message), time.Now().Add(writeWait))
915	return errors.New("websocket: " + message)
916}
917
918// NextReader returns the next data message received from the peer. The
919// returned messageType is either TextMessage or BinaryMessage.
920//
921// There can be at most one open reader on a connection. NextReader discards
922// the previous message if the application has not already consumed it.
923//
924// Applications must break out of the application's read loop when this method
925// returns a non-nil error value. Errors returned from this method are
926// permanent. Once this method returns a non-nil error, all subsequent calls to
927// this method return the same error.
928func (c *Conn) NextReader() (messageType int, r io.Reader, err error) {
929	// Close previous reader, only relevant for decompression.
930	if c.reader != nil {
931		c.reader.Close()
932		c.reader = nil
933	}
934
935	c.messageReader = nil
936	c.readLength = 0
937
938	for c.readErr == nil {
939		frameType, err := c.advanceFrame()
940		if err != nil {
941			c.readErr = hideTempErr(err)
942			break
943		}
944		if frameType == TextMessage || frameType == BinaryMessage {
945			c.messageReader = &messageReader{c}
946			c.reader = c.messageReader
947			if c.readDecompress {
948				c.reader = c.newDecompressionReader(c.reader)
949			}
950			return frameType, c.reader, nil
951		}
952	}
953
954	// Applications that do handle the error returned from this method spin in
955	// tight loop on connection failure. To help application developers detect
956	// this error, panic on repeated reads to the failed connection.
957	c.readErrCount++
958	if c.readErrCount >= 1000 {
959		panic("repeated read on failed websocket connection")
960	}
961
962	return noFrame, nil, c.readErr
963}
964
965type messageReader struct{ c *Conn }
966
967func (r *messageReader) Read(b []byte) (int, error) {
968	c := r.c
969	if c.messageReader != r {
970		return 0, io.EOF
971	}
972
973	for c.readErr == nil {
974
975		if c.readRemaining > 0 {
976			if int64(len(b)) > c.readRemaining {
977				b = b[:c.readRemaining]
978			}
979			n, err := c.br.Read(b)
980			c.readErr = hideTempErr(err)
981			if c.isServer {
982				c.readMaskPos = maskBytes(c.readMaskKey, c.readMaskPos, b[:n])
983			}
984			c.readRemaining -= int64(n)
985			if c.readRemaining > 0 && c.readErr == io.EOF {
986				c.readErr = errUnexpectedEOF
987			}
988			return n, c.readErr
989		}
990
991		if c.readFinal {
992			c.messageReader = nil
993			return 0, io.EOF
994		}
995
996		frameType, err := c.advanceFrame()
997		switch {
998		case err != nil:
999			c.readErr = hideTempErr(err)
1000		case frameType == TextMessage || frameType == BinaryMessage:
1001			c.readErr = errors.New("websocket: internal error, unexpected text or binary in Reader")
1002		}
1003	}
1004
1005	err := c.readErr
1006	if err == io.EOF && c.messageReader == r {
1007		err = errUnexpectedEOF
1008	}
1009	return 0, err
1010}
1011
1012func (r *messageReader) Close() error {
1013	return nil
1014}
1015
1016// ReadMessage is a helper method for getting a reader using NextReader and
1017// reading from that reader to a buffer.
1018func (c *Conn) ReadMessage() (messageType int, p []byte, err error) {
1019	var r io.Reader
1020	messageType, r, err = c.NextReader()
1021	if err != nil {
1022		return messageType, nil, err
1023	}
1024	p, err = ioutil.ReadAll(r)
1025	return messageType, p, err
1026}
1027
1028// SetReadDeadline sets the read deadline on the underlying network connection.
1029// After a read has timed out, the websocket connection state is corrupt and
1030// all future reads will return an error. A zero value for t means reads will
1031// not time out.
1032func (c *Conn) SetReadDeadline(t time.Time) error {
1033	return c.conn.SetReadDeadline(t)
1034}
1035
1036// SetReadLimit sets the maximum size for a message read from the peer. If a
1037// message exceeds the limit, the connection sends a close message to the peer
1038// and returns ErrReadLimit to the application.
1039func (c *Conn) SetReadLimit(limit int64) {
1040	c.readLimit = limit
1041}
1042
1043// CloseHandler returns the current close handler
1044func (c *Conn) CloseHandler() func(code int, text string) error {
1045	return c.handleClose
1046}
1047
1048// SetCloseHandler sets the handler for close messages received from the peer.
1049// The code argument to h is the received close code or CloseNoStatusReceived
1050// if the close message is empty. The default close handler sends a close
1051// message back to the peer.
1052//
1053// The handler function is called from the NextReader, ReadMessage and message
1054// reader Read methods. The application must read the connection to process
1055// close messages as described in the section on Control Messages above.
1056//
1057// The connection read methods return a CloseError when a close message is
1058// received. Most applications should handle close messages as part of their
1059// normal error handling. Applications should only set a close handler when the
1060// application must perform some action before sending a close message back to
1061// the peer.
1062func (c *Conn) SetCloseHandler(h func(code int, text string) error) {
1063	if h == nil {
1064		h = func(code int, text string) error {
1065			message := FormatCloseMessage(code, "")
1066			c.WriteControl(CloseMessage, message, time.Now().Add(writeWait))
1067			return nil
1068		}
1069	}
1070	c.handleClose = h
1071}
1072
1073// PingHandler returns the current ping handler
1074func (c *Conn) PingHandler() func(appData string) error {
1075	return c.handlePing
1076}
1077
1078// SetPingHandler sets the handler for ping messages received from the peer.
1079// The appData argument to h is the PING message application data. The default
1080// ping handler sends a pong to the peer.
1081//
1082// The handler function is called from the NextReader, ReadMessage and message
1083// reader Read methods. The application must read the connection to process
1084// ping messages as described in the section on Control Messages above.
1085func (c *Conn) SetPingHandler(h func(appData string) error) {
1086	if h == nil {
1087		h = func(message string) error {
1088			err := c.WriteControl(PongMessage, []byte(message), time.Now().Add(writeWait))
1089			if err == ErrCloseSent {
1090				return nil
1091			} else if e, ok := err.(net.Error); ok && e.Temporary() {
1092				return nil
1093			}
1094			return err
1095		}
1096	}
1097	c.handlePing = h
1098}
1099
1100// PongHandler returns the current pong handler
1101func (c *Conn) PongHandler() func(appData string) error {
1102	return c.handlePong
1103}
1104
1105// SetPongHandler sets the handler for pong messages received from the peer.
1106// The appData argument to h is the PONG message application data. The default
1107// pong handler does nothing.
1108//
1109// The handler function is called from the NextReader, ReadMessage and message
1110// reader Read methods. The application must read the connection to process
1111// pong messages as described in the section on Control Messages above.
1112func (c *Conn) SetPongHandler(h func(appData string) error) {
1113	if h == nil {
1114		h = func(string) error { return nil }
1115	}
1116	c.handlePong = h
1117}
1118
1119// UnderlyingConn returns the internal net.Conn. This can be used to further
1120// modifications to connection specific flags.
1121func (c *Conn) UnderlyingConn() net.Conn {
1122	return c.conn
1123}
1124
1125// EnableWriteCompression enables and disables write compression of
1126// subsequent text and binary messages. This function is a noop if
1127// compression was not negotiated with the peer.
1128func (c *Conn) EnableWriteCompression(enable bool) {
1129	c.enableWriteCompression = enable
1130}
1131
1132// SetCompressionLevel sets the flate compression level for subsequent text and
1133// binary messages. This function is a noop if compression was not negotiated
1134// with the peer. See the compress/flate package for a description of
1135// compression levels.
1136func (c *Conn) SetCompressionLevel(level int) error {
1137	if !isValidCompressionLevel(level) {
1138		return errors.New("websocket: invalid compression level")
1139	}
1140	c.compressionLevel = level
1141	return nil
1142}
1143
1144// FormatCloseMessage formats closeCode and text as a WebSocket close message.
1145// An empty message is returned for code CloseNoStatusReceived.
1146func FormatCloseMessage(closeCode int, text string) []byte {
1147	if closeCode == CloseNoStatusReceived {
1148		// Return empty message because it's illegal to send
1149		// CloseNoStatusReceived. Return non-nil value in case application
1150		// checks for nil.
1151		return []byte{}
1152	}
1153	buf := make([]byte, 2+len(text))
1154	binary.BigEndian.PutUint16(buf, uint16(closeCode))
1155	copy(buf[2:], text)
1156	return buf
1157}
1158