make memory safe

EwenQuim · Jun 20, 2024 · 54f834c · 54f834c
1 parent e8b2f87
commit 54f834c
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 27 deletions.
diff --git a/main.go b/main.go
@@ -36,22 +36,29 @@ type Entropy struct {
 
 // Entropies should be created with a size n using make()
 // it should not be written to manually, instead use Entropies.Add
-type Entropies []Entropy
+type Entropies struct {
+	sync.Mutex
+	Entropies []Entropy
+}
 
 // Add assumes that es contains an ordered set of entropies.
 // It preserves ordering, and inserts an additional value e, if it has high enough entropy.
 // In that case, the entry with lowest entropy is rejected.
-func (es Entropies) Add(e Entropy) {
-	if es[len(es)-1].Entropy >= e.Entropy {
+func (es *Entropies) Add(e Entropy) {
+	es.Lock()
+	defer es.Unlock()
+
+	entropies := es.Entropies
+	if entropies[len(entropies)-1].Entropy >= e.Entropy {
 		return
 	}
 
-	for i := range len(es) {
-		if es[i].Entropy < e.Entropy {
-			for j := len(es) - 1; j > i; j-- {
-				es[j] = es[j-1]
+	for i := range len(entropies) {
+		if entropies[i].Entropy < e.Entropy {
+			for j := len(entropies) - 1; j > i; j-- {
+				entropies[j] = entropies[j-1]
 			}
-			es[i] = e
+			entropies[i] = e
 			return
 		}
 	}
@@ -86,7 +93,9 @@ func main() {
 		fmt.Println("No files provided, defaults to current folder.")
 		fileNames = []string{"."}
 	}
-	entropies := make(Entropies, *resultCountFlag)
+	entropies := &Entropies{
+		Entropies: make([]Entropy, resultCount),
+	}
 	for _, fileName := range fileNames {
 		err := readFile(entropies, fileName)
 		if err != nil {
@@ -102,15 +111,15 @@ func main() {
 		resetMark = ""
 	}
 
-	for _, entropy := range entropies {
+	for _, entropy := range entropies.Entropies {
 		if entropy == (Entropy{}) {
 			return
 		}
 		fmt.Printf("%.2f: %s%s:%d%s %s\n", entropy.Entropy, redMark, entropy.File, entropy.LineNum, resetMark, entropy.Line)
 	}
 }
 
-func readFile(entropies Entropies, fileName string) error {
+func readFile(entropies *Entropies, fileName string) error {
 	// If file is a folder, walk inside the folder
 	fileInfo, err := os.Stat(fileName)
 	if err != nil {

diff --git a/main_test.go b/main_test.go
@@ -5,7 +5,7 @@ import (
 )
 
 func BenchmarkFile(b *testing.B) {
-	entropies := make(Entropies, 10)
+	entropies := &Entropies{Entropies: make([]Entropy, 10)}
 	for range b.N {
 		_ = readFile(entropies, "testdata")
 	}
@@ -47,27 +47,27 @@ func TestEntropy(t *testing.T) {
 
 func TestReadFile(t *testing.T) {
 	t.Run("random.js", func(t *testing.T) {
-		res := make(Entropies, 10)
+		res := &Entropies{Entropies: make([]Entropy, 10)}
 		err := readFile(res, "testdata/random.js")
 		if err != nil {
 			t.Errorf("expected nil, got %v", err)
 		}
 
-		ExpectFloat(t, res[0].Entropy, 5.53614242151549)
-		Expect(t, res[0].LineNum, 7) // The token is hidden here
-		ExpectFloat(t, res[4].Entropy, 3.321928094887362)
+		ExpectFloat(t, res.Entropies[0].Entropy, 5.53614242151549)
+		Expect(t, res.Entropies[0].LineNum, 7) // The token is hidden here
+		ExpectFloat(t, res.Entropies[4].Entropy, 3.321928094887362)
 	})
 
 	t.Run("testdata/folder", func(t *testing.T) {
-		res := make(Entropies, 10)
+		res := &Entropies{Entropies: make([]Entropy, 10)}
 		err := readFile(res, "testdata/folder")
 		if err != nil {
 			t.Errorf("expected nil, got %v", err)
 		}
 
-		ExpectFloat(t, res[0].Entropy, 3.7667029194153567)
-		Expect(t, res[0].LineNum, 7) // The token is hidden here
-		ExpectFloat(t, res[6].Entropy, 2.8553885422075336)
+		ExpectFloat(t, res.Entropies[0].Entropy, 3.7667029194153567)
+		Expect(t, res.Entropies[0].LineNum, 7) // The token is hidden here
+		ExpectFloat(t, res.Entropies[6].Entropy, 2.8553885422075336)
 	})
 
 	t.Run("dangling symlink in testdata folder", func(t *testing.T) {
@@ -120,16 +120,16 @@ func TestIsFileHidden(t *testing.T) {
 }
 
 func TestEntropies(t *testing.T) {
-	entropies := make(Entropies, 5)
+	res := &Entropies{Entropies: make([]Entropy, 5)}
 	for _, i := range []float64{1, 3, 5, 7, 2, 4, 6, 8} {
-		entropies.Add(Entropy{Entropy: i})
+		res.Add(Entropy{Entropy: i})
 	}
 
-	Expect(t, entropies[0].Entropy, 8)
-	Expect(t, entropies[1].Entropy, 7)
-	Expect(t, entropies[2].Entropy, 6)
-	Expect(t, entropies[3].Entropy, 5)
-	Expect(t, entropies[4].Entropy, 4)
+	Expect(t, res.Entropies[0].Entropy, 8)
+	Expect(t, res.Entropies[1].Entropy, 7)
+	Expect(t, res.Entropies[2].Entropy, 6)
+	Expect(t, res.Entropies[3].Entropy, 5)
+	Expect(t, res.Entropies[4].Entropy, 4)
 }
 
 func Expect[T comparable](t *testing.T, got, expected T) {