Update examples with modernized code

This commit is contained in:
Emil Hernvall
2020-06-18 01:47:09 +02:00
parent 31369696d9
commit f815075ae4
10 changed files with 708 additions and 627 deletions

View File

@@ -68,14 +68,14 @@ Let's go ahead and add them to our code! First we'll update our `QueryType`
enum:
```rust
#[derive(PartialEq,Eq,Debug,Clone,Hash,Copy)]
#[derive(PartialEq, Eq, Debug, Clone, Hash, Copy)]
pub enum QueryType {
UNKNOWN(u16),
A, // 1
NS, // 2
A, // 1
NS, // 2
CNAME, // 5
MX, // 15
AAAA, // 28
MX, // 15
AAAA, // 28
}
```
@@ -101,7 +101,7 @@ impl QueryType {
5 => QueryType::CNAME,
15 => QueryType::MX,
28 => QueryType::AAAA,
_ => QueryType::UNKNOWN(num)
_ => QueryType::UNKNOWN(num),
}
}
}
@@ -113,40 +113,40 @@ Now we need a way of holding the data for these records, so we'll make some
modifications to `DnsRecord`.
```rust
#[derive(Debug,Clone,PartialEq,Eq,Hash,PartialOrd,Ord)]
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[allow(dead_code)]
pub enum DnsRecord {
UNKNOWN {
domain: String,
qtype: u16,
data_len: u16,
ttl: u32
ttl: u32,
}, // 0
A {
domain: String,
addr: Ipv4Addr,
ttl: u32
ttl: u32,
}, // 1
NS {
domain: String,
host: String,
ttl: u32
ttl: u32,
}, // 2
CNAME {
domain: String,
host: String,
ttl: u32
ttl: u32,
}, // 5
MX {
domain: String,
priority: u16,
host: String,
ttl: u32
ttl: u32,
}, // 15
AAAA {
domain: String,
addr: Ipv6Addr,
ttl: u32
ttl: u32,
}, // 28
}
```
@@ -156,106 +156,101 @@ and reading records. Starting with read, we amend it with additional code for
each record type. First off, we've got the common preamble:
```rust
pub fn read(buffer: &mut BytePacketBuffer) -> Result<DnsRecord> {
let mut domain = String::new();
try!(buffer.read_qname(&mut domain));
impl DnsRecord {
pub fn read(buffer: &mut BytePacketBuffer) -> Result<DnsRecord> {
let mut domain = String::new();
buffer.read_qname(&mut domain)?;
let qtype_num = try!(buffer.read_u16());
let qtype = QueryType::from_num(qtype_num);
let _ = try!(buffer.read_u16());
let ttl = try!(buffer.read_u32());
let data_len = try!(buffer.read_u16());
let qtype_num = buffer.read_u16()?;
let qtype = QueryType::from_num(qtype_num);
let _ = buffer.read_u16()?;
let ttl = buffer.read_u32()?;
let data_len = buffer.read_u16()?;
match qtype {
match qtype {
QueryType::A => {
let raw_addr = buffer.read_u32()?;
let addr = Ipv4Addr::new(
((raw_addr >> 24) & 0xFF) as u8,
((raw_addr >> 16) & 0xFF) as u8,
((raw_addr >> 8) & 0xFF) as u8,
((raw_addr >> 0) & 0xFF) as u8,
);
// Handle each record type separately, starting with the A record
// type which remains the same as before.
QueryType::A => {
let raw_addr = try!(buffer.read_u32());
let addr = Ipv4Addr::new(((raw_addr >> 24) & 0xFF) as u8,
((raw_addr >> 16) & 0xFF) as u8,
((raw_addr >> 8) & 0xFF) as u8,
((raw_addr >> 0) & 0xFF) as u8);
Ok(DnsRecord::A {
domain: domain,
addr: addr,
ttl: ttl,
})
}
QueryType::AAAA => {
let raw_addr1 = buffer.read_u32()?;
let raw_addr2 = buffer.read_u32()?;
let raw_addr3 = buffer.read_u32()?;
let raw_addr4 = buffer.read_u32()?;
let addr = Ipv6Addr::new(
((raw_addr1 >> 16) & 0xFFFF) as u16,
((raw_addr1 >> 0) & 0xFFFF) as u16,
((raw_addr2 >> 16) & 0xFFFF) as u16,
((raw_addr2 >> 0) & 0xFFFF) as u16,
((raw_addr3 >> 16) & 0xFFFF) as u16,
((raw_addr3 >> 0) & 0xFFFF) as u16,
((raw_addr4 >> 16) & 0xFFFF) as u16,
((raw_addr4 >> 0) & 0xFFFF) as u16,
);
Ok(DnsRecord::A {
domain: domain,
addr: addr,
ttl: ttl
})
},
Ok(DnsRecord::AAAA {
domain: domain,
addr: addr,
ttl: ttl,
})
}
QueryType::NS => {
let mut ns = String::new();
buffer.read_qname(&mut ns)?;
// The AAAA record type follows the same logic, but with more numbers to keep
// track off.
QueryType::AAAA => {
let raw_addr1 = try!(buffer.read_u32());
let raw_addr2 = try!(buffer.read_u32());
let raw_addr3 = try!(buffer.read_u32());
let raw_addr4 = try!(buffer.read_u32());
let addr = Ipv6Addr::new(((raw_addr1 >> 16) & 0xFFFF) as u16,
((raw_addr1 >> 0) & 0xFFFF) as u16,
((raw_addr2 >> 16) & 0xFFFF) as u16,
((raw_addr2 >> 0) & 0xFFFF) as u16,
((raw_addr3 >> 16) & 0xFFFF) as u16,
((raw_addr3 >> 0) & 0xFFFF) as u16,
((raw_addr4 >> 16) & 0xFFFF) as u16,
((raw_addr4 >> 0) & 0xFFFF) as u16);
Ok(DnsRecord::NS {
domain: domain,
host: ns,
ttl: ttl,
})
}
QueryType::CNAME => {
let mut cname = String::new();
buffer.read_qname(&mut cname)?;
Ok(DnsRecord::AAAA {
domain: domain,
addr: addr,
ttl: ttl
})
},
Ok(DnsRecord::CNAME {
domain: domain,
host: cname,
ttl: ttl,
})
}
QueryType::MX => {
let priority = buffer.read_u16()?;
let mut mx = String::new();
buffer.read_qname(&mut mx)?;
// NS and CNAME both have the same structure.
QueryType::NS => {
let mut ns = String::new();
try!(buffer.read_qname(&mut ns));
Ok(DnsRecord::MX {
domain: domain,
priority: priority,
host: mx,
ttl: ttl,
})
}
QueryType::UNKNOWN(_) => {
buffer.step(data_len as usize)?;
Ok(DnsRecord::NS {
domain: domain,
host: ns,
ttl: ttl
})
},
QueryType::CNAME => {
let mut cname = String::new();
try!(buffer.read_qname(&mut cname));
Ok(DnsRecord::CNAME {
domain: domain,
host: cname,
ttl: ttl
})
},
// MX is almost like the previous two, but with one extra field for priority.
QueryType::MX => {
let priority = try!(buffer.read_u16());
let mut mx = String::new();
try!(buffer.read_qname(&mut mx));
Ok(DnsRecord::MX {
domain: domain,
priority: priority,
host: mx,
ttl: ttl
})
},
// And we end with some code for handling unknown record types, as before.
QueryType::UNKNOWN(_) => {
try!(buffer.step(data_len as usize));
Ok(DnsRecord::UNKNOWN {
domain: domain,
qtype: qtype_num,
data_len: data_len,
ttl: ttl
})
Ok(DnsRecord::UNKNOWN {
domain: domain,
qtype: qtype_num,
data_len: data_len,
ttl: ttl,
})
}
}
}
- snip -
}
```
@@ -280,8 +275,8 @@ impl BytePacketBuffer {
}
fn set_u16(&mut self, pos: usize, val: u16) -> Result<()> {
try!(self.set(pos,(val >> 8) as u8));
try!(self.set(pos+1,(val & 0xFF) as u8));
self.set(pos, (val >> 8) as u8)?;
self.set(pos + 1, (val & 0xFF) as u8)?;
Ok(())
}
@@ -289,89 +284,119 @@ impl BytePacketBuffer {
}
```
When writing the labels of a record, we don't know ahead of time the number of
bytes needed, since we might end up using jumps to compress the size. We'll
solve this by writing a zero size and then going back to fill in the size
needed.
### Extending DnsRecord for writing new record types
Now we can amend `DnsRecord::write`. Here's our new function:
```rust
pub fn write(&self, buffer: &mut BytePacketBuffer) -> Result<usize> {
impl DnsRecord {
let start_pos = buffer.pos();
- snip -
match *self {
DnsRecord::A { ref domain, ref addr, ttl } => {
try!(buffer.write_qname(domain));
try!(buffer.write_u16(QueryType::A.to_num()));
try!(buffer.write_u16(1));
try!(buffer.write_u32(ttl));
try!(buffer.write_u16(4));
pub fn write(&self, buffer: &mut BytePacketBuffer) -> Result<usize> {
let start_pos = buffer.pos();
let octets = addr.octets();
try!(buffer.write_u8(octets[0]));
try!(buffer.write_u8(octets[1]));
try!(buffer.write_u8(octets[2]));
try!(buffer.write_u8(octets[3]));
},
DnsRecord::NS { ref domain, ref host, ttl } => {
try!(buffer.write_qname(domain));
try!(buffer.write_u16(QueryType::NS.to_num()));
try!(buffer.write_u16(1));
try!(buffer.write_u32(ttl));
match *self {
DnsRecord::A {
ref domain,
ref addr,
ttl,
} => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::A.to_num())?;
buffer.write_u16(1)?;
buffer.write_u32(ttl)?;
buffer.write_u16(4)?;
let pos = buffer.pos();
try!(buffer.write_u16(0));
try!(buffer.write_qname(host));
let size = buffer.pos() - (pos + 2);
try!(buffer.set_u16(pos, size as u16));
},
DnsRecord::CNAME { ref domain, ref host, ttl } => {
try!(buffer.write_qname(domain));
try!(buffer.write_u16(QueryType::CNAME.to_num()));
try!(buffer.write_u16(1));
try!(buffer.write_u32(ttl));
let pos = buffer.pos();
try!(buffer.write_u16(0));
try!(buffer.write_qname(host));
let size = buffer.pos() - (pos + 2);
try!(buffer.set_u16(pos, size as u16));
},
DnsRecord::MX { ref domain, priority, ref host, ttl } => {
try!(buffer.write_qname(domain));
try!(buffer.write_u16(QueryType::MX.to_num()));
try!(buffer.write_u16(1));
try!(buffer.write_u32(ttl));
let pos = buffer.pos();
try!(buffer.write_u16(0));
try!(buffer.write_u16(priority));
try!(buffer.write_qname(host));
let size = buffer.pos() - (pos + 2);
try!(buffer.set_u16(pos, size as u16));
},
DnsRecord::AAAA { ref domain, ref addr, ttl } => {
try!(buffer.write_qname(domain));
try!(buffer.write_u16(QueryType::AAAA.to_num()));
try!(buffer.write_u16(1));
try!(buffer.write_u32(ttl));
try!(buffer.write_u16(16));
for octet in &addr.segments() {
try!(buffer.write_u16(*octet));
let octets = addr.octets();
buffer.write_u8(octets[0])?;
buffer.write_u8(octets[1])?;
buffer.write_u8(octets[2])?;
buffer.write_u8(octets[3])?;
}
},
DnsRecord::UNKNOWN { .. } => {
println!("Skipping record: {:?}", self);
}
}
DnsRecord::NS {
ref domain,
ref host,
ttl,
} => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::NS.to_num())?;
buffer.write_u16(1)?;
buffer.write_u32(ttl)?;
Ok(buffer.pos() - start_pos)
let pos = buffer.pos();
buffer.write_u16(0)?;
buffer.write_qname(host)?;
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::CNAME {
ref domain,
ref host,
ttl,
} => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::CNAME.to_num())?;
buffer.write_u16(1)?;
buffer.write_u32(ttl)?;
let pos = buffer.pos();
buffer.write_u16(0)?;
buffer.write_qname(host)?;
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::MX {
ref domain,
priority,
ref host,
ttl,
} => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::MX.to_num())?;
buffer.write_u16(1)?;
buffer.write_u32(ttl)?;
let pos = buffer.pos();
buffer.write_u16(0)?;
buffer.write_u16(priority)?;
buffer.write_qname(host)?;
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::AAAA {
ref domain,
ref addr,
ttl,
} => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::AAAA.to_num())?;
buffer.write_u16(1)?;
buffer.write_u32(ttl)?;
buffer.write_u16(16)?;
for octet in &addr.segments() {
buffer.write_u16(*octet)?;
}
}
DnsRecord::UNKNOWN { .. } => {
println!("Skipping record: {:?}", self);
}
}
Ok(buffer.pos() - start_pos)
}
}
```