From bd9d794a2c4386dc6c925221a2792cbc75bc0d54 Mon Sep 17 00:00:00 2001
From: Louis <lspgn@users.noreply.github.com>
Date: Sat, 20 Jan 2024 06:28:14 +0100
Subject: [PATCH] feat: add string renderer (#277)

---
 producer/proto/messages.go      |  4 +-
 producer/proto/messages_test.go | 67 +++++++++++++++++++++++++++++++++
 producer/proto/render.go        | 15 ++++++++
 3 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 producer/proto/messages_test.go

diff --git a/producer/proto/messages.go b/producer/proto/messages.go
index d1876835..5595b9a8 100644
--- a/producer/proto/messages.go
+++ b/producer/proto/messages.go
@@ -2,7 +2,6 @@ package protoproducer
 
 import (
 	"bytes"
-	"encoding/hex"
 	"fmt"
 	"hash"
 	"hash/fnv"
@@ -121,7 +120,8 @@ func (m *ProtoProducerMessage) mapUnknown() map[string]interface{} {
 				value = v
 			} else if dataType == protowire.BytesType {
 				v, _ := protowire.ConsumeString(data)
-				value = hex.EncodeToString([]byte(v))
+				//value = hex.EncodeToString([]byte(v)) // removed, this conversion is left to the renderer
+				value = []byte(v)
 			} else {
 				continue
 			}
diff --git a/producer/proto/messages_test.go b/producer/proto/messages_test.go
new file mode 100644
index 00000000..9310fbf9
--- /dev/null
+++ b/producer/proto/messages_test.go
@@ -0,0 +1,67 @@
+package protoproducer
+
+import (
+	"testing"
+
+	"google.golang.org/protobuf/encoding/protowire"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMarshalJSON(t *testing.T) {
+	var m ProtoProducerMessage
+
+	m.formatter = &FormatterConfigMapper{
+		fields: []string{"Etype", "test1", "test2", "test3"},
+		rename: map[string]string{
+			"Etype": "etype",
+		},
+		numToPb: map[int32]ProtobufFormatterConfig{
+			100: ProtobufFormatterConfig{
+				Name:  "test1",
+				Index: 100,
+				Type:  "varint",
+				Array: false,
+			},
+			101: ProtobufFormatterConfig{
+				Name:  "test2",
+				Index: 101,
+				Type:  "string",
+				Array: false,
+			},
+			102: ProtobufFormatterConfig{
+				Name:  "test3",
+				Index: 102,
+				Type:  "bytes",
+				Array: false,
+			},
+		},
+		render: map[string]RenderFunc{
+			"Etype": EtypeRenderer,
+			"test1": EtypeRenderer,
+			"test2": NilRenderer,
+			"test3": StringRenderer,
+		},
+	}
+
+	m.FlowMessage.Etype = 0x86dd
+
+	fmr := m.FlowMessage.ProtoReflect()
+	unk := fmr.GetUnknown()
+
+	unk = protowire.AppendTag(unk, protowire.Number(100), protowire.VarintType)
+	unk = protowire.AppendVarint(unk, 0x86dd)
+
+	unk = protowire.AppendTag(unk, protowire.Number(101), protowire.BytesType)
+	unk = protowire.AppendString(unk, string("testing"))
+
+	unk = protowire.AppendTag(unk, protowire.Number(102), protowire.BytesType)
+	unk = protowire.AppendString(unk, string([]byte{0x74, 0x65, 0x73, 0x74, 0x69, 0x6e, 0x67}))
+
+	fmr.SetUnknown(unk)
+
+	out, err := m.MarshalJSON()
+	assert.Nil(t, err)
+	t.Log(string(out))
+	assert.Equal(t, "{\"etype\":\"IPv6\",\"test1\":\"IPv6\",\"test2\":\"74657374696e67\",\"test3\":\"testing\"}", string(out))
+}
diff --git a/producer/proto/render.go b/producer/proto/render.go
index 07e99895..2e9070f2 100644
--- a/producer/proto/render.go
+++ b/producer/proto/render.go
@@ -22,6 +22,7 @@ const (
 	RendererNetwork      RendererID = "network"
 	RendererDateTime     RendererID = "datetime"
 	RendererDateTimeNano RendererID = "datetimenano"
+	RendererString       RendererID = "string"
 )
 
 var (
@@ -33,6 +34,7 @@ var (
 		RendererProto:        ProtoRenderer,
 		RendererDateTime:     DateTimeRenderer,
 		RendererDateTimeNano: DateTimeNanoRenderer,
+		RendererString:       StringRenderer,
 	}
 
 	defaultRenderers = map[string]RenderFunc{
@@ -95,6 +97,15 @@ func NilRenderer(msg *ProtoProducerMessage, fieldName string, data interface{})
 	return data
 }
 
+func StringRenderer(msg *ProtoProducerMessage, fieldName string, data interface{}) interface{} {
+	if dataC, ok := data.([]byte); ok {
+		return string(dataC)
+	} else if dataC, ok := data.(string); ok {
+		return string(dataC)
+	} // maybe should support uint64?
+	return NilRenderer(msg, fieldName, data)
+}
+
 func DateTimeRenderer(msg *ProtoProducerMessage, fieldName string, data interface{}) interface{} {
 	if dataC, ok := data.(uint64); ok {
 		ts := time.Unix(int64(dataC), 0).UTC()
@@ -151,6 +162,8 @@ func IPRenderer(msg *ProtoProducerMessage, fieldName string, data interface{}) i
 func EtypeRenderer(msg *ProtoProducerMessage, fieldName string, data interface{}) interface{} {
 	if dataC, ok := data.(uint32); ok {
 		return etypeName[dataC]
+	} else if dataC, ok := data.(uint64); ok { // supports protobuf mapped fields
+		return etypeName[uint32(dataC)]
 	}
 	return "unknown"
 }
@@ -158,6 +171,8 @@ func EtypeRenderer(msg *ProtoProducerMessage, fieldName string, data interface{}
 func ProtoRenderer(msg *ProtoProducerMessage, fieldName string, data interface{}) interface{} {
 	if dataC, ok := data.(uint32); ok {
 		return protoName[dataC]
+	} else if dataC, ok := data.(uint64); ok {
+		return protoName[uint32(dataC)]
 	}
 	return "unknown"
 }