From b01b1aacd2e450f9db47c8086c22253e5cfc9bc1 Mon Sep 17 00:00:00 2001
From: Louis <lspgn@users.noreply.github.com>
Date: Sun, 18 Aug 2024 20:06:30 -0700
Subject: [PATCH] fix(producer): mpls labels in NetFlow (#345)

Co-authored-by: Muhammad Iqbal Alaydrus <iqbal@bangunindo.com>
---
 producer/proto/producer_nf.go   |  8 +++++--
 producer/proto/producer_test.go | 40 ++++++++++++++++++++++++++++++---
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/producer/proto/producer_nf.go b/producer/proto/producer_nf.go
index 53cd9e96..bad74daf 100644
--- a/producer/proto/producer_nf.go
+++ b/producer/proto/producer_nf.go
@@ -518,7 +518,9 @@ func ConvertNetFlowDataSet(flowMessage *ProtoProducerMessage, version uint16, ba
 				return err
 			}
 			if len(flowMessage.MplsLabel) < 2 {
-				flowMessage.MplsLabel = make([]uint32, 2)
+				tmpLabels := make([]uint32, 2)
+				copy(tmpLabels, flowMessage.MplsLabel)
+				flowMessage.MplsLabel = tmpLabels
 			}
 			flowMessage.MplsLabel[1] = uint32(mplsLabel >> 4)
 		case netflow.IPFIX_FIELD_mplsLabelStackSection3:
@@ -527,7 +529,9 @@ func ConvertNetFlowDataSet(flowMessage *ProtoProducerMessage, version uint16, ba
 				return err
 			}
 			if len(flowMessage.MplsLabel) < 3 {
-				flowMessage.MplsLabel = make([]uint32, 3)
+				tmpLabels := make([]uint32, 3)
+				copy(tmpLabels, flowMessage.MplsLabel)
+				flowMessage.MplsLabel = tmpLabels
 			}
 			flowMessage.MplsLabel[2] = uint32(mplsLabel >> 4)
 		case netflow.IPFIX_FIELD_mplsTopLabelIPv4Address:
diff --git a/producer/proto/producer_test.go b/producer/proto/producer_test.go
index 0d009f3f..bbd14d5a 100644
--- a/producer/proto/producer_test.go
+++ b/producer/proto/producer_test.go
@@ -16,6 +16,31 @@ func TestProcessMessageNetFlow(t *testing.T) {
 					Type:  netflow.NFV9_FIELD_IPV4_SRC_ADDR,
 					Value: []byte{10, 0, 0, 1},
 				},
+				netflow.DataField{
+					Type: netflow.NFV9_FIELD_FIRST_SWITCHED,
+					// 218432176
+					Value: []byte{0x0d, 0x05, 0x02, 0xb0},
+				},
+				netflow.DataField{
+					Type: netflow.NFV9_FIELD_LAST_SWITCHED,
+					// 218432192
+					Value: []byte{0x0d, 0x05, 0x02, 0xc0},
+				},
+				netflow.DataField{
+					Type: netflow.NFV9_FIELD_MPLS_LABEL_1,
+					// 24041
+					Value: []byte{0x05, 0xde, 0x94},
+				},
+				netflow.DataField{
+					Type: netflow.NFV9_FIELD_MPLS_LABEL_2,
+					// 211992
+					Value: []byte{0x33, 0xc1, 0x85},
+				},
+				netflow.DataField{
+					Type: netflow.NFV9_FIELD_MPLS_LABEL_3,
+					// 48675
+					Value: []byte{0x0b, 0xe2, 0x35},
+				},
 			},
 		},
 	}
@@ -26,11 +51,20 @@ func TestProcessMessageNetFlow(t *testing.T) {
 	}
 
 	pktnf9 := netflow.NFv9Packet{
-		FlowSets: dfs,
+		SystemUptime: 218432000,
+		UnixSeconds:  1705732882,
+		FlowSets:     dfs,
 	}
 	testsr := &SingleSamplingRateSystem{1}
-	_, err := ProcessMessageNetFlowV9Config(&pktnf9, testsr, nil)
-	assert.Nil(t, err)
+	msgs, err := ProcessMessageNetFlowV9Config(&pktnf9, testsr, nil)
+	if assert.Nil(t, err) && assert.Len(t, msgs, 1) {
+		msg, ok := msgs[0].(*ProtoProducerMessage)
+		if assert.True(t, ok) {
+			assert.Equal(t, uint64(1705732882176*1e6), msg.TimeFlowStartNs)
+			assert.Equal(t, uint64(1705732882192*1e6), msg.TimeFlowEndNs)
+			assert.Equal(t, []uint32{24041, 211992, 48675}, msg.MplsLabel)
+		}
+	}
 
 	pktipfix := netflow.IPFIXPacket{
 		FlowSets: dfs,