ws/client.go

package websocket

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"git.xdrm.io/go/websocket/internal/http/upgrade/request"
	"net"
	"sync"
	"time"
)

// Represents a client socket utility (reader, writer, ..)
type clientIO struct {
	sock      net.Conn
	reader    *bufio.Reader
	kill      chan<- *client // unregisters client
	closing   bool
	closingMu sync.Mutex
	reading   sync.WaitGroup
	writing   bool
}

// Represents all channels that need a client
type clientChannelSet struct {
	receive chan Message
	send    chan Message
}

// Represents a websocket client
type client struct {
	io     clientIO
	iface  *Client
	ch     clientChannelSet
	status MessageError // close status ; 0 = nothing ; else -> must close
}

// Create creates a new client
func buildClient(s net.Conn, ctl ControllerSet, serverCh serverChannelSet) (*client, error) {

	/* (1) Manage UPGRADE request
	---------------------------------------------------------*/
	/* (1) Parse request */
	req, _ := request.Parse(s)

	/* (3) Build response */
	res := req.BuildResponse()

	/* (4) Write into socket */
	_, err := res.Send(s)
	if err != nil {
		return nil, fmt.Errorf("Upgrade write error: %s", err)
	}

	if res.GetStatusCode() != 101 {
		s.Close()
		return nil, fmt.Errorf("Upgrade error (HTTP %d)\n", res.GetStatusCode())
	}

	/* (2) Initialise client
	---------------------------------------------------------*/
	/* (1) Get upgrade data */
	clientURI := req.GetURI()
	clientProtocol := res.GetProtocol()

	/* (2) Initialise client */
	cli := &client{
		io: clientIO{
			sock:   s,
			reader: bufio.NewReader(s),
			kill:   serverCh.unregister,
		},

		iface: &Client{
			Protocol:  string(clientProtocol),
			Arguments: [][]string{[]string{clientURI}},
		},

		ch: clientChannelSet{
			receive: make(chan Message, 1),
			send:    make(chan Message, 1),
		},
	}

	/* (3) Find controller by URI
	---------------------------------------------------------*/
	/* (1) Try to find one */
	controller, arguments := ctl.Match(clientURI)

	/* (2) If nothing found -> error */
	if controller == nil {
		return nil, fmt.Errorf("No controller found, no default controller set\n")
	}

	/* (3) Copy arguments */
	cli.iface.Arguments = arguments

	/* (4) Launch client routines
	---------------------------------------------------------*/
	/* (1) Launch client controller */
	go controller.Fun(
		cli.iface,          // pass the client
		cli.ch.receive,     // the receiver
		cli.ch.send,        // the sender
		serverCh.broadcast, // broadcast sender
	)

	/* (2) Launch message reader */
	go clientReader(cli)

	/* (3) Launc writer */
	go clientWriter(cli)

	return cli, nil

}

