// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.

// Use of this source code is governed by a BSD-style

// license that can be found in the LICENSE file.

package websocket

import (

	"bufio"

	"encoding/binary"

	"errors"

	"io"

	"io/ioutil"

	"math/rand"

	"net"

	"strconv"

	"sync"

	"time"

	"unicode/utf8"

)

const (

	// Frame header byte 0 bits from Section 5.2 of RFC 6455

	finalBit = 1 << 7

	rsv1Bit  = 1 << 6

	rsv2Bit  = 1 << 5

	rsv3Bit  = 1 << 4

	// Frame header byte 1 bits from Section 5.2 of RFC 6455

	maskBit = 1 << 7

	maxFrameHeaderSize         = 2 + 8 + 4 // Fixed header + length + mask

	maxControlFramePayloadSize = 125

	writeWait = time.Second

	defaultReadBufferSize  = 4096

	defaultWriteBufferSize = 4096

	continuationFrame = 0

	noFrame           = -1

)

// Close codes defined in RFC 6455, section 11.7.

const (

	CloseNormalClosure           = 1000

	CloseGoingAway               = 1001

	CloseProtocolError           = 1002

	CloseUnsupportedData         = 1003

	CloseNoStatusReceived        = 1005

	CloseAbnormalClosure         = 1006

	CloseInvalidFramePayloadData = 1007

	ClosePolicyViolation         = 1008

	CloseMessageTooBig           = 1009

	CloseMandatoryExtension      = 1010

	CloseInternalServerErr       = 1011

	CloseServiceRestart          = 1012

	CloseTryAgainLater           = 1013

	CloseTLSHandshake            = 1015

)

// The message types are defined in RFC 6455, section 11.8.

const (

	// TextMessage denotes a text data message. The text message payload is

	// interpreted as UTF-8 encoded text data.

	TextMessage = 1

	// BinaryMessage denotes a binary data message.

	BinaryMessage = 2

	// CloseMessage denotes a close control message. The optional message

	// payload contains a numeric code and text. Use the FormatCloseMessage

	// function to format a close message payload.

	CloseMessage = 8

	// PingMessage denotes a ping control message. The optional message payload

	// is UTF-8 encoded text.

	PingMessage = 9

	// PongMessage denotes a pong control message. The optional message payload

	// is UTF-8 encoded text.

	PongMessage = 10

)

// ErrCloseSent is returned when the application writes a message to the

// connection after sending a close message.

var ErrCloseSent = errors.New("websocket: close sent")

// ErrReadLimit is returned when reading a message that is larger than the

// read limit set for the connection.

var ErrReadLimit = errors.New("websocket: read limit exceeded")

// netError satisfies the net Error interface.

type netError struct {

	msg       string

	temporary bool

	timeout   bool

}

func (e *netError) Error() string   { return e.msg }

func (e *netError) Temporary() bool { return e.temporary }

func (e *netError) Timeout() bool   { return e.timeout }

// CloseError represents a close message.

type CloseError struct {

	// Code is defined in RFC 6455, section 11.7.

	Code int

	// Text is the optional text payload.

	Text string

}

func (e *CloseError) Error() string {

	s := []byte("websocket: close ")

	s = strconv.AppendInt(s, int64(e.Code), 10)

	switch e.Code {

	case CloseNormalClosure:

		s = append(s, " (normal)"...)

	case CloseGoingAway:

		s = append(s, " (going away)"...)

	case CloseProtocolError:

		s = append(s, " (protocol error)"...)

	case CloseUnsupportedData:

		s = append(s, " (unsupported data)"...)

	case CloseNoStatusReceived:

		s = append(s, " (no status)"...)

	case CloseAbnormalClosure:

		s = append(s, " (abnormal closure)"...)

	case CloseInvalidFramePayloadData:

		s = append(s, " (invalid payload data)"...)

	case ClosePolicyViolation:

		s = append(s, " (policy violation)"...)

	case CloseMessageTooBig:

		s = append(s, " (message too big)"...)

	case CloseMandatoryExtension:

		s = append(s, " (mandatory extension missing)"...)

	case CloseInternalServerErr:

		s = append(s, " (internal server error)"...)

	case CloseTLSHandshake:

		s = append(s, " (TLS handshake error)"...)

	}

	if e.Text != "" {

		s = append(s, ": "...)

		s = append(s, e.Text...)

	}

	return string(s)

}

