From 6ccba029b0cd235978748ff231547645d3cacfc3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Geyslan=20Greg=C3=B3rio?= <geyslan@gmail.com>
Date: Tue, 8 Aug 2023 13:49:45 -0300
Subject: [PATCH] fix(selftest): use channel for sync (#364)

attachgenericfd and perfbuffers selftests were using a goroutine without
any syncronization mechanism.
---
 selftest/attachgenericfd/main.go | 32 ++++++++++++++++++++++++++++---
 selftest/perfbuffers/main.go     | 33 +++++++++++++++++++-------------
 2 files changed, 49 insertions(+), 16 deletions(-)

diff --git a/selftest/attachgenericfd/main.go b/selftest/attachgenericfd/main.go
index 132a7c16..103caf7d 100644
--- a/selftest/attachgenericfd/main.go
+++ b/selftest/attachgenericfd/main.go
@@ -3,9 +3,11 @@ package main
 import "C"
 
 import (
+	"bytes"
 	"fmt"
 	"net"
 	"os"
+	"time"
 	"unsafe"
 
 	bpf "github.com/aquasecurity/libbpfgo"
@@ -32,6 +34,7 @@ func main() {
 		os.Exit(-1)
 	}
 	defer unix.Close(serverFD)
+
 	serverAddr := &unix.SockaddrInet4{
 		Port: 22345,
 		Addr: [4]byte{127, 0, 0, 1},
@@ -40,6 +43,7 @@ func main() {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(-1)
 	}
+
 	if err := unix.Listen(serverFD, 100); err != nil {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(-1)
@@ -60,6 +64,7 @@ func main() {
 			os.Exit(-1)
 		}
 	}()
+
 	prog2, err := bpfModule.GetProgram("bpf_prog_verdict")
 	prog2.AttachGenericFD(sockMapRx.FileDescriptor(),
 		bpf.BPFAttachTypeSKSKBStreamVerdict, bpf.BPFFNone)
@@ -70,6 +75,8 @@ func main() {
 		}
 	}()
 
+	mapUpdateChan := make(chan struct{}, 1)
+
 	go func() {
 		acceptedFD, _, err := unix.Accept(serverFD)
 		if err != nil {
@@ -82,6 +89,8 @@ func main() {
 			fmt.Fprintln(os.Stderr, err)
 			os.Exit(-1)
 		}
+
+		mapUpdateChan <- struct{}{}
 	}()
 
 	c, err := net.Dial("tcp", "127.0.0.1:22345")
@@ -90,13 +99,30 @@ func main() {
 		os.Exit(-1)
 	}
 	defer c.Close()
-	if _, err = c.Write([]byte("foobar")); err != nil {
+
+	// wait for the bpf map to be updated
+	select {
+	case <-mapUpdateChan:
+	// continue with write/read
+	case <-time.After(15 * time.Second): // Same of the selftest
+		fmt.Fprintln(os.Stderr, "bpf map timeout")
+		os.Exit(-1)
+	}
+
+	input := []byte("foobar")
+	if _, err = c.Write(input); err != nil {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(-1)
 	}
-	data := make([]byte, 10)
-	if _, err = c.Read(data); err != nil {
+
+	output := make([]byte, 6)
+	if _, err = c.Read(output); err != nil {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(-1)
 	}
+
+	if !bytes.Equal(output, input) {
+		fmt.Fprintln(os.Stderr, "data mismatch")
+		os.Exit(-1)
+	}
 }
diff --git a/selftest/perfbuffers/main.go b/selftest/perfbuffers/main.go
index b100bad4..c01dd26a 100644
--- a/selftest/perfbuffers/main.go
+++ b/selftest/perfbuffers/main.go
@@ -5,6 +5,7 @@ import "C"
 import (
 	"os"
 	"runtime"
+	"time"
 
 	"encoding/binary"
 	"fmt"
@@ -67,26 +68,32 @@ func main() {
 
 	pb.Poll(300)
 
-	numberOfEventsReceived := 0
+	stop := make(chan struct{})
 
 	go func() {
 		for {
-			syscall.Mmap(999, 999, 999, 1, 1)
+			select {
+			case <-stop:
+				return
+			case b := <-eventsChannel:
+				if binary.LittleEndian.Uint32(b) != 2021 {
+					fmt.Fprintf(os.Stderr, "invalid data retrieved\n")
+					os.Exit(-1)
+				}
+			}
 		}
 	}()
-recvLoop:
-	for {
-		b := <-eventsChannel
-		if binary.LittleEndian.Uint32(b) != 2021 {
-			fmt.Fprintf(os.Stderr, "invalid data retrieved\n")
-			os.Exit(-1)
-		}
-		numberOfEventsReceived++
-		if numberOfEventsReceived > 5 {
-			break recvLoop
-		}
+
+	// give some time for the upper goroutine to start
+	time.Sleep(100 * time.Millisecond)
+
+	for sent := 0; sent < 5; sent++ {
+		syscall.Mmap(999, 999, 999, 1, 1)
+		time.Sleep(100 * time.Millisecond)
 	}
 
+	close(stop)
+
 	// Test that it won't cause a panic or block if Stop or Close called multiple times
 	pb.Stop()
 	pb.Stop()