// reader reads and parses messages from the buffer
func clientReader(c *client) {
	var frag *Message

	closeStatus := NORMAL
	clientAck := true

	c.io.reading.Add(1)

	for {

		/* (1) if currently closing -> exit */
		if c.io.closing {
			fmt.Printf("[reader] killed because closing")
			break
		}

		/* (2) Parse message */
		msg, err := readMessage(c.io.reader)

		if err == ErrUnmaskedFrame || err == ErrReservedBits {
			closeStatus = PROTOCOL_ERR
		}
		if err != nil {
			break
		}

		/* (3) Fail on invalid message */
		msgErr := msg.check(frag != nil)
		if msgErr != nil {

			mustClose := false

			switch msgErr {

			// Fail
			case ErrUnexpectedContinuation:
				closeStatus = NONE
				clientAck = false
				mustClose = true

			// proper close
			case CloseFrame:
				closeStatus = NORMAL
				clientAck = true
				mustClose = true

			// invalid payload proper close
			case ErrInvalidPayload:
				closeStatus = INVALID_PAYLOAD
				clientAck = true
				mustClose = true

			// any other error -> protocol error
			default:
				closeStatus = PROTOCOL_ERR
				clientAck = true
				mustClose = true
			}

			if mustClose {
				break
			}

		}

		/* (4) Ping <-> Pong */
		if msg.Type == PING && c.io.writing {
			msg.Final = true
			msg.Type = PONG
			c.ch.send <- *msg
			continue
		}

		/* (5) Store first fragment */
		if frag == nil && !msg.Final {
			frag = &Message{
				Type:  msg.Type,
				Final: msg.Final,
				Data:  msg.Data,
				Size:  msg.Size,
			}
			continue
		}

		/* (6) Store fragments */
		if frag != nil {
			frag.Final = msg.Final
			frag.Size += msg.Size
			frag.Data = append(frag.Data, msg.Data...)

			if !frag.Final { // continue if not last fragment
				continue
			}

			// check message errors
			fragErr := frag.check(false)
			if fragErr == ErrInvalidPayload {
				closeStatus = INVALID_PAYLOAD
				break
			} else if fragErr != nil {
				closeStatus = PROTOCOL_ERR
				break
			}

			msg = frag
			frag = nil

		}

		/* (7) Dispatch to receiver */
		if msg.Type == TEXT || msg.Type == BINARY {
			c.ch.receive <- *msg
		}

	}

	close(c.ch.receive)
	c.io.reading.Done()

	/* (8) close channel (if not already done) */
	// fmt.Printf("[reader] end\n")
	c.close(closeStatus, clientAck)

}

// writer writes into websocket
// and is triggered by client.ch.send channel
func clientWriter(c *client) {

	c.io.writing = true // if channel still exists

	for msg := range c.ch.send {

		/* (2) Send message */
		err := msg.Send(c.io.sock)

		/* (3) Fail on error */
		if err != nil {
			fmt.Printf(" [writer] %s\n", err)
			c.io.writing = false
			break
		}

	}

	c.io.writing = false

	/* (4) close channel (if not already done) */
	// fmt.Printf("[writer] end\n")
	c.close(NORMAL, true)

}

// closes the connection
// send CLOSE frame is 'status' is not NONE
// wait for the next message (CLOSE acknowledge) if 'clientACK'
// then delete client
func (c *client) close(status MessageError, clientACK bool) {

	/* (1) Fail if already closing */
	alreadyClosing := false
	c.io.closingMu.Lock()
	alreadyClosing = c.io.closing
	c.io.closing = true
	c.io.closingMu.Unlock()

	if alreadyClosing {
		return
	}

	/* (2) kill writer' if still running */
	if c.io.writing {
		close(c.ch.send)
	}

	/* (3) kill reader if still running */
	c.io.sock.SetReadDeadline(time.Now().Add(time.Second * -1))
	c.io.reading.Wait()

	if status != NONE {

		/* (3) Build message */
		msg := &Message{
			Final: true,
			Type:  CLOSE,
			Size:  2,
			Data:  make([]byte, 2),
		}
		binary.BigEndian.PutUint16(msg.Data, uint16(status))

		/* (4) Send message */
		msg.Send(c.io.sock)
		// if err != nil {
		// 	fmt.Printf("[close] send error (%s0\n", err)
		// }

	}

	/* (2) Wait for client CLOSE if needed */
	if clientACK {

		c.io.sock.SetReadDeadline(time.Now().Add(time.Millisecond))

		/* Wait for message */
		readMessage(c.io.reader)
		// if err != nil || msg.Type != CLOSE {
		// 	if err == nil {
		// 		fmt.Printf("[close] received OpCode = %d\n", msg.Type)
		// 	} else {
		// 		fmt.Printf("[close] read error (%v)\n", err)
		// 	}
		// }

		// fmt.Printf("[close] received ACK\n")

	}

	/* (3) Close socket */
	c.io.sock.Close()
	// fmt.Printf("[close] socket closed\n")

	/* (4) Unregister */
	c.io.kill <- c

	return

}