refactor: refactor keyboard, mouse, and WebSocket

server/service/hid/hid.go

@@ -18,9 +18,9 @@ type Hid struct {
 }
 
 const (
-	HID0 = "/dev/hidg0"
-	HID1 = "/dev/hidg1"
-	HID2 = "/dev/hidg2"
+	HID0 = "/dev/hidg0" // Keyboard
+	HID1 = "/dev/hidg1" // Mouse (Relative Mode)
+	HID2 = "/dev/hidg2" // Touchpad (Absolute Mode)
 )
 
 var (
@@ -95,16 +95,22 @@ func (h *Hid) Close() {
 }
 
 func (h *Hid) WriteHid0(data []byte) {
+	deadline := time.Now().Add(8 * time.Millisecond)
+
 	h.kbMutex.Lock()
+	_ = h.g0.SetWriteDeadline(deadline)
 	_, err := h.g0.Write(data)
 	h.kbMutex.Unlock()
 
 	if err != nil {
-		if errors.Is(err, os.ErrClosed) {
+		switch {
+		case errors.Is(err, os.ErrClosed):
 			log.Errorf("hid already closed, reopen it...")
 			h.OpenNoLock()
-		} else {
-			log.Debugf("write to %s failed: %s", HID0, err)
+		case errors.Is(err, os.ErrDeadlineExceeded):
+			log.Debugf("write to %s timeout", HID0)
+		default:
+			log.Errorf("write to %s failed: %s", HID0, err)
 		}
 		return
 	}

server/service/hid/keyboard.go

@@ -1,15 +1,16 @@
 package hid
 
-func (h *Hid) Keyboard(queue <-chan []int) {
-	for event := range queue {
-		code := byte(event[0])
+import (
+	log "github.com/sirupsen/logrus"
+)
 
-		var modifier byte = 0x00
-		if code > 0 {
-			modifier = byte(event[1]) | byte(event[2]) | byte(event[3]) | byte(event[4])
+func (h *Hid) Keyboard(queue <-chan []byte) {
+	for event := range queue {
+		if len(event) != 8 {
+			log.Debugf("invalid keyboard event: %v", event)
+			continue
 		}
 
-		data := []byte{modifier, 0x00, code, 0x00, 0x00, 0x00, 0x00, 0x00}
-		h.WriteHid0(data)
+		h.WriteHid0(event)
 	}
 }

server/service/hid/mouse.go

@@ -1,82 +1,18 @@
 package hid
 
 import (
-	"encoding/binary"
-
 	log "github.com/sirupsen/logrus"
 )
 
-const (
-	MouseUp = iota
-	MouseDown
-	MouseMoveAbsolute
-	MouseMoveRelative
-	MouseScroll
-)
-
-var mouseButtonMap = map[byte]bool{
-	0x01: true,
-	0x02: true,
-	0x04: true,
-}
-
-func (h *Hid) Mouse(queue <-chan []int) {
+func (h *Hid) Mouse(queue <-chan []byte) {
 	for event := range queue {
-
-		switch event[0] {
-		case MouseDown:
-			h.mouseDown(event)
-		case MouseUp:
-			h.mouseUp()
-		case MouseMoveAbsolute:
-			h.mouseMoveAbsolute(event)
-		case MouseMoveRelative:
-			h.mouseMoveRelative(event)
-		case MouseScroll:
-			h.mouseScroll(event)
+		switch len(event) {
+		case 4:
+			h.WriteHid1(event)
+		case 6:
+			h.WriteHid2(event)
 		default:
 			log.Debugf("invalid mouse event: %v", event)
 		}
 	}
 }
-
-func (h *Hid) mouseDown(event []int) {
-	button := byte(event[1])
-
-	if _, ok := mouseButtonMap[button]; !ok {
-		log.Errorf("invalid mouse button: %v", event)
-		return
-	}
-
-	data := []byte{button, 0, 0, 0}
-	h.WriteHid1(data)
-}
-
-func (h *Hid) mouseUp() {
-	data := []byte{0, 0, 0, 0}
-	h.WriteHid1(data)
-}
-
-func (h *Hid) mouseScroll(event []int) {
-	direction := 0x01
-	if event[3] < 0 {
-		direction = -0x1
-	}
-
-	data := []byte{0, 0, 0, byte(direction)}
-	h.WriteHid1(data)
-}
-func (h *Hid) mouseMoveAbsolute(event []int) {
-	x := make([]byte, 2)
-	y := make([]byte, 2)
-	binary.LittleEndian.PutUint16(x, uint16(event[2]))
-	binary.LittleEndian.PutUint16(y, uint16(event[3]))
-
-	data := []byte{0, x[0], x[1], y[0], y[1], 0}
-	h.WriteHid2(data)
-}
-
-func (h *Hid) mouseMoveRelative(event []int) {
-	data := []byte{byte(event[1]), byte(event[2]), byte(event[3]), 0}
-	h.WriteHid1(data)
-}

server/service/ws/client.go

@@ -0,0 +1,109 @@
+package ws
+
+import (
+	"encoding/json"
+	"time"
+
+	"NanoKVM-Server/service/hid"
+	"NanoKVM-Server/service/vm/jiggler"
+
+	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	Heartbeat = iota
+	KeyboardEvent
+	MouseEvent
+)
+
+func NewClient(ws *websocket.Conn) *Client {
+	client := &Client{
+		ws:            ws,
+		hid:           hid.GetHid(),
+		keyboard:      make(chan []byte, 200),
+		mouse:         make(chan []byte, 200),
+		lastHeartbeat: time.Time{},
+	}
+
+	client.hid.Open()
+
+	return client
+}
+
+func (c *Client) Start() {
+	defer c.Close()
+
+	go c.hid.Keyboard(c.keyboard)
+	go c.hid.Mouse(c.mouse)
+
+	_ = c.Read()
+}
+
+func (c *Client) Read() error {
+	var zeroTime time.Time
+	_ = c.ws.SetReadDeadline(zeroTime)
+
+	for {
+		messageType, data, err := c.ws.ReadMessage()
+		if err != nil {
+			return err
+		}
+
+		log.Debugf("received message %d: %v", messageType, data)
+
+		switch data[0] {
+		case Heartbeat:
+			c.UpdateHeartbeat()
+		case KeyboardEvent:
+			writeQueue(c.keyboard, data[1:])
+		case MouseEvent:
+			writeQueue(c.mouse, data[1:])
+		}
+	}
+}
+
+func (c *Client) Write(event string, data string) error {
+	message := &Message{
+		Type: event,
+		Data: data,
+	}
+
+	messageByte, err := json.Marshal(message)
+	if err != nil {
+		log.Errorf("failed to marshal message: %s", err)
+		return err
+	}
+
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	_ = c.ws.SetWriteDeadline(time.Now().Add(10 * time.Second))
+	return c.ws.WriteMessage(websocket.TextMessage, messageByte)
+}
+
+func (c *Client) UpdateHeartbeat() {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+	c.lastHeartbeat = time.Now()
+}
+
+func (c *Client) Close() {
+	_ = c.ws.Close()
+
+	closeQueue(c.keyboard)
+	closeQueue(c.mouse)
+
+	log.Debug("websocket disconnected")
+}
+
+func writeQueue(queue chan []byte, data []byte) {
+	queue <- data
+	jiggler.GetJiggler().Update()
+}
+
+func closeQueue(queue chan []byte) {
+	for range queue {
+	}
+	close(queue)
+}

server/service/ws/manager.go

@@ -0,0 +1,46 @@
+package ws
+
+import (
+	"sync"
+
+	"github.com/gorilla/websocket"
+)
+
+var (
+	globalManager *Manager
+	managerOnce   sync.Once
+)
+
+func GetManager() *Manager {
+	managerOnce.Do(func() {
+		globalManager = &Manager{
+			clients: make(map[*websocket.Conn]*Client),
+			mutex:   sync.RWMutex{},
+		}
+	})
+	return globalManager
+}
+
+func (m *Manager) AddClient(ws *websocket.Conn, client *Client) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	m.clients[ws] = client
+}
+
+func (m *Manager) RemoveClient(ws *websocket.Conn) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	delete(m.clients, ws)
+}
+
+func (m *Manager) GetClients() []*Client {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	clients := make([]*Client, 0, len(m.clients))
+	for _, c := range m.clients {
+		clients = append(clients, c)
+	}
+
+	return clients
+}

server/service/ws/message.go

@@ -1,6 +0,0 @@
-package ws
-
-type Stream struct {
-	Type  string `json:"type"`
-	State int    `json:"state"`
-}

server/service/ws/service.go

@@ -1,7 +1,41 @@
 package ws
 
+import (
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
+)
+
 type Service struct{}
 
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+	CheckOrigin: func(r *http.Request) bool {
+		return true
+	},
+}
+
 func NewService() *Service {
 	return &Service{}
 }
+
+func (s *Service) Connect(c *gin.Context) {
+	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		log.Errorf("create websocket failed: %s", err)
+		return
+	}
+
+	log.Debug("websocket connected")
+
+	client := NewClient(ws)
+
+	manager := GetManager()
+	manager.AddClient(ws, client)
+	defer manager.RemoveClient(ws)
+
+	client.Start()
+}

server/service/ws/types.go

@@ -0,0 +1,29 @@
+package ws
+
+import (
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+
+	"NanoKVM-Server/service/hid"
+)
+
+type Manager struct {
+	clients map[*websocket.Conn]*Client
+	mutex   sync.RWMutex
+}
+
+type Client struct {
+	ws            *websocket.Conn
+	hid           *hid.Hid
+	keyboard      chan []byte
+	mouse         chan []byte
+	lastHeartbeat time.Time
+	mutex         sync.Mutex
+}
+
+type Message struct {
+	Type string `json:"type"`
+	Data string `json:"data"`
+}

server/service/ws/ws.go

@@ -1,117 +0,0 @@
-package ws
-
-import (
-	"encoding/json"
-	"net/http"
-	"time"
-
-	"github.com/gin-gonic/gin"
-	"github.com/gorilla/websocket"
-	log "github.com/sirupsen/logrus"
-
-	"NanoKVM-Server/service/hid"
-	"NanoKVM-Server/service/vm/jiggler"
-)
-
-const (
-	KeyboardEvent int = 1
-	MouseEvent    int = 2
-)
-
-type WsClient struct {
-	conn     *websocket.Conn
-	hid      *hid.Hid
-	keyboard chan []int
-	mouse    chan []int
-}
-
-var upgrader = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 1024,
-	CheckOrigin: func(r *http.Request) bool {
-		return true
-	},
-}
-
-func (s *Service) Connect(c *gin.Context) {
-	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
-	if err != nil {
-		log.Errorf("create websocket failed: %s", err)
-		return
-	}
-	log.Debug("websocket connected")
-
-	client := &WsClient{
-		hid:      hid.GetHid(),
-		conn:     conn,
-		keyboard: make(chan []int, 200),
-		mouse:    make(chan []int, 200),
-	}
-
-	go client.Start()
-}
-
-func (c *WsClient) Start() {
-	defer c.Clean()
-
-	c.hid.Open()
-
-	go c.hid.Keyboard(c.keyboard)
-	go c.hid.Mouse(c.mouse)
-
-	_ = c.Read()
-}
-
-func (c *WsClient) Read() error {
-	var zeroTime time.Time
-	_ = c.conn.SetReadDeadline(zeroTime)
-
-	for {
-		_, message, err := c.conn.ReadMessage()
-		if err != nil {
-			return err
-		}
-
-		log.Debugf("received message: %s", message)
-
-		var event []int
-		err = json.Unmarshal(message, &event)
-		if err != nil {
-			log.Debugf("received invalid message: %s", message)
-			continue
-		}
-
-		if event[0] == KeyboardEvent {
-			c.keyboard <- event[1:]
-		} else if event[0] == MouseEvent {
-			c.mouse <- event[1:]
-		}
-
-		// update latest HID operation time
-		jiggler.GetJiggler().Update()
-	}
-}
-
-func (c *WsClient) Write(message []byte) error {
-	_ = c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
-	return c.conn.WriteMessage(websocket.TextMessage, message)
-}
-
-func (c *WsClient) Clean() {
-	_ = c.conn.Close()
-
-	go clearQueue(c.keyboard)
-	close(c.keyboard)
-
-	go clearQueue(c.mouse)
-	close(c.mouse)
-
-	c.hid.Close()
-
-	log.Debug("websocket disconnected")
-}
-
-func clearQueue(queue chan []int) {
-	for range queue {
-	}
-}