// IsCloseError returns boolean indicating whether the error is a *CloseError

// with one of the specified codes.

func IsCloseError(err error, codes ...int) bool {

	if e, ok := err.(*CloseError); ok {

		for _, code := range codes {

			if e.Code == code {

				return true

			}

		}

	}

	return false

}

// IsUnexpectedCloseError returns boolean indicating whether the error is a

// *CloseError with a code not in the list of expected codes.

func IsUnexpectedCloseError(err error, expectedCodes ...int) bool {

	if e, ok := err.(*CloseError); ok {

		for _, code := range expectedCodes {

			if e.Code == code {

				return false

			}

		}

		return true

	}

	return false

}

var (

	errWriteTimeout        = &netError{msg: "websocket: write timeout", timeout: true, temporary: true}

	errUnexpectedEOF       = &CloseError{Code: CloseAbnormalClosure, Text: io.ErrUnexpectedEOF.Error()}

	errBadWriteOpCode      = errors.New("websocket: bad write message type")

	errWriteClosed         = errors.New("websocket: write closed")

	errInvalidControlFrame = errors.New("websocket: invalid control frame")

)

func newMaskKey() [4]byte {

	n := rand.Uint32()

	return [4]byte{byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24)}

}

func hideTempErr(err error) error {

	if e, ok := err.(net.Error); ok && e.Temporary() {

		err = &netError{msg: e.Error(), timeout: e.Timeout()}

	}

	return err

}

func isControl(frameType int) bool {

	return frameType == CloseMessage || frameType == PingMessage || frameType == PongMessage

}

func isData(frameType int) bool {

	return frameType == TextMessage || frameType == BinaryMessage

}

var validReceivedCloseCodes = map[int]bool{

	// see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number

	CloseNormalClosure:           true,

	CloseGoingAway:               true,

	CloseProtocolError:           true,

	CloseUnsupportedData:         true,

	CloseNoStatusReceived:        false,

	CloseAbnormalClosure:         false,

	CloseInvalidFramePayloadData: true,

	ClosePolicyViolation:         true,

	CloseMessageTooBig:           true,

	CloseMandatoryExtension:      true,

	CloseInternalServerErr:       true,

	CloseServiceRestart:          true,

	CloseTryAgainLater:           true,

	CloseTLSHandshake:            false,

}

func isValidReceivedCloseCode(code int) bool {

	return validReceivedCloseCodes[code] || (code >= 3000 && code <= 4999)

}

// BufferPool represents a pool of buffers. The *sync.Pool type satisfies this

// interface.  The type of the value stored in a pool is not specified.

type BufferPool interface {

	// Get gets a value from the pool or returns nil if the pool is empty.

	Get() interface{}

	// Put adds a value to the pool.

	Put(interface{})

}

// writePoolData is the type added to the write buffer pool. This wrapper is

// used to prevent applications from peeking at and depending on the values

// added to the pool.

type writePoolData struct{ buf []byte }

// The Conn type represents a WebSocket connection.

type Conn struct {

	conn        net.Conn

	isServer    bool

	subprotocol string

	// Write fields

	mu            chan bool // used as mutex to protect write to conn

	writeBuf      []byte    // frame is constructed in this buffer.

	writePool     BufferPool

	writeBufSize  int

	writeDeadline time.Time

	writer        io.WriteCloser // the current writer returned to the application

	isWriting     bool           // for best-effort concurrent write detection

	writeErrMu sync.Mutex

	writeErr   error

	enableWriteCompression bool

	compressionLevel       int

	newCompressionWriter   func(io.WriteCloser, int) io.WriteCloser

	// Read fields

	reader        io.ReadCloser // the current reader returned to the application

	readErr       error

	br            *bufio.Reader

	readRemaining int64 // bytes remaining in current frame.

	readFinal     bool  // true the current message has more frames.

	readLength    int64 // Message size.

	readLimit     int64 // Maximum message size.

	readMaskPos   int

	readMaskKey   [4]byte

	handlePong    func(string) error

	handlePing    func(string) error

	handleClose   func(int, string) error

	readErrCount  int

	messageReader *messageReader // the current low-level reader

	readDecompress         bool // whether last read frame had RSV1 set

	newDecompressionReader func(io.Reader) io.ReadCloser

}

