// Package ws implements the server side of RFC 6455 WebSockets.
//
// Only what a relay needs is here: the opening handshake, masked client
// frames, fragmentation, ping/pong, close and a maximum message size. Server
// frames are never masked. Reads and writes may run concurrently (one reader
// goroutine, one writer goroutine); Write is internally serialized.
package ws

import (
	"bufio"
	"crypto/sha1"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"io"
	"net"
	"strings"
	"sync"
	"time"
	"unicode/utf8"
)

// Frame opcodes.
const (
	OpContinuation = 0x0
	OpText         = 0x1
	OpBinary       = 0x2
	OpClose        = 0x8
	OpPing         = 0x9
	OpPong         = 0xa
)

const acceptGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

var (
	// ErrProtocol is returned for a framing violation.
	ErrProtocol = errors.New("ws: protocol error")
	// ErrTooLarge is returned when a message exceeds MaxMessage.
	ErrTooLarge = errors.New("ws: message too large")
	// ErrClosed is returned after the peer closed the connection.
	ErrClosed = errors.New("ws: closed")
)

// Conn is a server-side WebSocket connection.
type Conn struct {
	conn net.Conn
	br   *bufio.Reader
	bw   *bufio.Writer

	wmu  sync.Mutex
	dead bool

	// MaxMessage bounds one (possibly fragmented) message. Zero means 1 MiB.
	MaxMessage int64
}

// Upgrade validates the handshake headers and writes the 101 response. The
// caller must have checked that the request is a GET.
func Upgrade(conn net.Conn, brw *bufio.ReadWriter, header map[string]string) (*Conn, error) {
	if !headerContains(header["upgrade"], "websocket") {
		return nil, errors.New("ws: missing Upgrade: websocket")
	}
	if !headerContains(header["connection"], "upgrade") {
		return nil, errors.New("ws: missing Connection: Upgrade")
	}
	if header["sec-websocket-version"] != "13" {
		return nil, errors.New("ws: unsupported websocket version")
	}
	key := header["sec-websocket-key"]
	if key == "" {
		return nil, errors.New("ws: missing Sec-WebSocket-Key")
	}
	sum := sha1.Sum([]byte(key + acceptGUID))
	accept := base64.StdEncoding.EncodeToString(sum[:])

	if _, err := brw.WriteString("HTTP/1.1 101 Switching Protocols\r\n" +
		"Upgrade: websocket\r\n" +
		"Connection: Upgrade\r\n" +
		"Sec-WebSocket-Accept: " + accept + "\r\n\r\n"); err != nil {
		return nil, err
	}
	if err := brw.Flush(); err != nil {
		return nil, err
	}
	return &Conn{conn: conn, br: brw.Reader, bw: brw.Writer, MaxMessage: 1 << 20}, nil
}

func headerContains(value, token string) bool {
	for _, part := range strings.Split(value, ",") {
		if strings.EqualFold(strings.TrimSpace(part), token) {
			return true
		}
	}
	return false
}

// SetReadDeadline sets the deadline for the next read.
func (c *Conn) SetReadDeadline(t time.Time) error { return c.conn.SetReadDeadline(t) }

// SetWriteDeadline sets the deadline for the next write.
func (c *Conn) SetWriteDeadline(t time.Time) error { return c.conn.SetWriteDeadline(t) }

// RemoteAddr returns the peer address.
func (c *Conn) RemoteAddr() string { return c.conn.RemoteAddr().String() }

