From 65c40ee92dd5bcad1ab52cbafa1afd62cf669e48 Mon Sep 17 00:00:00 2001
From: Pavel Odintsov <pavel.odintsov@gmail.com>
Date: Fri, 13 Dec 2024 02:32:10 +0300
Subject: [PATCH] Added capping logic for sFlow counter and flow samples to
 reduce chances of DoS

---
 src/libsflow/libsflow.hpp            |  6 ++++++
 src/sflow_plugin/sflow_collector.cpp | 19 +++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/src/libsflow/libsflow.hpp b/src/libsflow/libsflow.hpp
index f32a0388..611a9e94 100644
--- a/src/libsflow/libsflow.hpp
+++ b/src/libsflow/libsflow.hpp
@@ -18,6 +18,12 @@ const uint32_t max_udp_packet_size = 65535;
 // We need to limit number of samples by reasonable number
 const int32_t max_sflow_sample_number = 256;
 
+// We need to limit number of counter samples by reasonable number
+const uint32_t max_number_of_counter_records = 256;
+
+// We need to limit number of flow samples by reasonable number
+const uint32_t max_number_of_flow_records = 256;
+
 enum class sflow_sample_type_t : unsigned int {
     FLOW_SAMPLE             = 1,
     COUNTER_SAMPLE          = 2,
diff --git a/src/sflow_plugin/sflow_collector.cpp b/src/sflow_plugin/sflow_collector.cpp
index 98754fa8..01c71dd4 100644
--- a/src/sflow_plugin/sflow_collector.cpp
+++ b/src/sflow_plugin/sflow_collector.cpp
@@ -330,6 +330,17 @@ bool process_sflow_flow_sample(const uint8_t* data_pointer,
         return false;
     }
 
+    if (sflow_sample_header_unified_accessor.get_number_of_flow_records() > max_number_of_flow_records) {
+        logger << log4cpp::Priority::ERROR << plugin_log_prefix << "flow records number "
+            << sflow_sample_header_unified_accessor.get_number_of_flow_records()
+            << " exceeds maximum value "
+            << max_number_of_flow_records;
+
+        sflow_bad_flow_samples++;
+
+        return false;
+    }
+
     const uint8_t* flow_record_zone_start = data_pointer + sflow_sample_header_unified_accessor.get_original_payload_length();
 
     std::vector<record_tuple_t> vector_tuple;
@@ -609,6 +620,14 @@ bool process_sflow_counter_sample(const uint8_t* data_pointer,
         return false;
     }
 
+    if (sflow_counter_header_unified_accessor.get_number_of_counter_records() > max_number_of_counter_records) {
+        logger << log4cpp::Priority::ERROR << plugin_log_prefix << "number of counter records "
+                << sflow_counter_header_unified_accessor.get_number_of_counter_records()
+                << " exceeds maximum value "
+                << max_number_of_counter_records;
+        return false;
+    }
+
     std::vector<counter_record_sample_t> counter_record_sample_vector;
     counter_record_sample_vector.reserve(sflow_counter_header_unified_accessor.get_number_of_counter_records());