Add protocol tweak to set UDP payload size

completions/dog.bash

@@ -19,7 +19,7 @@ _dog()
             ;;
 
         -Z)
-            COMPREPLY=( $( compgen -W 'aa ad cd' -- "$cur" ) )
+            COMPREPLY=( $( compgen -W 'aa ad bufsize= cd' -- "$cur" ) )
             return
             ;;
 

completions/dog.fish

@@ -19,6 +19,7 @@ complete -c dog        -l 'txid'       -d "Set the transaction ID to a specific
 complete -c dog -s 'Z'                 -d "Configure uncommon protocol-level tweaks" -x -a "
     aa\t'Set the AA (Authoritative Answers) query bit'
     ad\t'Set the AD (Authentic Data) query bit'
+    bufsize=\t'Set the UDP payload size'
     cd\t'Set the CD (Checking Disabled) query bit'
 "
 

completions/dog.zsh

@@ -10,7 +10,7 @@ __dog() {
         --class"[Network class of the DNS record being queried]:(network class):(IN CH HS)" \
         --edns"[Whether to OPT in to EDNS]:(edns setting):(disable hide show)" \
         --txid"[Set the transaction ID to a specific value]" \
-        -Z"[Configure uncommon protocol-level tweaks]:(protocol tweak):(aa ad cd)" \
+        -Z"[Configure uncommon protocol-level tweaks]:(protocol tweak):(aa ad bufsize= cd)" \
         {-U,--udp}"[Use the DNS protocol over UDP]" \
         {-T,--tcp}"[Use the DNS protocol over TCP]" \
         {-S,--tls}"[Use the DNS-over-TLS protocol]" \

src/options.rs

@@ -380,6 +380,18 @@ impl ProtocolTweaks {
                     tweaks.set_checking_disabled_flag = true;
                 }
                 otherwise => {
+                    if let Some(remaining_num) = tweak_str.strip_prefix("bufsize=") {
+                        match remaining_num.parse() {
+                            Ok(parsed_bufsize) => {
+                                tweaks.udp_payload_size = Some(parsed_bufsize);
+                                continue;
+                            }
+                            Err(e) => {
+                                warn!("Failed to parse buffer size: {}", e);
+                            }
+                        }
+                    }
+
                     return Err(OptionsError::InvalidTweak(otherwise.into()));
                 }
             }
@@ -677,6 +689,24 @@ mod test {
         assert_eq!(options.requests.protocol_tweaks.set_checking_disabled_flag, true);
     }
 
+    #[test]
+    fn udp_size() {
+        let options = Options::getopts(&[ "dom.ain", "-Z", "bufsize=4096" ]).unwrap();
+        assert_eq!(options.requests.protocol_tweaks.udp_payload_size, Some(4096));
+    }
+
+    #[test]
+    fn udp_size_invalid() {
+        assert_eq!(Options::getopts(&[ "-Z", "bufsize=null" ]),
+                   OptionsResult::InvalidOptions(OptionsError::InvalidTweak("bufsize=null".into())));
+    }
+
+    #[test]
+    fn udp_size_too_big() {
+        assert_eq!(Options::getopts(&[ "-Z", "bufsize=999999999" ]),
+                   OptionsResult::InvalidOptions(OptionsError::InvalidTweak("bufsize=999999999".into())));
+    }
+
     #[test]
     fn short_mode() {
         let tf = TextFormat { format_durations: true };

src/requests.rs

@@ -57,6 +57,9 @@ pub struct ProtocolTweaks {
 
     /// Set the `CD` (Checking Disabled) flag in the header of each request.
     pub set_checking_disabled_flag: bool,
+
+    /// Set the buffer size field in the OPT record of each request.
+    pub udp_payload_size: Option<u16>,
 }
 
 /// Whether to send or display OPT packets.
@@ -98,7 +101,9 @@ impl RequestGenerator {
 
                             let mut additional = None;
                             if self.edns.should_send() {
-                                additional = Some(dns::Request::additional_record());
+                                let mut opt = dns::Request::additional_record();
+                                self.protocol_tweaks.set_request_opt_fields(&mut opt);
+                                additional = Some(opt);
                             }
 
                             let query = dns::Query { qname: domain.clone(), qtype, qclass };
@@ -143,4 +148,10 @@ impl ProtocolTweaks {
             flags.checking_disabled = true;
         }
     }
+
+    pub fn set_request_opt_fields(self, opt: &mut dns::record::OPT) {
+        if let Some(bufsize) = self.udp_payload_size {
+            opt.udp_payload_size = bufsize;
+        }
+    }
 }