From 3d1c950dff0c096ac14ae8d50d265248949ff9ff Mon Sep 17 00:00:00 2001
From: Vladimir Rachkin <33569237+robozmey@users.noreply.github.com>
Date: Tue, 5 Nov 2024 20:25:45 +0300
Subject: [PATCH] Add error forwarding from to yp-client from yproxy (#75)

* Add error forwarding

* Add ErrorMessage type of message

* Add test for ErrorMessage

* Remove error check from ReadPacket()

* Update proto.go

---------

Co-authored-by: reshke <reshke@double.cloud>
---
 pkg/client/client.go         |  6 ++---
 pkg/message/error_message.go | 49 ++++++++++++++++++++++++++++++++++++
 pkg/message/message.go       |  3 +++
 pkg/message/message_test.go  | 20 +++++++++++++++
 pkg/proc/proto.go            |  3 +--
 5 files changed, 75 insertions(+), 6 deletions(-)
 create mode 100644 pkg/message/error_message.go

diff --git a/pkg/client/client.go b/pkg/client/client.go
index 9bb698d..418ba13 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -1,11 +1,11 @@
 package client
 
 import (
-	"fmt"
 	"io"
 	"net"
 	"reflect"
 
+	"github.com/yezzey-gp/yproxy/pkg/message"
 	"github.com/yezzey-gp/yproxy/pkg/ylogger"
 )
 
@@ -76,9 +76,7 @@ func NewYClient(c net.Conn) YproxyClient {
 func (y *YClient) ReplyError(err error, msg string) error {
 	ylogger.Zero.Error().Err(err).Msg(msg)
 
-	_, _ = y.Conn.Write([]byte(
-		fmt.Sprintf("%s: %v", msg, err),
-	))
+	_, _ = y.Conn.Write(message.NewErrorMessage(err, msg).Encode())
 	return nil
 }
 
diff --git a/pkg/message/error_message.go b/pkg/message/error_message.go
new file mode 100644
index 0000000..4cc6abc
--- /dev/null
+++ b/pkg/message/error_message.go
@@ -0,0 +1,49 @@
+package message
+
+import (
+	"encoding/binary"
+)
+
+type ErrorMessage struct {
+	Error   string
+	Message string
+}
+
+var _ ProtoMessage = &ListMessage{}
+
+func NewErrorMessage(err error, msg string) *ErrorMessage {
+	return &ErrorMessage{
+		Error:   err.Error(),
+		Message: msg,
+	}
+}
+
+func (c *ErrorMessage) Encode() []byte {
+	encodedMessage := []byte{
+		byte(MessageTypeError),
+		0,
+		0,
+		0,
+	}
+
+	byteError := []byte(c.Error)
+	byteLen := make([]byte, 8)
+	binary.BigEndian.PutUint64(byteLen, uint64(len(byteError)))
+	encodedMessage = append(encodedMessage, byteLen...)
+	encodedMessage = append(encodedMessage, byteError...)
+
+	byteMessage := []byte(c.Message)
+	binary.BigEndian.PutUint64(byteLen, uint64(len(byteMessage)))
+	encodedMessage = append(encodedMessage, byteLen...)
+	encodedMessage = append(encodedMessage, byteMessage...)
+
+	binary.BigEndian.PutUint64(byteLen, uint64(len(encodedMessage)+8))
+	return append(byteLen, encodedMessage...)
+}
+
+func (c *ErrorMessage) Decode(data []byte) {
+	errorLen := binary.BigEndian.Uint64(data[4:12])
+	c.Error = string(data[12 : 12+errorLen])
+	messageLen := binary.BigEndian.Uint64(data[12+errorLen : 12+errorLen+8])
+	c.Message = string(data[12+errorLen+8 : 12+errorLen+8+messageLen])
+}
diff --git a/pkg/message/message.go b/pkg/message/message.go
index 4392171..f2f2daa 100644
--- a/pkg/message/message.go
+++ b/pkg/message/message.go
@@ -23,6 +23,7 @@ const (
 	MessageTypeGool            = MessageType(52)
 	MessageTypePutV2           = MessageType(53)
 	MessageTypeCatV2           = MessageType(54)
+	MessageTypeError           = MessageType(55)
 
 	DecryptMessage   = RequestEncryption(1)
 	NoDecryptMessage = RequestEncryption(0)
@@ -59,6 +60,8 @@ func (m MessageType) String() string {
 		return "COPY"
 	case MessageTypeGool:
 		return "GOOL"
+	case MessageTypeError:
+		return "ERROR"
 	}
 	return "UNKNOWN"
 }
diff --git a/pkg/message/message_test.go b/pkg/message/message_test.go
index 2ea1c85..3aa6e49 100644
--- a/pkg/message/message_test.go
+++ b/pkg/message/message_test.go
@@ -2,6 +2,7 @@ package message_test
 
 import (
 	"encoding/binary"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -303,3 +304,22 @@ func TestDeleteMsg(t *testing.T) {
 	assert.True(msg2.Confirm)
 	assert.True(msg2.Garbage)
 }
+
+func TestErrorMsg(t *testing.T) {
+	assert := assert.New(t)
+
+	errString := "could not list objects: NoCredentialProviders: no valid providers in chain."
+	messageString := "Deprecated. For verbose messaging see aws.Config.CredentialsChainVerboseErrors"
+	err := fmt.Errorf("%s", errString)
+
+	msg := message.NewErrorMessage(err, messageString)
+	body := msg.Encode()
+
+	assert.Equal(body[8], byte(message.MessageTypeError))
+
+	msg2 := message.ErrorMessage{}
+	msg2.Decode(body[8:])
+
+	assert.Equal(errString, msg2.Error)
+	assert.Equal(messageString, msg2.Message)
+}
diff --git a/pkg/proc/proto.go b/pkg/proc/proto.go
index af534dc..156bf73 100644
--- a/pkg/proc/proto.go
+++ b/pkg/proc/proto.go
@@ -50,6 +50,5 @@ func (r *ProtoReader) ReadPacket() (message.MessageType, []byte, error) {
 		return 0, nil, err
 	}
 
-	msgType := message.MessageType(data[0])
-	return msgType, data, nil
+	return message.MessageType(data[0]), data, nil
 }