From dfea7fca427e7e0cc9df4ae6da0408de0544b610 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Mon, 10 Apr 2023 23:12:11 +0200
Subject: [PATCH] refactor handler.

---
 internal/webrtc/handler.go         | 18 ++++++++----------
 internal/webrtc/payload/receive.go |  8 --------
 internal/webrtc/payload/send.go    |  4 ----
 internal/webrtc/peer.go            | 27 +++++++++++++++++++--------
 4 files changed, 27 insertions(+), 30 deletions(-)

diff --git a/internal/webrtc/handler.go b/internal/webrtc/handler.go
index 28a64a1e..bcc9a501 100644
--- a/internal/webrtc/handler.go
+++ b/internal/webrtc/handler.go
@@ -21,13 +21,8 @@ func (manager *WebRTCManagerCtx) handle(data []byte, dataChannel *webrtc.DataCha
 
 	buffer := bytes.NewBuffer(data)
 
-	hbytes := make([]byte, 3)
-	if _, err := buffer.Read(hbytes); err != nil {
-		return err
-	}
-
 	header := &payload.Header{}
-	if err := binary.Read(bytes.NewBuffer(hbytes), binary.BigEndian, header); err != nil {
+	if err := binary.Read(buffer, binary.BigEndian, header); err != nil {
 		return err
 	}
 
@@ -35,8 +30,6 @@ func (manager *WebRTCManagerCtx) handle(data []byte, dataChannel *webrtc.DataCha
 	// parse body
 	//
 
-	buffer = bytes.NewBuffer(data)
-
 	// handle cursor move event
 	if header.Event == payload.OP_MOVE {
 		payload := &payload.Move{}
@@ -64,8 +57,8 @@ func (manager *WebRTCManagerCtx) handle(data []byte, dataChannel *webrtc.DataCha
 			return err
 		}
 
-		// change header event to pong
-		ping.Header = payload.Header{
+		// create pong header
+		header := payload.Header{
 			Event:  payload.OP_PONG,
 			Length: 19,
 		}
@@ -81,6 +74,11 @@ func (manager *WebRTCManagerCtx) handle(data []byte, dataChannel *webrtc.DataCha
 		}
 
 		buffer := &bytes.Buffer{}
+
+		if err := binary.Write(buffer, binary.BigEndian, header); err != nil {
+			return err
+		}
+
 		if err := binary.Write(buffer, binary.BigEndian, pong); err != nil {
 			return err
 		}
diff --git a/internal/webrtc/payload/receive.go b/internal/webrtc/payload/receive.go
index 856c067c..68ef8f77 100644
--- a/internal/webrtc/payload/receive.go
+++ b/internal/webrtc/payload/receive.go
@@ -13,28 +13,20 @@ const (
 )
 
 type Move struct {
-	Header
-
 	X uint16
 	Y uint16
 }
 
 type Scroll struct {
-	Header
-
 	X int16
 	Y int16
 }
 
 type Key struct {
-	Header
-
 	Key uint32
 }
 
 type Ping struct {
-	Header
-
 	// client's timestamp split into two uint32
 	ClientTs1 uint32
 	ClientTs2 uint32
diff --git a/internal/webrtc/payload/send.go b/internal/webrtc/payload/send.go
index b539084e..949e27a3 100644
--- a/internal/webrtc/payload/send.go
+++ b/internal/webrtc/payload/send.go
@@ -9,15 +9,11 @@ const (
 )
 
 type CursorPosition struct {
-	Header
-
 	X uint16
 	Y uint16
 }
 
 type CursorImage struct {
-	Header
-
 	Width  uint16
 	Height uint16
 	Xhot   uint16
diff --git a/internal/webrtc/peer.go b/internal/webrtc/peer.go
index f250aee3..ca9b01ad 100644
--- a/internal/webrtc/peer.go
+++ b/internal/webrtc/peer.go
@@ -202,16 +202,22 @@ func (peer *WebRTCPeerCtx) SendCursorPosition(x, y int) error {
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
+	header := payload.Header{
+		Event:  payload.OP_CURSOR_POSITION,
+		Length: 7,
+	}
+
 	data := payload.CursorPosition{
-		Header: payload.Header{
-			Event:  payload.OP_CURSOR_POSITION,
-			Length: 7,
-		},
 		X: uint16(x),
 		Y: uint16(y),
 	}
 
 	buffer := &bytes.Buffer{}
+
+	if err := binary.Write(buffer, binary.BigEndian, header); err != nil {
+		return err
+	}
+
 	if err := binary.Write(buffer, binary.BigEndian, data); err != nil {
 		return err
 	}
@@ -223,11 +229,12 @@ func (peer *WebRTCPeerCtx) SendCursorImage(cur *types.CursorImage, img []byte) e
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
+	header := payload.Header{
+		Event:  payload.OP_CURSOR_IMAGE,
+		Length: uint16(11 + len(img)),
+	}
+
 	data := payload.CursorImage{
-		Header: payload.Header{
-			Event:  payload.OP_CURSOR_IMAGE,
-			Length: uint16(11 + len(img)),
-		},
 		Width:  cur.Width,
 		Height: cur.Height,
 		Xhot:   cur.Xhot,
@@ -236,6 +243,10 @@ func (peer *WebRTCPeerCtx) SendCursorImage(cur *types.CursorImage, img []byte) e
 
 	buffer := &bytes.Buffer{}
 
+	if err := binary.Write(buffer, binary.BigEndian, header); err != nil {
+		return err
+	}
+
 	if err := binary.Write(buffer, binary.BigEndian, data); err != nil {
 		return err
 	}