mchess-server/connection/type.go

package connection

import (
	"context"
	"log"

	"nhooyr.io/websocket"
)

type Connection struct {
	ws                      *websocket.Conn
	wsConnectionEstablished chan bool
	ctx                     context.Context
	buffer                  MessageBuffer
	disconnectCallback      func()
}

func NewConnection(options ...func(*Connection)) *Connection {
	connection := Connection{
		buffer:                  *newMessageBuffer(100),
		wsConnectionEstablished: make(chan bool),
	}

	for _, option := range options {
		option(&connection)
	}

	return &connection
}

func WithWebsocket(ws *websocket.Conn) func(*Connection) {
	return func(c *Connection) {
		c.ws = ws
	}
}

func WithContext(ctx context.Context) func(*Connection) {
	return func(c *Connection) {
		c.ctx = ctx
	}
}

func WithDisconnectCallback(cb func()) func(*Connection) {
	return func(c *Connection) {
		if cb != nil {
			c.disconnectCallback = cb
		}
	}
}

func (conn *Connection) SetDisconnectCallback(cb func()) {
	conn.disconnectCallback = cb
}

func (conn *Connection) HasWebsocketConnection() bool {
	return conn.ws != nil
}

func (conn *Connection) SetWebsocketConnection(ws *websocket.Conn) {
	if ws == nil {
		return
	}

	conn.ws = ws

	select {
	case conn.wsConnectionEstablished <- true:
	default:
	}

	go func() {
		for {
			_, msg, err := conn.ws.Read(conn.ctx)
			if err != nil {
				log.Println("while reading from websocket: %w", err)
				if conn.disconnectCallback != nil {
					conn.disconnectCallback()
				}
				return
			}
			conn.buffer.Insert(string(msg))
		}
	}()
}

func (conn *Connection) Write(msg []byte) error {
	if conn.ws == nil { //if ws is not yet set, we wait for it
		<-conn.wsConnectionEstablished
	}

	log.Printf("Writing message: %s", string(msg))
	return conn.ws.Write(conn.ctx, websocket.MessageText, msg)
}

func (conn *Connection) Read() ([]byte, error) {
	msg, err := conn.buffer.Get()
	if err != nil {
		conn.ws = nil
		return nil, err // TODO: Tell game-handler that connection was lost
	}

	return []byte(msg), err
}

func (conn *Connection) Close(msg string) {
	conn.ws.Close(websocket.StatusCode(400), msg)
	conn.ws = nil
}