From 800d4e9e60d007b7179055d74919a573c430d8f8 Mon Sep 17 00:00:00 2001
From: 1998-felix <felix.gateru@gmail.com>
Date: Fri, 26 Apr 2024 13:57:35 +0300
Subject: [PATCH] fix: Change context source for client.handler

Signed-off-by: 1998-felix <felix.gateru@gmail.com>
---
 coap/client.go               |  3 +-
 internal/server/coap/coap.go | 56 +++++++++---------------------------
 2 files changed, 16 insertions(+), 43 deletions(-)

diff --git a/coap/client.go b/coap/client.go
index 73ab01c87c8..ae3c0107a50 100644
--- a/coap/client.go
+++ b/coap/client.go
@@ -5,6 +5,7 @@ package coap
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"log/slog"
 	"sync/atomic"
@@ -71,7 +72,7 @@ func (c *client) Token() string {
 }
 
 func (c *client) Handle(msg *messaging.Message) error {
-	pm := pool.NewMessage(c.conn.Context())
+	pm := pool.NewMessage(context.Background())
 	pm.SetCode(codes.Content)
 	pm.SetToken(c.token)
 	pm.SetBody(bytes.NewReader(msg.GetPayload()))
diff --git a/internal/server/coap/coap.go b/internal/server/coap/coap.go
index 5b9ea11a98e..6bdc496e498 100644
--- a/internal/server/coap/coap.go
+++ b/internal/server/coap/coap.go
@@ -10,16 +10,13 @@ import (
 	"time"
 
 	"github.com/absmach/magistrala/internal/server"
+	gocoap "github.com/plgd-dev/go-coap/v3"
 	"github.com/plgd-dev/go-coap/v3/mux"
-	"github.com/plgd-dev/go-coap/v3/net"
-	"github.com/plgd-dev/go-coap/v3/options"
-	"github.com/plgd-dev/go-coap/v3/udp"
-	udpClient "github.com/plgd-dev/go-coap/v3/udp/client"
 )
 
-// const (
-// 	stopWaitTime = 5 * time.Second
-// )
+const (
+	stopWaitTime = 5 * time.Second
+)
 
 type Server struct {
 	server.BaseServer
@@ -47,51 +44,26 @@ func (s *Server) Start() error {
 	errCh := make(chan error)
 	s.Logger.Info(fmt.Sprintf("%s service started using http, exposed port %s", s.Name, s.Address))
 	s.Logger.Info(fmt.Sprintf("%s service %s server listening at %s without TLS", s.Name, s.Protocol, s.Address))
-	l, err := net.NewListenUDP("udp", s.Address)
-	if err != nil {
-		return err
-	}
-	defer l.Close()
-
-	s.Logger.Info(fmt.Sprintf("CoAP proxy server started at %s without DTLS", s.Address))
-
-	dummyInactiveFunc := func(cc *udpClient.Conn) {
-		// This function intentionally left blank.
-	}
-	cs := udp.NewServer(
-		options.WithMux(mux.HandlerFunc(s.handler)),
-		options.WithKeepAlive(10, 10*time.Minute, dummyInactiveFunc),
-	)
-
 	go func() {
-		errCh <- cs.Serve(l)
+		errCh <- gocoap.ListenAndServe("udp", s.Address, s.handler)
 	}()
 
 	select {
 	case <-s.Ctx.Done():
-		s.Logger.Info(fmt.Sprintf("CoAP proxy server at %s without DTLS exiting ...", s.Address))
-		l.Close()
+		return s.Stop()
 	case err := <-errCh:
-		s.Logger.Error(fmt.Sprintf("CoAP proxy server at %s without DTLS exiting with errors: %s", s.Address, err.Error()))
 		return err
 	}
-	return nil
-	// select {
-	// case <-s.Ctx.Done():
-	// 	return s.Stop()
-	// case err := <-errCh:
-	// 	return err
-	// }
 }
 
 func (s *Server) Stop() error {
-	// 	defer s.Cancel()
-	// 	c := make(chan bool)
-	// 	defer close(c)
-	// 	select {
-	// 	case <-c:
-	// 	case <-time.After(stopWaitTime):
-	// 	}
-	// 	s.Logger.Info(fmt.Sprintf("%s service shutdown of http at %s", s.Name, s.Address))
+	defer s.Cancel()
+	c := make(chan bool)
+	defer close(c)
+	select {
+	case <-c:
+	case <-time.After(stopWaitTime):
+	}
+	s.Logger.Info(fmt.Sprintf("%s service shutdown of http at %s", s.Name, s.Address))
 	return nil
 }