From 0d895750b4663eb81c9d7093754de641c7df28a2 Mon Sep 17 00:00:00 2001
From: crn4 <vladimir.klementyev@gmail.com>
Date: Tue, 16 Sep 2025 18:49:08 +0200
Subject: [PATCH] implemented suppressor with loging only - no drop messages

---
 signal/peer/peer.go                  | 14 ++++++++-----
 signal/server/signal.go              | 30 ++++++++++++++++++++--------
 signal/suppressor/suppressor.go      |  9 ++++-----
 signal/suppressor/suppressor_test.go |  4 ++--
 4 files changed, 37 insertions(+), 20 deletions(-)

diff --git a/signal/peer/peer.go b/signal/peer/peer.go
index c9dd60fc0..0e8b7ed21 100644
--- a/signal/peer/peer.go
+++ b/signal/peer/peer.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/netbirdio/netbird/shared/signal/proto"
 	"github.com/netbirdio/netbird/signal/metrics"
+	"github.com/netbirdio/netbird/signal/suppressor"
 )
 
 var (
@@ -20,13 +21,11 @@ var (
 // Peer representation of a connected Peer
 type Peer struct {
 	// a unique id of the Peer (e.g. sha256 fingerprint of the Wireguard public key)
-	Id string
-
-	StreamID int64
-
+	Id         string
+	suppressor *suppressor.Suppressor
+	StreamID   int64
 	// a gRpc connection stream to the Peer
 	Stream proto.SignalExchange_ConnectStreamServer
-
 	// registration time
 	RegisteredAt time.Time
 
@@ -41,6 +40,7 @@ func NewPeer(id string, stream proto.SignalExchange_ConnectStreamServer, cancel
 		StreamID:     time.Now().UnixNano(),
 		RegisteredAt: time.Now(),
 		Cancel:       cancel,
+		suppressor:   suppressor.NewSuppressor(nil),
 	}
 }
 
@@ -117,3 +117,7 @@ func (registry *Registry) Deregister(peer *Peer) {
 		registry.metrics.Deregistrations.Add(context.Background(), 1)
 	}
 }
+
+func (peer *Peer) SendMessageAllowed(destination string, size int, arrivedTime time.Time) bool {
+	return peer.suppressor.PackageReceived(suppressor.PeerID(destination), size, arrivedTime)
+}
diff --git a/signal/server/signal.go b/signal/server/signal.go
index 47f01edae..11db01cbd 100644
--- a/signal/server/signal.go
+++ b/signal/server/signal.go
@@ -23,14 +23,16 @@ import (
 )
 
 const (
-	labelType              = "type"
-	labelTypeError         = "error"
-	labelTypeNotConnected  = "not_connected"
-	labelTypeNotRegistered = "not_registered"
-	labelTypeStream        = "stream"
-	labelTypeMessage       = "message"
-	labelTypeTimeout       = "timeout"
-	labelTypeDisconnected  = "disconnected"
+	labelType                    = "type"
+	labelTypeError               = "error"
+	labelTypeNotConnected        = "not_connected"
+	labelTypeNotRegistered       = "not_registered"
+	labelTypeSenderNotRegistered = "sender_not_registered"
+	labelTypeMessageSuppressed   = "message_suppressed"
+	labelTypeStream              = "stream"
+	labelTypeMessage             = "message"
+	labelTypeTimeout             = "timeout"
+	labelTypeDisconnected        = "disconnected"
 
 	labelError                   = "error"
 	labelErrorMissingId          = "missing_id"
@@ -95,6 +97,18 @@ func NewServer(ctx context.Context, meter metric.Meter) (*Server, error) {
 func (s *Server) Send(ctx context.Context, msg *proto.EncryptedMessage) (*proto.EncryptedMessage, error) {
 	log.Tracef("received a new message to send from peer [%s] to peer [%s]", msg.Key, msg.RemoteKey)
 
+	peer, found := s.registry.Get(msg.Key)
+	if !found {
+		s.metrics.MessageForwardFailures.Add(ctx, 1, metric.WithAttributes(attribute.String(labelType, labelTypeSenderNotRegistered)))
+		log.Tracef("message from peer [%s] can't be forwarded to peer [%s] because sender peer is not registered", msg.Key, msg.RemoteKey)
+		return nil, status.Errorf(codes.FailedPrecondition, "peer not registered")
+	}
+
+	if !peer.SendMessageAllowed(msg.RemoteKey, len(msg.Body), time.Now()) {
+		s.metrics.MessageForwardFailures.Add(ctx, 1, metric.WithAttributes(attribute.String(labelType, labelTypeMessageSuppressed)))
+		log.Tracef("message from peer [%s] to peer [%s] suppressed due to repetition", msg.Key, msg.RemoteKey)
+	}
+
 	if _, found := s.registry.Get(msg.RemoteKey); found {
 		s.forwardMessageToPeer(ctx, msg)
 		return &proto.EncryptedMessage{}, nil
diff --git a/signal/suppressor/suppressor.go b/signal/suppressor/suppressor.go
index 7ce2244c4..92c6214ac 100644
--- a/signal/suppressor/suppressor.go
+++ b/signal/suppressor/suppressor.go
@@ -1,7 +1,6 @@
 package suppressor
 
 import (
-	"fmt"
 	"time"
 )
 
@@ -63,20 +62,20 @@ type Suppressor struct {
 	peers               map[PeerID]*packageStat
 }
 
-func NewSuppressor(opts *Opts) (*Suppressor, error) {
+func NewSuppressor(opts *Opts) *Suppressor {
 	threshold := DefaultRepetitionThreshold
 	if opts != nil {
+		threshold = opts.RepetitionThreshold
 		if opts.RepetitionThreshold < minRepetitionThreshold {
-			return nil, fmt.Errorf("invalid repetition threshold")
+			threshold = DefaultRepetitionThreshold
 		}
 
-		threshold = opts.RepetitionThreshold
 	}
 
 	return &Suppressor{
 		repetitionThreshold: threshold,
 		peers:               make(map[PeerID]*packageStat),
-	}, nil
+	}
 }
 
 // PackageReceived handles a newly received package from a peer.
diff --git a/signal/suppressor/suppressor_test.go b/signal/suppressor/suppressor_test.go
index 2e96f88bd..5eca1d20c 100644
--- a/signal/suppressor/suppressor_test.go
+++ b/signal/suppressor/suppressor_test.go
@@ -7,7 +7,7 @@ import (
 
 func TestSuppressor_PackageReceived(t *testing.T) {
 	destID := PeerID("remote")
-	s, _ := NewSuppressor(&Opts{RepetitionThreshold: 3})
+	s := NewSuppressor(&Opts{RepetitionThreshold: 3})
 
 	// Define sequence with base deltas (s ±10% tolerance)
 	deltas := []time.Duration{
@@ -68,7 +68,7 @@ func TestSuppressor_PackageReceived(t *testing.T) {
 
 func TestSuppressor_PackageReceivedReset(t *testing.T) {
 	destID := PeerID("remote")
-	s, _ := NewSuppressor(&Opts{RepetitionThreshold: 5})
+	s := NewSuppressor(&Opts{RepetitionThreshold: 5})
 
 	// Define sequence with base deltas (s ±10% tolerance)
 	deltas := []time.Duration{