From 716b8e4efbe65b6fe6dcf1c4de9bb7b756ae1e83 Mon Sep 17 00:00:00 2001
From: Ivan Babrou <github@ivan.computer>
Date: Mon, 5 Feb 2024 17:06:05 -0800
Subject: [PATCH] Protect decoder cache with a mutex

---
 decoder/decoder.go      | 11 ++++----
 decoder/decoder_test.go | 56 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/decoder/decoder.go b/decoder/decoder.go
index 15837738..029c9e9c 100644
--- a/decoder/decoder.go
+++ b/decoder/decoder.go
@@ -61,8 +61,8 @@ func NewSet() (*Set, error) {
 	}, nil
 }
 
-// Decode transforms input byte field into a string according to configuration
-func (s *Set) Decode(in []byte, label config.Label) ([]byte, error) {
+// decode transforms input byte field into a string according to configuration
+func (s *Set) decode(in []byte, label config.Label) ([]byte, error) {
 	result := in
 
 	for _, decoder := range label.Decoders {
@@ -70,9 +70,7 @@ func (s *Set) Decode(in []byte, label config.Label) ([]byte, error) {
 			return result, fmt.Errorf("unknown decoder %q", decoder.Name)
 		}
 
-		s.mu.Lock()
 		decoded, err := s.decoders[decoder.Name].Decode(result, decoder)
-		s.mu.Unlock()
 		if err != nil {
 			if err == ErrSkipLabelSet {
 				return decoded, err
@@ -89,6 +87,9 @@ func (s *Set) Decode(in []byte, label config.Label) ([]byte, error) {
 // DecodeLabels transforms eBPF map key bytes into a list of label values
 // according to configuration
 func (s *Set) DecodeLabels(in []byte, labels []config.Label) ([]string, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	// string(in) must not be a variable to avoid allocation:
 	// * https://github.com/golang/go/commit/f5f5a8b6209f8
 	if cached, ok := s.cache[string(in)]; ok {
@@ -132,7 +133,7 @@ func (s *Set) decodeLabels(in []byte, labels []config.Label) ([]string, error) {
 
 		size := label.Size
 
-		decoded, err := s.Decode(in[off:off+size], label)
+		decoded, err := s.decode(in[off:off+size], label)
 		if err != nil {
 			return nil, err
 		}
diff --git a/decoder/decoder_test.go b/decoder/decoder_test.go
index 5871f628..d9debaa6 100644
--- a/decoder/decoder_test.go
+++ b/decoder/decoder_test.go
@@ -2,6 +2,7 @@ package decoder
 
 import (
 	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/cloudflare/ebpf_exporter/v2/config"
@@ -151,6 +152,61 @@ func TestDecodeLabels(t *testing.T) {
 	}
 }
 
+func TestConcurrency(t *testing.T) {
+	in := append([]byte{0x8, 0x0, 0x0, 0x0}, zeroPaddedString("bananas", 32)...)
+
+	labels := []config.Label{
+		{
+			Name: "number",
+			Size: 4,
+			Decoders: []config.Decoder{
+				{
+					Name: "uint",
+				},
+			},
+		},
+		{
+			Name: "fruit",
+			Size: 32,
+			Decoders: []config.Decoder{
+				{
+					Name: "string",
+				},
+				{
+					Name: "regexp",
+					Regexps: []string{
+						"^bananas$",
+						"$is-banana-even-fruit$",
+					},
+				},
+			},
+		},
+	}
+
+	s, err := NewSet()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	count := 1000
+
+	wg := sync.WaitGroup{}
+	wg.Add(count)
+
+	for i := 0; i < count; i++ {
+		go func() {
+			defer wg.Done()
+
+			_, err := s.DecodeLabels(in, labels)
+			if err != nil {
+				t.Error(err)
+			}
+		}()
+	}
+
+	wg.Wait()
+}
+
 func BenchmarkCache(b *testing.B) {
 	in := []byte{
 		0x8, 0xab, 0xce, 0xef,