From f0958a79e1a0dedebd96957043599a3fc9109592 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maciej=20Krzy=C5=BCanowski?= <krzmaciek@gmail.com>
Date: Sun, 29 Dec 2024 03:06:56 +0100
Subject: [PATCH] Add handler context and echo updates handling

---
 daemonconn/messages.go            | 18 ++++++++++++++++--
 main.go                           |  2 +-
 msghandlers/echorequesthandler.go | 23 +++++++++++++++++++++--
 msghandlers/msghandlers.go        | 13 +++++++++----
 4 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/daemonconn/messages.go b/daemonconn/messages.go
index de890aa..b0f9656 100644
--- a/daemonconn/messages.go
+++ b/daemonconn/messages.go
@@ -3,18 +3,32 @@ package daemonconn
 const (
 	MessageTypeEchoRequest = iota
 	MessageTypeEchoResponse
+	MessageTypeEchoUpdate
 )
 
 type EchoRequest struct {
-	EchoByte byte
+	EchoByte     byte
+	UpdatesCount byte
 }
 
 func (EchoRequest) TypeID() uint32 {
 	return MessageTypeEchoRequest
 }
 
-type EchoResponse EchoRequest
+type EchoResponse struct {
+	EchoByte byte
+}
 
 func (EchoResponse) TypeID() uint32 {
 	return MessageTypeEchoResponse
 }
+
+// EchoUpdate should be sent from the side, which received EchoRequest
+// UpdatesCount times (which is specified in the EchoRequest message)
+type EchoUpdate struct {
+	EchoByte byte
+}
+
+func (EchoUpdate) TypeID() uint32 {
+	return MessageTypeEchoUpdate
+}
diff --git a/main.go b/main.go
index e420385..3d533b1 100644
--- a/main.go
+++ b/main.go
@@ -56,7 +56,7 @@ func main() {
 		panic(err)
 	}
 
-	err = conn.WriteMessage(daemonconn.EchoRequest{EchoByte: 123})
+	err = conn.WriteMessage(daemonconn.EchoRequest{EchoByte: 123, UpdatesCount: 16})
 	if err != nil {
 		panic(err)
 	}
diff --git a/msghandlers/echorequesthandler.go b/msghandlers/echorequesthandler.go
index 579b92d..5962762 100644
--- a/msghandlers/echorequesthandler.go
+++ b/msghandlers/echorequesthandler.go
@@ -3,11 +3,30 @@ package msghandlers
 import (
 	"daemonSocketExample/daemonconn"
 	"log/slog"
+	"time"
 )
 
-func handleEchoRequest(msg daemonconn.Message) error {
+func handleEchoRequest(ctx HandlerContext, msg daemonconn.Message) error {
 	echoRequest := msg.(daemonconn.EchoRequest)
-	slog.Info("Got echo request", "echoByte", echoRequest.EchoByte)
+	slog.Info("Got echo request", "echoByte", echoRequest.EchoByte, "updatesCount", echoRequest.UpdatesCount)
+
+	for i := 0; i < int(echoRequest.UpdatesCount); i++ {
+		err := ctx.conn.WriteMessage(daemonconn.EchoUpdate{EchoByte: echoRequest.EchoByte})
+		if err != nil {
+			return err
+		}
+
+		slog.Info("Sent echo update", "i", i)
+		time.Sleep(1 * time.Second)
+	}
+
+	echoResponse := daemonconn.EchoResponse{EchoByte: echoRequest.EchoByte}
+	err := ctx.conn.WriteMessage(echoResponse)
+	if err != nil {
+		return err
+	}
+
+	slog.Info("Sent echo response", "echoByte", echoResponse.EchoByte)
 
 	return nil
 }
diff --git a/msghandlers/msghandlers.go b/msghandlers/msghandlers.go
index 618f5b6..73bfd9d 100644
--- a/msghandlers/msghandlers.go
+++ b/msghandlers/msghandlers.go
@@ -4,17 +4,22 @@ import (
 	"daemonSocketExample/daemonconn"
 )
 
-type MessageHandler func(daemonconn.Message) error
+type HandlerContext struct {
+	conn daemonconn.DaemonConn
+}
+
+type MessageHandler func(HandlerContext, daemonconn.Message) error
 
 var Handlers = map[uint32]MessageHandler{
 	daemonconn.MessageTypeEchoRequest: handleEchoRequest,
 }
 
 func HandleDaemonConn(conn daemonconn.DaemonConn) error {
+	ctx := HandlerContext{conn: conn}
 	msg, err := conn.ReadMessage()
 
 	for err == nil {
-		err := handleMessage(msg)
+		err := handleMessage(ctx, msg)
 		if err != nil {
 			return err
 		}
@@ -25,8 +30,8 @@ func HandleDaemonConn(conn daemonconn.DaemonConn) error {
 	return err
 }
 
-func handleMessage(msg daemonconn.Message) error {
+func handleMessage(ctx HandlerContext, msg daemonconn.Message) error {
 	handlerFunc := Handlers[msg.TypeID()]
 
-	return handlerFunc(msg)
+	return handlerFunc(ctx, msg)
 }