From bf42e95368982f87a273bdf86d85022eaf3e4680 Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Mon, 23 Dec 2019 23:27:21 +0100
Subject: [PATCH] edns0 padding

TODO: check that a padding pseudorecord is not already present
---
 src/dns.rs   | 59 +++++++++++++++++++++++++++++++++++++++++++++++++++-
 src/main.rs  |  4 ++--
 src/utils.rs |  6 ------
 3 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/src/dns.rs b/src/dns.rs
index d6401ac..54b9b1c 100644
--- a/src/dns.rs
+++ b/src/dns.rs
@@ -3,10 +3,12 @@ use byteorder::{BigEndian, ByteOrder};
 
 const DNS_HEADER_SIZE: usize = 12;
 const DNS_MAX_HOSTNAME_SIZE: usize = 256;
-const DNS_MAX_PACKET_SIZE: usize = 65_535;
+const DNS_MAX_PACKET_SIZE: usize = 4096;
 const DNS_OFFSET_QUESTION: usize = DNS_HEADER_SIZE;
 const DNS_TYPE_OPT: u16 = 41;
 
+const DNS_PTYPE_PADDING: u16 = 12;
+
 const DNS_RCODE_SERVFAIL: u8 = 2;
 const DNS_RCODE_REFUSED: u8 = 5;
 
@@ -201,3 +203,58 @@ pub fn set_edns_max_payload_size(packet: &mut Vec<u8>, max_payload_size: u16) ->
     add_edns_section(packet, max_payload_size)?;
     Ok(())
 }
+
+pub fn add_edns_padding(packet: &mut Vec<u8>, block_size: usize) -> Result<(), Error> {
+    ensure!(qdcount(packet) == 1, "Unsupported number of questions");
+    let mut packet_len = packet.len();
+    ensure!(packet_len > DNS_OFFSET_QUESTION, "Short packet");
+    ensure!(packet_len <= DNS_MAX_PACKET_SIZE, "Large packet");
+
+    let mut offset = skip_name(packet, DNS_OFFSET_QUESTION)?;
+    assert!(offset > DNS_OFFSET_QUESTION);
+    ensure!(packet_len - offset >= 4, "Short packet");
+    offset += 4;
+    let (ancount, nscount, arcount) = (ancount(packet), nscount(packet), arcount(packet));
+    offset = traverse_rrs(packet, offset, ancount + nscount, |_offset| Ok(()))?;
+    let mut edns_offset = None;
+    traverse_rrs_mut(packet, offset, arcount, |packet, offset| {
+        let qtype = BigEndian::read_u16(&packet[offset..]);
+        if qtype == DNS_TYPE_OPT {
+            ensure!(edns_offset.is_none(), "Duplicate OPT RR found");
+            edns_offset = Some(offset)
+        }
+        Ok(())
+    })?;
+    let edns_offset = match edns_offset {
+        Some(edns_offset) => edns_offset,
+        None => {
+            let edns_offset = packet.len() + 1;
+            add_edns_section(packet, DNS_MAX_PACKET_SIZE as _)?;
+            packet_len = packet.len();
+            edns_offset
+        }
+    };
+    ensure!(packet_len < DNS_MAX_PACKET_SIZE, "Large packet");
+    let pad_len = (block_size - 1) - ((packet_len + (block_size - 1)) & (block_size - 1));
+    let mut edns_padding_prr = vec![0u8; 2 + pad_len];
+    edns_padding_prr[0] = (DNS_PTYPE_PADDING >> 8) as u8;
+    edns_padding_prr[1] = DNS_PTYPE_PADDING as u8;
+    let edns_padding_prr_len = edns_padding_prr.len();
+    let edns_rdlen_offset: usize = edns_offset + 8;
+    ensure!(packet_len - edns_rdlen_offset >= 2, "Short packet");
+    let edns_rdlen = BigEndian::read_u16(&packet[edns_rdlen_offset..]);
+    ensure!(
+        0xffff - edns_rdlen as usize >= edns_padding_prr_len,
+        "EDNS section too large for padding"
+    );
+    BigEndian::write_u16(
+        &mut packet[edns_rdlen_offset..],
+        edns_rdlen + edns_padding_prr_len as u16,
+    );
+    ensure!(
+        DNS_MAX_PACKET_SIZE - packet_len >= edns_padding_prr_len,
+        "Large packet"
+    );
+    packet.extend(&edns_padding_prr);
+    Ok(())
+}
diff --git a/src/main.rs b/src/main.rs
index db3c783..914b412 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -159,7 +159,7 @@ impl DoH {
         if query.len() < MIN_DNS_PACKET_LEN {
             return Err(DoHError::Incomplete);
         }
-        let _ = dns::set_edns_max_payload_size(&mut query, MAX_DNS_RESPONSE_LEN as u16);
+        let _ = dns::set_edns_max_payload_size(&mut query, MAX_DNS_RESPONSE_LEN as _);
         let globals = &self.globals;
         let mut socket = UdpSocket::bind(&globals.local_bind_address)
             .await
@@ -185,11 +185,11 @@ impl DoH {
                 Ok(ttl) => ttl,
             }
         };
+        dns::add_edns_padding(&mut packet, BLOCK_SIZE).map_err(|_| DoHError::TooLarge)?;
         let packet_len = packet.len();
         let response = Response::builder()
             .header(hyper::header::CONTENT_LENGTH, packet_len)
             .header(hyper::header::CONTENT_TYPE, "application/dns-message")
-            .header("X-Padding", utils::padding_string(packet_len, BLOCK_SIZE))
             .header(
                 hyper::header::CACHE_CONTROL,
                 format!("max-age={}", ttl).as_str(),
diff --git a/src/utils.rs b/src/utils.rs
index b170fa8..0214ed5 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,11 +1,5 @@
 use std::net::{SocketAddr, ToSocketAddrs};
 
-pub(crate) fn padding_string(input_size: usize, block_size: usize) -> String {
-    let block_size_ = block_size - 1;
-    let padding_len = block_size_ - ((input_size + block_size_) & block_size_);
-    String::from_utf8(vec![b'X'; padding_len]).unwrap()
-}
-
 // functions to verify the startup arguments as correct
 pub(crate) fn verify_sock_addr(arg_val: String) -> Result<(), String> {
     match arg_val.parse::<SocketAddr>() {