// ReadMessage returns the next complete data message. Text messages are
// validated as UTF-8. Control frames are handled internally.
func (c *Conn) ReadMessage() (opcode byte, payload []byte, err error) {
	if c.MaxMessage <= 0 {
		c.MaxMessage = 1 << 20
	}
	var fragOp byte
	var message []byte
	inFrag := false

	for {
		hdr := make([]byte, 2)
		if _, err := io.ReadFull(c.br, hdr); err != nil {
			return 0, nil, err
		}
		fin := hdr[0]&0x80 != 0
		rsv := hdr[0] & 0x70
		op := hdr[0] & 0x0f
		masked := hdr[1]&0x80 != 0
		length := int64(hdr[1] & 0x7f)
		if rsv != 0 {
			c.fail(1002, "reserved bits set")
			return 0, nil, ErrProtocol
		}
		if !masked {
			c.fail(1002, "client frames must be masked")
			return 0, nil, ErrProtocol
		}
		switch length {
		case 126:
			ext := make([]byte, 2)
			if _, err := io.ReadFull(c.br, ext); err != nil {
				return 0, nil, err
			}
			length = int64(binary.BigEndian.Uint16(ext))
		case 127:
			ext := make([]byte, 8)
			if _, err := io.ReadFull(c.br, ext); err != nil {
				return 0, nil, err
			}
			v := binary.BigEndian.Uint64(ext)
			if v > 1<<62 {
				c.fail(1009, "frame too large")
				return 0, nil, ErrTooLarge
			}
			length = int64(v)
		}

		control := op&0x08 != 0
		if control {
			if !fin {
				c.fail(1002, "fragmented control frame")
				return 0, nil, ErrProtocol
			}
			if length > 125 {
				c.fail(1002, "oversized control frame")
				return 0, nil, ErrProtocol
			}
		}
		if length > c.MaxMessage || int64(len(message))+length > c.MaxMessage {
			c.fail(1009, "message too large")
			return 0, nil, ErrTooLarge
		}

		mask := make([]byte, 4)
		if _, err := io.ReadFull(c.br, mask); err != nil {
			return 0, nil, err
		}
		data := make([]byte, length)
		if _, err := io.ReadFull(c.br, data); err != nil {
			return 0, nil, err
		}
		for i := range data {
			data[i] ^= mask[i&3]
		}

		switch op {
		case OpPing:
			if err := c.WriteFrame(OpPong, data); err != nil {
				return 0, nil, err
			}
			continue
		case OpPong:
			continue
		case OpClose:
			code := uint16(1000)
			if len(data) >= 2 {
				code = binary.BigEndian.Uint16(data[:2])
			}
			_ = c.Close(code, "")
			return 0, nil, ErrClosed
		case OpText, OpBinary:
			if inFrag {
				c.fail(1002, "new data frame inside a fragmented message")
				return 0, nil, ErrProtocol
			}
			if fin {
				if op == OpText && !utf8.Valid(data) {
					c.fail(1007, "invalid UTF-8")
					return 0, nil, ErrProtocol
				}
				return op, data, nil
			}
			fragOp = op
			message = append(message[:0], data...)
			inFrag = true
		case OpContinuation:
			if !inFrag {
				c.fail(1002, "continuation without a message")
				return 0, nil, ErrProtocol
			}
			message = append(message, data...)
			if fin {
				if fragOp == OpText && !utf8.Valid(message) {
					c.fail(1007, "invalid UTF-8")
					return 0, nil, ErrProtocol
				}
				return fragOp, message, nil
			}
		default:
			c.fail(1002, "unknown opcode")
			return 0, nil, ErrProtocol
		}
	}
}

// WriteText sends a text message.
func (c *Conn) WriteText(data []byte) error { return c.WriteFrame(OpText, data) }

// WriteFrame sends one unfragmented frame.
func (c *Conn) WriteFrame(op byte, payload []byte) error {
	c.wmu.Lock()
	defer c.wmu.Unlock()
	if c.dead {
		return ErrClosed
	}
	var hdr [10]byte
	hdr[0] = 0x80 | op
	n := 2
	switch {
	case len(payload) < 126:
		hdr[1] = byte(len(payload))
	case len(payload) <= 0xffff:
		hdr[1] = 126
		binary.BigEndian.PutUint16(hdr[2:4], uint16(len(payload)))
		n = 4
	default:
		hdr[1] = 127
		binary.BigEndian.PutUint64(hdr[2:10], uint64(len(payload)))
		n = 10
	}
	if _, err := c.bw.Write(hdr[:n]); err != nil {
		return err
	}
	if len(payload) > 0 {
		if _, err := c.bw.Write(payload); err != nil {
			return err
		}
	}
	return c.bw.Flush()
}

// Close sends a close frame. It is not an error to call it more than once.
func (c *Conn) Close(code uint16, reason string) error {
	c.wmu.Lock()
	defer c.wmu.Unlock()
	if c.dead {
		return nil
	}
	c.dead = true
	payload := make([]byte, 2+len(reason))
	binary.BigEndian.PutUint16(payload[:2], code)
	copy(payload[2:], reason)
	var hdr [4]byte
	hdr[0] = 0x80 | OpClose
	hdr[1] = byte(len(payload))
	binary.BigEndian.PutUint16(hdr[2:4], 0)
	if _, err := c.bw.Write(hdr[:2]); err != nil {
		return err
	}
	if _, err := c.bw.Write(payload); err != nil {
		return err
	}
	return c.bw.Flush()
}

func (c *Conn) fail(code uint16, reason string) {
	_ = c.Close(code, reason)
}

// CloseSocket closes the underlying TCP connection.
func (c *Conn) CloseSocket() error { return c.conn.Close() }

// HandshakeResponse validates the opening handshake headers and returns the
// complete 101 response. Reader workers use it so that only their paired
// writer domain writes to the socket.
func HandshakeResponse(header map[string]string) ([]byte, error) {
	if !headerContains(header["upgrade"], "websocket") {
		return nil, errors.New("ws: missing Upgrade: websocket")
	}
	if !headerContains(header["connection"], "upgrade") {
		return nil, errors.New("ws: missing Connection: Upgrade")
	}
	if header["sec-websocket-version"] != "13" {
		return nil, errors.New("ws: unsupported websocket version")
	}
	key := header["sec-websocket-key"]
	if key == "" {
		return nil, errors.New("ws: missing Sec-WebSocket-Key")
	}
	sum := sha1.Sum([]byte(key + acceptGUID))
	accept := base64.StdEncoding.EncodeToString(sum[:])
	out := []byte("HTTP/1.1 101 Switching Protocols\r\n" +
		"Upgrade: websocket\r\n" +
		"Connection: Upgrade\r\n" +
		"Sec-WebSocket-Accept: " + accept + "\r\n\r\n")
	return out, nil
}