ws/ws/client.go

package ws

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

// 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){

	errorCode := NORMAL
	clientAck := true
	c.io.reading.Add(1)

	for {

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

		/*** Parse message ***/
		msg, err := readMessage(c.io.reader)
		if err != nil {
			// fmt.Printf(" [reader] %s\n", err)
			break
		}

		/* (4) CLOSE */
		if msg.Type == CLOSE {
			// fmt.Printf(" [reader] CLOSE ; size %d\n", msg.Size)
			// if msg.Size >= 2 {
			// 	errCode := binary.BigEndian.Uint16(msg.Data[0:2])
			// 	fmt.Printf("                ; status %d\n", errCode)
			// 	fmt.Printf("                ; msg '%s'\n", msg.Data[2:])
			// }
			clientAck = false
			break

		}

		/* (5) PING size error */
		if msg.Type == PING && msg.Size > 125 {
			fmt.Printf(" [reader] PING payload too big\n")
			// fmt.Printf("[reader] PING err\n")
			errorCode = PROTOCOL_ERR
			break
		}

		/* (6) Send PONG */
		if msg.Type == PING {
			// fmt.Printf("[reader] PING -> PONG\n")
			if c.io.writing {
				msg.Final = true
				msg.Type  = PONG
				c.ch.send <- *msg
			}
			continue
		}

		/* (7) Invalid UTF8 */
		if msg.Type == TEXT && !utf8.Valid(msg.Data) {
			fmt.Printf(" [reader] invalid utf-8\n")
			errorCode = INVALID_PAYLOAD
			break
		}

		/* (8) Unknown opcode */
		if msg.Type != TEXT && msg.Type != BINARY {
			fmt.Printf(" [reader] unknown OpCode %d\n", msg.Type)
			errorCode = PROTOCOL_ERR
			break
		}

		/* (9) Dispatch to receiver */
		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(errorCode, 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)

}




// close writes the error message (if needed)
// and it closes the socket
// if 'clientACK' is true, reads the next message (CLOSE acknowledge)
// before closing the socket
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 {
		status = NORMAL
	}

	/* (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 */
	err := msg.Send(c.io.sock)
	if err != nil {
		fmt.Printf("[close] send error (%s0\n", err)
	}
	// fmt.Printf("[close] frame sent\n")


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

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

		/* Wait for message */
		msg, err := 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

}