Repository: open-telemetry/opamp-go (go)

Issue:
[Feature] Callback on message read error
I want to store the error in the database when there is an error reading the message.


Pull request description:
feat: Supports callbacks when reading a message fails
Resolves #330

This adds a new Server callback OnReadMessageError that is called when a read or deserialization error occurs.

Patch:
diff --git a/client/internal/packagessyncer.go b/client/internal/packagessyncer.go
index c02828d..fbc23d5 100644
--- a/client/internal/packagessyncer.go
+++ b/client/internal/packagessyncer.go
@@ -280,6 +280,12 @@ func (s *packagesSyncer) downloadFile(ctx context.Context, pkgName string, file
 		return fmt.Errorf("cannot download file from %s: %v", file.DownloadUrl, err)
 	}
 
+	if file.Headers != nil {
+		for _, h := range file.Headers.Headers {
+			req.Header.Add(h.GetKey(), h.GetValue())
+		}
+	}
+
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return fmt.Errorf("cannot download file from %s: %v", file.DownloadUrl, err)
diff --git a/server/serverimpl.go b/server/serverimpl.go
index dcdd044..ff7ff4a 100644
--- a/server/serverimpl.go
+++ b/server/serverimpl.go
@@ -5,6 +5,7 @@ import (
 	"compress/gzip"
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"net"
 	"net/http"
@@ -226,27 +227,39 @@ func (s *server) handleWSConnection(reqCtx context.Context, wsConn *websocket.Co
 	// Loop until fail to read from the WebSocket connection.
 	for {
 		msgContext := context.Background()
+		request := protobufs.AgentToServer{}
+
 		// Block until the next message can be read.
 		mt, msgBytes, err := wsConn.ReadMessage()
-		if err != nil {
-			if !websocket.IsUnexpectedCloseError(err) {
-				s.logger.Errorf(msgContext, "Cannot read a message from WebSocket: %v", err)
-				break
+		isBreak, err := func() (bool, error) {
+			if err != nil {
+				if !websocket.IsUnexpectedCloseError(err) {
+					s.logger.Errorf(msgContext, "Cannot read a message from WebSocket: %v", err)
+					return true, err
+				}
+				// This is a normal closing of the WebSocket connection.
+				s.logger.Debugf(msgContext, "Agent disconnected: %v", err)
+				return true, err
+			}
+			if mt != websocket.BinaryMessage {
+				err = fmt.Errorf("unexpected message type: %v, must be binary message", mt)
+				s.logger.Errorf(msgContext, "Cannot process a message from WebSocket: %v", err)
+				return false, err
 			}
-			// This is a normal closing of the WebSocket connection.
-			s.logger.Debugf(msgContext, "Agent disconnected: %v", err)
-			break
-		}
-		if mt != websocket.BinaryMessage {
-			s.logger.Errorf(msgContext, "Received unexpected message type from WebSocket: %v", mt)
-			continue
-		}
 
-		// Decode WebSocket message as a Protobuf message.
-		var request protobufs.AgentToServer
-		err = internal.DecodeWSMessage(msgBytes, &request)
+			// Decode WebSocket message as a Protobuf message.
+			err = internal.DecodeWSMessage(msgBytes, &request)
+			if err != nil {
+				s.logger.Errorf(msgContext, "Cannot decode message from WebSocket: %v", err)
+				return false, err
+			}
+			return false, nil
+		}()
 		if err != nil {
-			s.logger.Errorf(msgContext, "Cannot decode message from WebSocket: %v", err)
+			connectionCallbacks.OnReadMessageError(agentConn, mt, msgBytes, err)
+			if isBreak {
+				break
+			}
 			continue
 		}
 
@@ -366,7 +379,6 @@ func (s *server) handlePlainHTTPRequest(req *http.Request, w http.ResponseWriter
 		w.Header().Set(headerContentEncoding, contentEncodingGzip)
 	}
 	_, err = w.Write(bodyBytes)
-
 	if err != nil {
 		s.logger.Debugf(req.Context(), "Cannot send HTTP response: %v", err)
 	}
diff --git a/server/types/callbacks.go b/server/types/callbacks.go
index b25e818..9df3f21 100644
--- a/server/types/callbacks.go
+++ b/server/types/callbacks.go
@@ -62,6 +62,9 @@ type ConnectionCallbacks struct {
 
 	// OnConnectionClose is called when the OpAMP connection is closed.
 	OnConnectionClose func(conn Connection)
+
+	// OnReadMessageError is called when an error occurs while reading or deserializing a message.
+	OnReadMessageError func(conn Connection, mt int, msgByte []byte, err error)
 }
 
 func defaultOnConnected(ctx context.Context, conn Connection) {}
@@ -77,6 +80,8 @@ func defaultOnMessage(
 
 func defaultOnConnectionClose(conn Connection) {}
 
+func defaultOnReadMessageError(conn Connection, mt int, msgByte []byte, err error) {}
+
 func (c *ConnectionCallbacks) SetDefaults() {
 	if c.OnConnected == nil {
 		c.OnConnected = defaultOnConnected
@@ -89,4 +94,8 @@ func (c *ConnectionCallbacks) SetDefaults() {
 	if c.OnConnectionClose == nil {
 		c.OnConnectionClose = defaultOnConnectionClose
 	}
+
+	if c.OnReadMessageError == nil {
+		c.OnReadMessageError = defaultOnReadMessageError
+	}
 }


Test patch:
diff --git a/client/clientimpl_test.go b/client/clientimpl_test.go
index 87d3e5e..cf59a90 100644
--- a/client/clientimpl_test.go
+++ b/client/clientimpl_test.go
@@ -1439,10 +1439,16 @@ const packageFileURL = "/validfile.pkg"
 
 var packageFileContent = []byte("Package File Content")
 
+var optionalAuthHeaders = protobufs.Header{Key: "Authorization", Value: "Basic YWxhZGRpbjpvcGVuc2VzYW1l"}
+
 func createDownloadSrv(t *testing.T) *httptest.Server {
 	m := http.NewServeMux()
 	m.HandleFunc(packageFileURL,
 		func(w http.ResponseWriter, r *http.Request) {
+			if r.Header.Get(optionalAuthHeaders.GetKey()) != "" {
+				assert.Equal(t, r.Header.Get(optionalAuthHeaders.GetKey()), optionalAuthHeaders.GetValue())
+			}
+
 			w.WriteHeader(http.StatusOK)
 			_, err := w.Write(packageFileContent)
 			assert.NoError(t, err)
@@ -1527,6 +1533,11 @@ func TestUpdatePackages(t *testing.T) {
 	errorOnCallback.errorOnCallback = true
 	tests = append(tests, errorOnCallback)
 
+	// A case where we send optional headers
+	withHeaders := createPackageTestCase("with optional HTTP headers", downloadSrv)
+	withHeaders.available.Packages["package1"].File.Headers = &protobufs.Headers{Headers: []*protobufs.Header{&optionalAuthHeaders}}
+	tests = append(tests, withHeaders)
+
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			verifyUpdatePackages(t, test)
diff --git a/server/serverimpl_test.go b/server/serverimpl_test.go
index ca7a5c3..b2d7062 100644
--- a/server/serverimpl_test.go
+++ b/server/serverimpl_test.go
@@ -315,6 +315,52 @@ func TestServerReceiveSendMessage(t *testing.T) {
 	assert.EqualValues(t, settings.CustomCapabilities, response.CustomCapabilities.Capabilities)
 }
 
+func TestServerReceiveSendErrorMessage(t *testing.T) {
+	var rcvMsg atomic.Value
+	type ErrorInfo struct {
+		mt      int
+		msgByte []byte
+		err     error
+	}
+	callbacks := types.Callbacks{
+		OnConnecting: func(request *http.Request) types.ConnectionResponse {
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: types.ConnectionCallbacks{
+				OnReadMessageError: func(conn types.Connection, mt int, msgByte []byte, err error) {
+					rcvMsg.Store(ErrorInfo{
+						mt:      mt,
+						msgByte: msgByte,
+						err:     err,
+					})
+				},
+			}}
+		},
+	}
+
+	// Start a Server.
+	settings := &StartSettings{Settings: Settings{
+		Callbacks:          callbacks,
+		CustomCapabilities: []string{"local.test.capability"},
+	}}
+	srv := startServer(t, settings)
+	defer srv.Stop(context.Background())
+
+	// Connect using a WebSocket client.
+	conn, _, _ := dialClient(settings)
+	require.NotNil(t, conn)
+	defer conn.Close()
+
+	// Send an invalid message to the Server. This should result in calling OnReadMessageError().
+	err := conn.WriteMessage(websocket.TextMessage, []byte("abc"))
+	require.NoError(t, err)
+
+	// Wait until Server receives the message.
+	eventually(t, func() bool { return rcvMsg.Load() != nil })
+	errInfo := rcvMsg.Load().(ErrorInfo)
+	assert.EqualValues(t, websocket.TextMessage, errInfo.mt)
+	assert.EqualValues(t, []byte("abc"), errInfo.msgByte)
+	assert.NotNil(t, errInfo.err)
+}
+
 func TestServerReceiveSendMessageWithCompression(t *testing.T) {
 	// Use highly compressible config body.
 	uncompressedCfg := []byte(strings.Repeat("test", 10000))
@@ -620,7 +666,6 @@ func TestServerAttachSendMessagePlainHTTP(t *testing.T) {
 }
 
 func TestServerHonoursClientRequestContentEncoding(t *testing.T) {
-
 	hc := http.Client{}
 	var rcvMsg atomic.Value
 	var onConnectedCalled, onCloseCalled int32
@@ -698,7 +743,6 @@ func TestServerHonoursClientRequestContentEncoding(t *testing.T) {
 }
 
 func TestServerHonoursAcceptEncoding(t *testing.T) {
-
 	hc := http.Client{}
 	var rcvMsg atomic.Value
 	var onConnectedCalled, onCloseCalled int32
@@ -985,7 +1029,6 @@ func BenchmarkSendToClient(b *testing.B) {
 	}
 	srv := New(&sharedinternal.NopLogger{})
 	err := srv.Start(*settings)
-
 	if err != nil {
 		b.Error(err)
 	}
@@ -1017,7 +1060,6 @@ func BenchmarkSendToClient(b *testing.B) {
 
 	for _, conn := range serverConnections {
 		err := conn.Send(context.Background(), &protobufs.ServerToAgent{})
-
 		if err != nil {
 			b.Error(err)
 		}
@@ -1026,5 +1068,4 @@ func BenchmarkSendToClient(b *testing.B) {
 	for _, conn := range clientConnections {
 		conn.Close()
 	}
-
 }