From b30d2b59c736d4239389d8fdb6d0a530231ebcd9 Mon Sep 17 00:00:00 2001
From: Vladislav Grishenko <themiron@yandex-team.ru>
Date: Sun, 31 Mar 2024 16:06:34 +0500
Subject: [PATCH] Fix ipfix buffer use and mirror test

---
 vflow/ipfix_test.go | 30 +++++++++++++++++++++++-------
 vflow/ipfix_unix.go |  4 +++-
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/vflow/ipfix_test.go b/vflow/ipfix_test.go
index e3ac5f88..5c6dd7f3 100644
--- a/vflow/ipfix_test.go
+++ b/vflow/ipfix_test.go
@@ -30,7 +30,7 @@ import (
 )
 
 func init() {
-	opts = &Options{}
+	opts = NewOptions()
 }
 
 func TestMirrorIPFIX(t *testing.T) {
@@ -39,6 +39,7 @@ func TestMirrorIPFIX(t *testing.T) {
 		fb    = make(chan IPFIXUDPMsg)
 		dst   = net.ParseIP("127.0.0.1")
 		ready = make(chan struct{})
+		errCh = make(chan struct{})
 	)
 
 	go func() {
@@ -46,14 +47,19 @@ func TestMirrorIPFIX(t *testing.T) {
 		if err != nil {
 			if strings.Contains(err.Error(), "not permitted") {
 				t.Log(err)
-				ready <- struct{}{}
+				errCh <- struct{}{}
 			} else {
 				t.Fatal("unexpected error", err)
 			}
 		}
 	}()
 
-	time.Sleep(2 * time.Second)
+	select {
+	case <-errCh:
+		return
+	case <-time.After(2 * time.Second):
+		break
+	}
 
 	go func() {
 		b := make([]byte, 1500)
@@ -86,9 +92,11 @@ func TestMirrorIPFIX(t *testing.T) {
 
 	}()
 
-	_, ok := <-ready
-	if ok {
-		return
+	select {
+	case <-ready:
+		break
+	case <-time.After(2 * time.Second):
+		t.Fatal("timeout")
 	}
 
 	body := []byte("hello")
@@ -100,7 +108,15 @@ func TestMirrorIPFIX(t *testing.T) {
 		},
 	}
 
-	feedback := <-fb
+	var feedback IPFIXUDPMsg
+	select {
+	case feedback = <-fb:
+		break
+	case <-errCh:
+		return
+	case <-time.After(2 * time.Second):
+		t.Fatal("timeout")
+	}
 
 	if string(feedback.body) != "hello" {
 		t.Error("expect body is hello, got", string(feedback.body))
diff --git a/vflow/ipfix_unix.go b/vflow/ipfix_unix.go
index 20041e4b..37ac6cfb 100644
--- a/vflow/ipfix_unix.go
+++ b/vflow/ipfix_unix.go
@@ -113,7 +113,9 @@ func mirrorIPFIX(dst net.IP, port int, ch chan IPFIXUDPMsg) error {
 		copy(packet[ipHLen:ipHLen+8], udpHdr)
 		copy(packet[ipHLen+8:], msg.body)
 
-		ipfixBuffer.Put(msg.body[:opts.IPFIXUDPSize])
+		if cap(msg.body) >= opts.IPFIXUDPSize {
+			ipfixBuffer.Put(msg.body[:opts.IPFIXUDPSize])
+		}
 
 		if err = conn.Send(packet[0 : ipHLen+8+pLen]); err != nil {
 			return err