From 9d9a901b68d3bc5b2996c0e78400283087cb686d Mon Sep 17 00:00:00 2001
From: Jason Ish <jason.ish@oisf.net>
Date: Sun, 7 Jun 2020 14:07:26 +0200
Subject: [PATCH] dns: conditional logging

Apply config to newly created response TX.
---
 rust/src/dns/dns.rs | 62 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 60 insertions(+), 2 deletions(-)

diff --git a/rust/src/dns/dns.rs b/rust/src/dns/dns.rs
index eed5864048..e788106797 100644
--- a/rust/src/dns/dns.rs
+++ b/rust/src/dns/dns.rs
@@ -20,6 +20,8 @@ extern crate nom;
 use std;
 use std::ffi::CString;
 use std::mem::transmute;
+use std::collections::HashMap;
+use std::collections::VecDeque;
 
 use crate::log::*;
 use crate::applayer::*;
@@ -318,6 +320,36 @@ impl Drop for DNSTransaction {
     }
 }
 
+struct ConfigTracker {
+    map: HashMap<u16, AppLayerTxConfig>,
+    queue: VecDeque<u16>,
+}
+
+impl ConfigTracker {
+    fn new() -> ConfigTracker {
+        ConfigTracker {
+            map: HashMap::new(),
+            queue: VecDeque::new(),
+        }
+    }
+
+    fn add(&mut self, id: u16, config: AppLayerTxConfig) {
+        // If at size limit, remove the oldest entry.
+        if self.queue.len() > 499 {
+            if let Some(id) = self.queue.pop_front() {
+                self.map.remove(&id);
+            }
+        }
+
+        self.map.insert(id, config);
+        self.queue.push_back(id);
+    }
+
+    fn remove(&mut self, id: &u16) -> Option<AppLayerTxConfig> {
+        self.map.remove(id)
+    }
+}
+
 pub struct DNSState {
     // Internal transaction ID.
     pub tx_id: u64,
@@ -330,6 +362,8 @@ pub struct DNSState {
     pub request_buffer: Vec<u8>,
     pub response_buffer: Vec<u8>,
 
+    config: Option<ConfigTracker>,
+
     gap: bool,
 }
 
@@ -342,6 +376,7 @@ impl DNSState {
             events: 0,
             request_buffer: Vec::new(),
             response_buffer: Vec::new(),
+            config: None,
             gap: false,
         };
     }
@@ -355,6 +390,7 @@ impl DNSState {
             events: 0,
             request_buffer: Vec::with_capacity(0xffff),
             response_buffer: Vec::with_capacity(0xffff),
+            config: None,
             gap: false,
         };
     }
@@ -479,6 +515,11 @@ impl DNSState {
                 }
 
                 let mut tx = self.new_tx();
+                if let Some(ref mut config) = &mut self.config {
+                    if let Some(config) = config.remove(&response.header.tx_id) {
+                        tx.tx_data.config = config;
+                    }
+                }
                 tx.response = Some(response);
                 self.transactions.push(tx);
                 return true;
@@ -944,6 +985,23 @@ pub extern "C" fn rs_dns_probe_tcp(
     return 0;
 }
 
+#[no_mangle]
+pub extern "C" fn rs_dns_apply_tx_config(
+    _state: *mut std::os::raw::c_void, _tx: *mut std::os::raw::c_void,
+    _mode: std::os::raw::c_int, config: AppLayerTxConfig
+) {
+    let tx = cast_pointer!(_tx, DNSTransaction);
+    let state = cast_pointer!(_state, DNSState);
+    if let Some(request) = &tx.request {
+        if state.config.is_none() {
+            state.config = Some(ConfigTracker::new());
+        }
+        if let Some(ref mut tracker) = &mut state.config {
+            tracker.add(request.header.tx_id, config);
+        }
+    }
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn rs_dns_init(proto: AppProto) {
     ALPROTO_DNS = proto;
@@ -979,7 +1037,7 @@ pub unsafe extern "C" fn rs_dns_udp_register_parser() {
         get_de_state: rs_dns_state_get_tx_detect_state,
         set_de_state: rs_dns_state_set_tx_detect_state,
         get_tx_data: rs_dns_state_get_tx_data,
-        apply_tx_config: None,
+        apply_tx_config: Some(rs_dns_apply_tx_config),
     };
 
     let ip_proto_str = CString::new("udp").unwrap();
@@ -1024,7 +1082,7 @@ pub unsafe extern "C" fn rs_dns_tcp_register_parser() {
         get_de_state: rs_dns_state_get_tx_detect_state,
         set_de_state: rs_dns_state_set_tx_detect_state,
         get_tx_data: rs_dns_state_get_tx_data,
-        apply_tx_config: None,
+        apply_tx_config: Some(rs_dns_apply_tx_config),
     };
 
     let ip_proto_str = CString::new("tcp").unwrap();