Restructure protobufs and use them everywhere

packages/relay/internal/common/safebufio.go

@@ -3,16 +3,28 @@ package common
 import (
 	"bufio"
 	"encoding/binary"
-	"encoding/json"
 	"errors"
 	"io"
+	gen "relay/internal/proto"
 	"sync"
 
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protoreflect"
+	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
-// MaxSize is the maximum allowed data size (1MB)
-const MaxSize = 1024 * 1024
+// readUvarint reads an unsigned varint from the reader
+func readUvarint(r io.ByteReader) (uint64, error) {
+	return binary.ReadUvarint(r)
+}
+
+// writeUvarint writes an unsigned varint to the writer
+func writeUvarint(w io.Writer, x uint64) error {
+	buf := make([]byte, binary.MaxVarintLen64)
+	n := binary.PutUvarint(buf, x)
+	_, err := w.Write(buf[:n])
+	return err
+}
 
 // SafeBufioRW wraps a bufio.ReadWriter for sending and receiving JSON and protobufs safely
 type SafeBufioRW struct {
@@ -24,83 +36,6 @@ func NewSafeBufioRW(brw *bufio.ReadWriter) *SafeBufioRW {
 	return &SafeBufioRW{brw: brw}
 }
 
-// SendJSON serializes the given data as JSON and sends it with a 4-byte length prefix
-func (bu *SafeBufioRW) SendJSON(data interface{}) error {
-	bu.mutex.Lock()
-	defer bu.mutex.Unlock()
-
-	jsonData, err := json.Marshal(data)
-	if err != nil {
-		return err
-	}
-
-	if len(jsonData) > MaxSize {
-		return errors.New("JSON data exceeds maximum size")
-	}
-
-	// Write the 4-byte length prefix
-	if err = binary.Write(bu.brw, binary.BigEndian, uint32(len(jsonData))); err != nil {
-		return err
-	}
-
-	// Write the JSON data
-	if _, err = bu.brw.Write(jsonData); err != nil {
-		return err
-	}
-
-	// Flush the writer to ensure data is sent
-	return bu.brw.Flush()
-}
-
-// ReceiveJSON reads a 4-byte length prefix, then reads and unmarshals the JSON
-func (bu *SafeBufioRW) ReceiveJSON(dest interface{}) error {
-	bu.mutex.RLock()
-	defer bu.mutex.RUnlock()
-
-	// Read the 4-byte length prefix
-	var length uint32
-	if err := binary.Read(bu.brw, binary.BigEndian, &length); err != nil {
-		return err
-	}
-
-	if length > MaxSize {
-		return errors.New("received JSON data exceeds maximum size")
-	}
-
-	// Read the JSON data
-	data := make([]byte, length)
-	if _, err := io.ReadFull(bu.brw, data); err != nil {
-		return err
-	}
-
-	return json.Unmarshal(data, dest)
-}
-
-// Receive reads a 4-byte length prefix, then reads the raw data
-func (bu *SafeBufioRW) Receive() ([]byte, error) {
-	bu.mutex.RLock()
-	defer bu.mutex.RUnlock()
-
-	// Read the 4-byte length prefix
-	var length uint32
-	if err := binary.Read(bu.brw, binary.BigEndian, &length); err != nil {
-		return nil, err
-	}
-
-	if length > MaxSize {
-		return nil, errors.New("received data exceeds maximum size")
-	}
-
-	// Read the raw data
-	data := make([]byte, length)
-	if _, err := io.ReadFull(bu.brw, data); err != nil {
-		return nil, err
-	}
-
-	return data, nil
-}
-
-// SendProto serializes the given protobuf message and sends it with a 4-byte length prefix
 func (bu *SafeBufioRW) SendProto(msg proto.Message) error {
 	bu.mutex.Lock()
 	defer bu.mutex.Unlock()
@@ -110,12 +45,8 @@ func (bu *SafeBufioRW) SendProto(msg proto.Message) error {
 		return err
 	}
 
-	if len(protoData) > MaxSize {
-		return errors.New("protobuf data exceeds maximum size")
-	}
-
-	// Write the 4-byte length prefix
-	if err = binary.Write(bu.brw, binary.BigEndian, uint32(len(protoData))); err != nil {
+	// Write varint length prefix
+	if err := writeUvarint(bu.brw, uint64(len(protoData))); err != nil {
 		return err
 	}
 
@@ -124,25 +55,19 @@ func (bu *SafeBufioRW) SendProto(msg proto.Message) error {
 		return err
 	}
 
-	// Flush the writer to ensure data is sent
 	return bu.brw.Flush()
 }
 
-// ReceiveProto reads a 4-byte length prefix, then reads and unmarshals the protobuf
 func (bu *SafeBufioRW) ReceiveProto(msg proto.Message) error {
 	bu.mutex.RLock()
 	defer bu.mutex.RUnlock()
 
-	// Read the 4-byte length prefix
-	var length uint32
-	if err := binary.Read(bu.brw, binary.BigEndian, &length); err != nil {
+	// Read varint length prefix
+	length, err := readUvarint(bu.brw)
+	if err != nil {
 		return err
 	}
 
-	if length > MaxSize {
-		return errors.New("received Protobuf data exceeds maximum size")
-	}
-
 	// Read the Protobuf data
 	data := make([]byte, length)
 	if _, err := io.ReadFull(bu.brw, data); err != nil {
@@ -152,24 +77,51 @@ func (bu *SafeBufioRW) ReceiveProto(msg proto.Message) error {
 	return proto.Unmarshal(data, msg)
 }
 
-// Write writes raw data to the underlying buffer
-func (bu *SafeBufioRW) Write(data []byte) (int, error) {
-	bu.mutex.Lock()
-	defer bu.mutex.Unlock()
-
-	if len(data) > MaxSize {
-		return 0, errors.New("data exceeds maximum size")
-	}
-
-	n, err := bu.brw.Write(data)
-	if err != nil {
-		return n, err
-	}
-
-	// Flush the writer to ensure data is sent
-	if err = bu.brw.Flush(); err != nil {
-		return n, err
-	}
-
-	return n, nil
+type CreateMessageOptions struct {
+	SequenceID string
+	Latency    *gen.ProtoLatencyTracker
+}
+
+func CreateMessage(payload proto.Message, payloadType string, opts *CreateMessageOptions) (*gen.ProtoMessage, error) {
+	msg := &gen.ProtoMessage{
+		MessageBase: &gen.ProtoMessageBase{
+			PayloadType: payloadType,
+		},
+	}
+
+	if opts != nil {
+		if opts.Latency != nil {
+			msg.MessageBase.Latency = opts.Latency
+		} else if opts.SequenceID != "" {
+			msg.MessageBase.Latency = &gen.ProtoLatencyTracker{
+				SequenceId: opts.SequenceID,
+				Timestamps: []*gen.ProtoTimestampEntry{
+					{
+						Stage: "created",
+						Time:  timestamppb.Now(),
+					},
+				},
+			}
+		}
+	}
+
+	// Use reflection to set the oneof field automatically
+	msgReflect := msg.ProtoReflect()
+	payloadReflect := payload.ProtoReflect()
+
+	oneofDesc := msgReflect.Descriptor().Oneofs().ByName("payload")
+	if oneofDesc == nil {
+		return nil, errors.New("payload oneof not found")
+	}
+
+	fields := oneofDesc.Fields()
+	for i := 0; i < fields.Len(); i++ {
+		field := fields.Get(i)
+		if field.Message() != nil && field.Message().FullName() == payloadReflect.Descriptor().FullName() {
+			msgReflect.Set(field, protoreflect.ValueOfMessage(payloadReflect))
+			return msg, nil
+		}
+	}
+
+	return nil, errors.New("payload type not found in oneof")
 }