func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int, writeBufferPool BufferPool, br *bufio.Reader, writeBuf []byte) *Conn {

	if br == nil {

		if readBufferSize == 0 {

			readBufferSize = defaultReadBufferSize

		} else if readBufferSize < maxControlFramePayloadSize {

			// must be large enough for control frame

			readBufferSize = maxControlFramePayloadSize

		}

		br = bufio.NewReaderSize(conn, readBufferSize)

	}

	if writeBufferSize <= 0 {

		writeBufferSize = defaultWriteBufferSize

	}

	writeBufferSize += maxFrameHeaderSize

	if writeBuf == nil && writeBufferPool == nil {

		writeBuf = make([]byte, writeBufferSize)

	}

	mu := make(chan bool, 1)

	mu <- true

	c := &Conn{

		isServer:               isServer,

		br:                     br,

		conn:                   conn,

		mu:                     mu,

		readFinal:              true,

		writeBuf:               writeBuf,

		writePool:              writeBufferPool,

		writeBufSize:           writeBufferSize,

		enableWriteCompression: true,

		compressionLevel:       defaultCompressionLevel,

	}

	c.SetCloseHandler(nil)

	c.SetPingHandler(nil)

	c.SetPongHandler(nil)

	return c

}

// Subprotocol returns the negotiated protocol for the connection.

func (c *Conn) Subprotocol() string {

	return c.subprotocol

}

// Close closes the underlying network connection without sending or waiting

// for a close message.

func (c *Conn) Close() error {

	return c.conn.Close()

}

// LocalAddr returns the local network address.

func (c *Conn) LocalAddr() net.Addr {

	return c.conn.LocalAddr()

}

// RemoteAddr returns the remote network address.

func (c *Conn) RemoteAddr() net.Addr {

	return c.conn.RemoteAddr()

}

// Write methods

func (c *Conn) writeFatal(err error) error {

	err = hideTempErr(err)

	c.writeErrMu.Lock()

	if c.writeErr == nil {

		c.writeErr = err

	}

	c.writeErrMu.Unlock()

	return err

}

func (c *Conn) read(n int) ([]byte, error) {

	p, err := c.br.Peek(n)

	if err == io.EOF {

		err = errUnexpectedEOF

	}

	c.br.Discard(len(p))

	return p, err

}

func (c *Conn) write(frameType int, deadline time.Time, buf0, buf1 []byte) error {

	<-c.mu

	defer func() { c.mu <- true }()

	c.writeErrMu.Lock()

	err := c.writeErr

	c.writeErrMu.Unlock()

	if err != nil {

		return err

	}

	c.conn.SetWriteDeadline(deadline)

	if len(buf1) == 0 {

		_, err = c.conn.Write(buf0)

	} else {

		err = c.writeBufs(buf0, buf1)

	}

	if err != nil {

		return c.writeFatal(err)

	}

	if frameType == CloseMessage {

		c.writeFatal(ErrCloseSent)

	}

	return nil

}

// WriteControl writes a control message with the given deadline. The allowed

// message types are CloseMessage, PingMessage and PongMessage.

func (c *Conn) WriteControl(messageType int, data []byte, deadline time.Time) error {

	if !isControl(messageType) {

		return errBadWriteOpCode

	}

	if len(data) > maxControlFramePayloadSize {

		return errInvalidControlFrame

	}

	b0 := byte(messageType) | finalBit

	b1 := byte(len(data))

	if !c.isServer {

		b1 |= maskBit

	}

	buf := make([]byte, 0, maxFrameHeaderSize+maxControlFramePayloadSize)

	buf = append(buf, b0, b1)

	if c.isServer {

		buf = append(buf, data...)

	} else {

		key := newMaskKey()

		buf = append(buf, key[:]...)

		buf = append(buf, data...)

		maskBytes(key, 0, buf[6:])

	}

	d := time.Hour * 1000