From 4202cae860a066d408512af6ccfc46e7e75dc898 Mon Sep 17 00:00:00 2001 From: Emil Hernvall Date: Thu, 18 Jun 2020 02:23:42 +0200 Subject: [PATCH] Eliminate all useless allocation --- chapter2.md | 2 +- chapter4.md | 21 ++++++++-------- chapter5.md | 59 ++++++++++++++++++++++++++------------------- examples/sample2.rs | 2 +- examples/sample3.rs | 2 +- examples/sample4.rs | 25 +++++++++---------- examples/sample5.rs | 45 +++++++++++++++++----------------- 7 files changed, 81 insertions(+), 75 deletions(-) diff --git a/chapter2.md b/chapter2.md index 0a4588e..7cd02f5 100644 --- a/chapter2.md +++ b/chapter2.md @@ -97,7 +97,7 @@ impl DnsHeader { )?; buffer.write_u8( - (self.rescode.clone() as u8) + (self.rescode as u8) | ((self.checking_disabled as u8) << 4) | ((self.authed_data as u8) << 5) | ((self.z as u8) << 6) diff --git a/chapter4.md b/chapter4.md index bd851d4..e4210a7 100644 --- a/chapter4.md +++ b/chapter4.md @@ -190,7 +190,7 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { // Next, `DnsPacket::from_buffer` is used to parse the raw bytes into // a `DnsPacket`. - let request = DnsPacket::from_buffer(&mut req_buffer)?; + let mut request = DnsPacket::from_buffer(&mut req_buffer)?; // Create and initialize the response packet let mut packet = DnsPacket::new(); @@ -199,15 +199,8 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { packet.header.recursion_available = true; packet.header.response = true; - // Being mindful of how unreliable input data from arbitrary senders can be, we - // need make sure that a question is actually present. If not, we return `FORMERR` - // to indicate that the sender made something wrong. - if request.questions.is_empty() { - packet.header.rescode = ResultCode::FORMERR; - } - // Usually a question will be present, though. - else { - let question = &request.questions[0]; + // In the normal case, exactly one question is present + if let Some(question) = request.questions.pop() { println!("Received query: {:?}", question); // Since all is set up and as expected, the query can be forwarded to the @@ -216,7 +209,7 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { // as much to the client. If rather everything goes as planned, the // question and response records as copied into our response packet. if let Ok(result) = lookup(&question.name, question.qtype) { - packet.questions.push(question.clone()); + packet.questions.push(question); packet.header.rescode = result.header.rescode; for rec in result.answers { @@ -235,6 +228,12 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { packet.header.rescode = ResultCode::SERVFAIL; } } + // Being mindful of how unreliable input data from arbitrary senders can be, we + // need make sure that a question is actually present. If not, we return `FORMERR` + // to indicate that the sender made something wrong. + else { + packet.header.rescode = ResultCode::FORMERR; + } // The only thing remaining is to encode our response and send it off! let mut res_buffer = BytePacketBuffer::new(); diff --git a/chapter5.md b/chapter5.md index b85a3f9..9813fc7 100644 --- a/chapter5.md +++ b/chapter5.md @@ -171,11 +171,11 @@ impl DnsPacket { /// It's useful to be able to pick a random A record from a packet. When we /// get multiple IP's for a single name, it doesn't matter which one we /// choose, so in those cases we can now pick one at random. - pub fn get_random_a(&self) -> Option { + pub fn get_random_a(&self) -> Option { self.answers .iter() .filter_map(|record| match record { - DnsRecord::A { ref addr, .. } => Some(addr.to_string()), + DnsRecord::A { addr, .. } => Some(*addr), _ => None, }) .next() @@ -183,8 +183,9 @@ impl DnsPacket { /// A helper function which returns an iterator over all name servers in /// the authorities section, represented as (domain, host) tuples - fn get_ns<'a>(&'a self, qname: &'a str) -> impl Iterator { - self.authorities.iter() + fn get_ns<'a>(&'a self, qname: &'a str) -> impl Iterator { + self.authorities + .iter() // In practice, these are always NS records in well formed packages. // Convert the NS records to a tuple which has only the data we need // to make it easy to work with. @@ -196,17 +197,18 @@ impl DnsPacket { .filter(move |(domain, _)| qname.ends_with(*domain)) } - /// When there is a NS record in the authorities section, there may also - /// be a matching A record in the additional section. This saves us - /// from doing a separate query to resolve the IP of the name server. - pub fn get_resolved_ns(&self, qname: &str) -> Option { + /// We'll use the fact that name servers often bundle the corresponding + /// A records when replying to an NS query to implement a function that + /// returns the actual IP for an NS record if possible. + pub fn get_resolved_ns(&self, qname: &str) -> Option { // Get an iterator over the nameservers in the authorities section self.get_ns(qname) // Now we need to look for a matching A record in the additional // section. Since we just want the first valid record, we can just // build a stream of matching records. .flat_map(|(_, host)| { - self.resources.iter() + self.resources + .iter() // Filter for A records where the domain match the host // of the NS record that we are currently processing .filter_map(move |record| match record { @@ -214,19 +216,19 @@ impl DnsPacket { _ => None, }) }) - .map(|addr| addr.to_string()) + .map(|addr| *addr) // Finally, pick the first valid entry .next() } /// However, not all name servers are as that nice. In certain cases there won't /// be any A records in the additional section, and we'll have to perform *another* - /// lookup in the midst of our first. For this, we introduce a method for - /// returning the hostname of an appropriate name server. - pub fn get_unresolved_ns(&self, qname: &str) -> Option { + /// lookup in the midst. For this, we introduce a method for returning the host + /// name of an appropriate name server. + pub fn get_unresolved_ns<'a>(&'a self, qname: &'a str) -> Option<&'a str> { // Get an iterator over the nameservers in the authorities section self.get_ns(qname) - .map(|(_, host)| host.to_string()) + .map(|(_, host)| host) // Finally, pick the first valid entry .next() } @@ -240,36 +242,35 @@ We move swiftly on to our new `recursive_lookup` function: ```rust fn recursive_lookup(qname: &str, qtype: QueryType) -> Result { - // For now we're always starting with *a.root-servers.net*. - let mut ns = "198.41.0.4".to_string(); + let mut ns = "198.41.0.4".parse::().unwrap(); // Since it might take an arbitrary number of steps, we enter an unbounded loop. loop { println!("attempting lookup of {:?} {} with ns {}", qtype, qname, ns); // The next step is to send the query to the active server. - let ns_copy = ns.clone(); + let ns_copy = ns; - let server = (ns_copy.as_str(), 53); - let response = lookup(qname, qtype.clone(), server)?; + let server = (ns_copy, 53); + let response = lookup(qname, qtype, server)?; // If there are entries in the answer section, and no errors, we are done! if !response.answers.is_empty() && response.header.rescode == ResultCode::NOERROR { - return Ok(response.clone()); + return Ok(response); } // We might also get a `NXDOMAIN` reply, which is the authoritative name servers // way of telling us that the name doesn't exist. if response.header.rescode == ResultCode::NXDOMAIN { - return Ok(response.clone()); + return Ok(response); } // Otherwise, we'll try to find a new nameserver based on NS and a corresponding A // record in the additional section. If this succeeds, we can switch name server // and retry the loop. if let Some(new_ns) = response.get_resolved_ns(qname) { - ns = new_ns.clone(); + ns = new_ns; continue; } @@ -278,7 +279,7 @@ fn recursive_lookup(qname: &str, qtype: QueryType) -> Result { // we'll go with what the last server told us. let new_ns_name = match response.get_unresolved_ns(qname) { Some(x) => x, - None => return Ok(response.clone()), + None => return Ok(response), }; // Here we go down the rabbit hole by starting _another_ lookup sequence in the @@ -289,14 +290,22 @@ fn recursive_lookup(qname: &str, qtype: QueryType) -> Result { // Finally, we pick a random ip from the result, and restart the loop. If no such // record is available, we again return the last result we got. if let Some(new_ns) = recursive_response.get_random_a() { - ns = new_ns.clone(); + ns = new_ns; } else { - return Ok(response.clone()); + return Ok(response); } } } ``` +This also requires a small change to the `lookup` function, as we need to +pass which server to use. We add a server parameter to the function signature, +and remove the hardcoded `server` variable we used in chapter 4: + +```rust +fn lookup(qname: &str, qtype: QueryType, server: (Ipv4Addr, u16)) -> Result { +``` + ### Trying out recursive lookup The only thing remaining is to change our `handle_query` function to use diff --git a/examples/sample2.rs b/examples/sample2.rs index cf0dba4..345291e 100644 --- a/examples/sample2.rs +++ b/examples/sample2.rs @@ -290,7 +290,7 @@ impl DnsHeader { )?; buffer.write_u8( - (self.rescode.clone() as u8) + (self.rescode as u8) | ((self.checking_disabled as u8) << 4) | ((self.authed_data as u8) << 5) | ((self.z as u8) << 6) diff --git a/examples/sample3.rs b/examples/sample3.rs index e995562..fa8c18c 100644 --- a/examples/sample3.rs +++ b/examples/sample3.rs @@ -303,7 +303,7 @@ impl DnsHeader { )?; buffer.write_u8( - (self.rescode.clone() as u8) + (self.rescode as u8) | ((self.checking_disabled as u8) << 4) | ((self.authed_data as u8) << 5) | ((self.z as u8) << 6) diff --git a/examples/sample4.rs b/examples/sample4.rs index b446362..eee856b 100644 --- a/examples/sample4.rs +++ b/examples/sample4.rs @@ -303,7 +303,7 @@ impl DnsHeader { )?; buffer.write_u8( - (self.rescode.clone() as u8) + (self.rescode as u8) | ((self.checking_disabled as u8) << 4) | ((self.authed_data as u8) << 5) | ((self.z as u8) << 6) @@ -643,7 +643,7 @@ impl DnsPacket { result.header.read(buffer)?; for _ in 0..result.header.questions { - let mut question = DnsQuestion::new("".to_string(), QueryType::UNKNOWN(0)); + let mut question = DnsQuestion::new(String::new(), QueryType::UNKNOWN(0)); question.read(buffer)?; result.questions.push(question); } @@ -728,7 +728,7 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { // Next, `DnsPacket::from_buffer` is used to parse the raw bytes into // a `DnsPacket`. - let request = DnsPacket::from_buffer(&mut req_buffer)?; + let mut request = DnsPacket::from_buffer(&mut req_buffer)?; // Create and initialize the response packet let mut packet = DnsPacket::new(); @@ -737,15 +737,8 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { packet.header.recursion_available = true; packet.header.response = true; - // Being mindful of how unreliable input data from arbitrary senders can be, we - // need make sure that a question is actually present. If not, we return `FORMERR` - // to indicate that the sender made something wrong. - if request.questions.is_empty() { - packet.header.rescode = ResultCode::FORMERR; - } - // Usually a question will be present, though. - else { - let question = &request.questions[0]; + // In the normal case, exactly one question is present + if let Some(question) = request.questions.pop() { println!("Received query: {:?}", question); // Since all is set up and as expected, the query can be forwarded to the @@ -754,7 +747,7 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { // as much to the client. If rather everything goes as planned, the // question and response records as copied into our response packet. if let Ok(result) = lookup(&question.name, question.qtype) { - packet.questions.push(question.clone()); + packet.questions.push(question); packet.header.rescode = result.header.rescode; for rec in result.answers { @@ -773,6 +766,12 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { packet.header.rescode = ResultCode::SERVFAIL; } } + // Being mindful of how unreliable input data from arbitrary senders can be, we + // need make sure that a question is actually present. If not, we return `FORMERR` + // to indicate that the sender made something wrong. + else { + packet.header.rescode = ResultCode::FORMERR; + } // The only thing remaining is to encode our response and send it off! let mut res_buffer = BytePacketBuffer::new(); diff --git a/examples/sample5.rs b/examples/sample5.rs index fe3fd49..15abd8c 100644 --- a/examples/sample5.rs +++ b/examples/sample5.rs @@ -303,7 +303,7 @@ impl DnsHeader { )?; buffer.write_u8( - (self.rescode.clone() as u8) + (self.rescode as u8) | ((self.checking_disabled as u8) << 4) | ((self.authed_data as u8) << 5) | ((self.z as u8) << 6) @@ -691,11 +691,11 @@ impl DnsPacket { /// It's useful to be able to pick a random A record from a packet. When we /// get multiple IP's for a single name, it doesn't matter which one we /// choose, so in those cases we can now pick one at random. - pub fn get_random_a(&self) -> Option { + pub fn get_random_a(&self) -> Option { self.answers .iter() .filter_map(|record| match record { - DnsRecord::A { ref addr, .. } => Some(addr.to_string()), + DnsRecord::A { addr, .. } => Some(*addr), _ => None, }) .next() @@ -720,7 +720,7 @@ impl DnsPacket { /// We'll use the fact that name servers often bundle the corresponding /// A records when replying to an NS query to implement a function that /// returns the actual IP for an NS record if possible. - pub fn get_resolved_ns(&self, qname: &str) -> Option { + pub fn get_resolved_ns(&self, qname: &str) -> Option { // Get an iterator over the nameservers in the authorities section self.get_ns(qname) // Now we need to look for a matching A record in the additional @@ -736,7 +736,7 @@ impl DnsPacket { _ => None, }) }) - .map(|addr| addr.to_string()) + .map(|addr| *addr) // Finally, pick the first valid entry .next() } @@ -745,16 +745,16 @@ impl DnsPacket { /// be any A records in the additional section, and we'll have to perform *another* /// lookup in the midst. For this, we introduce a method for returning the host /// name of an appropriate name server. - pub fn get_unresolved_ns(&self, qname: &str) -> Option { + pub fn get_unresolved_ns<'a>(&'a self, qname: &'a str) -> Option<&'a str> { // Get an iterator over the nameservers in the authorities section self.get_ns(qname) - .map(|(_, host)| host.to_string()) + .map(|(_, host)| host) // Finally, pick the first valid entry .next() } } -fn lookup(qname: &str, qtype: QueryType, server: (&str, u16)) -> Result { +fn lookup(qname: &str, qtype: QueryType, server: (Ipv4Addr, u16)) -> Result { let socket = UdpSocket::bind(("0.0.0.0", 43210))?; let mut packet = DnsPacket::new(); @@ -778,34 +778,34 @@ fn lookup(qname: &str, qtype: QueryType, server: (&str, u16)) -> Result Result { // For now we're always starting with *a.root-servers.net*. - let mut ns = "198.41.0.4".to_string(); + let mut ns = "198.41.0.4".parse::().unwrap(); // Since it might take an arbitrary number of steps, we enter an unbounded loop. loop { println!("attempting lookup of {:?} {} with ns {}", qtype, qname, ns); // The next step is to send the query to the active server. - let ns_copy = ns.clone(); + let ns_copy = ns; - let server = (ns_copy.as_str(), 53); - let response = lookup(qname, qtype.clone(), server)?; + let server = (ns_copy, 53); + let response = lookup(qname, qtype, server)?; // If there are entries in the answer section, and no errors, we are done! if !response.answers.is_empty() && response.header.rescode == ResultCode::NOERROR { - return Ok(response.clone()); + return Ok(response); } // We might also get a `NXDOMAIN` reply, which is the authoritative name servers // way of telling us that the name doesn't exist. if response.header.rescode == ResultCode::NXDOMAIN { - return Ok(response.clone()); + return Ok(response); } // Otherwise, we'll try to find a new nameserver based on NS and a corresponding A // record in the additional section. If this succeeds, we can switch name server // and retry the loop. if let Some(new_ns) = response.get_resolved_ns(qname) { - ns = new_ns.clone(); + ns = new_ns; continue; } @@ -814,7 +814,7 @@ fn recursive_lookup(qname: &str, qtype: QueryType) -> Result { // we'll go with what the last server told us. let new_ns_name = match response.get_unresolved_ns(qname) { Some(x) => x, - None => return Ok(response.clone()), + None => return Ok(response), }; // Here we go down the rabbit hole by starting _another_ lookup sequence in the @@ -825,9 +825,9 @@ fn recursive_lookup(qname: &str, qtype: QueryType) -> Result { // Finally, we pick a random ip from the result, and restart the loop. If no such // record is available, we again return the last result we got. if let Some(new_ns) = recursive_response.get_random_a() { - ns = new_ns.clone(); + ns = new_ns; } else { - return Ok(response.clone()); + return Ok(response); } } } @@ -836,7 +836,7 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { let mut req_buffer = BytePacketBuffer::new(); let (_, src) = socket.recv_from(&mut req_buffer.buf)?; - let request = DnsPacket::from_buffer(&mut req_buffer)?; + let mut request = DnsPacket::from_buffer(&mut req_buffer)?; let mut packet = DnsPacket::new(); packet.header.id = request.header.id; @@ -844,10 +844,7 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { packet.header.recursion_available = true; packet.header.response = true; - if request.questions.is_empty() { - packet.header.rescode = ResultCode::FORMERR; - } else { - let question = &request.questions[0]; + if let Some(question) = request.questions.pop() { println!("Received query: {:?}", question); if let Ok(result) = recursive_lookup(&question.name, question.qtype) { @@ -869,6 +866,8 @@ fn handle_query(socket: &UdpSocket) -> Result<()> { } else { packet.header.rescode = ResultCode::SERVFAIL; } + } else { + packet.header.rescode = ResultCode::FORMERR; } let mut res_buffer = BytePacketBuffer::new();