From 8e383d77d25684bcf8761547735e920fba14ab11 Mon Sep 17 00:00:00 2001
From: Alistair King <alistair@kentik.com>
Date: Mon, 22 Jul 2024 09:56:36 -0700
Subject: [PATCH] Add a limit on max DNS batch size

Otherwise we can exceed max payload sizes at ingest.
---
 src/mode/dns.rs | 30 +++++++++++++++++-------------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/src/mode/dns.rs b/src/mode/dns.rs
index a9e520a..54df8e8 100644
--- a/src/mode/dns.rs
+++ b/src/mode/dns.rs
@@ -1,3 +1,4 @@
+use std::char::MAX;
 use std::mem::swap;
 use std::net::IpAddr;
 use anyhow::Result;
@@ -17,6 +18,8 @@ use crate::protocol::dns::parser::{self, Rdata};
 use crate::reasm::Reassembler;
 use crate::time::Timestamp;
 
+const MAX_BUFFER_LEN: usize = 10000;
+
 pub struct Dns {
     asm:    Reassembler,
     buffer: Vec<Response>,
@@ -157,20 +160,21 @@ impl Dns {
     }
 
     fn flush(&mut self, ts: Timestamp) {
-        if (ts - self.last) >= Duration::seconds(1) {
-            let mut rs = Vec::with_capacity(self.buffer.len());
-            swap(&mut self.buffer, &mut rs);
-
-            let timeout = Duration::milliseconds(10).unsigned_abs();
-            let len = rs.len();
-            match self.client.send(rs, timeout) {
-                Ok(..) => debug!("DNS batch sent: {}", len),
-                Err(e) => warn!("DNS queue full: {:?}", e),
-            };
-
-            self.asm.flush(ts);
-            self.last = ts;
+        if (ts - self.last) < Duration::seconds(1) || self.buffer.len() >= MAX_BUFFER_LEN {
+            return;
         }
+        let mut rs = Vec::with_capacity(self.buffer.len());
+        swap(&mut self.buffer, &mut rs);
+
+        let timeout = Duration::milliseconds(10).unsigned_abs();
+        let len = rs.len();
+        match self.client.send(rs, timeout) {
+            Ok(..) => debug!("DNS batch sent: {}", len),
+            Err(e) => warn!("DNS queue full: {:?}", e),
+        };
+
+        self.asm.flush(ts);
+        self.last = ts;
     }
 
     fn tcp<'a>(&self, p: &Packet, tcp: &'a TcpPacket) -> (Addr, Addr, &'a [u8]) {