Skip to content

Commit

Permalink
build: more intuitive EDNS API
Browse files Browse the repository at this point in the history
  • Loading branch information
@LEXUGE
LEXUGE committed Jun 23, 2022
1 parent 71269f3 commit b4d1c21
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 30 deletions.
13 changes: 5 additions & 8 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,20 +61,17 @@ script:
}
```
And another script that removes and adds EDNS Client Subnet record into the OPT pseudo-section:
And another script that adds EDNS Client Subnet record into the OPT pseudo-section:
```yaml
script:
route: |
let query = query;
// Remove existing OPT section
for (record, counter) in query.additional {
if record.rtype == "OPT" {
query.remove_additional(counter);
}
}
query.push_opt(create_opt_section(), create_client_subnet(15, 0, "23.62.93.233"));
// Optionally remove all the existing OPT pseudo-section(s)
// query.clear_opt();
query.push_opt(create_client_subnet(15, 0, "23.62.93.233"));
upstreams.send("secure", query)
```
Expand Down
54 changes: 45 additions & 9 deletions droute/src/router/script/message/helpers.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
use super::{MessageError, MessageResult as Result};
use bytes::{Bytes, BytesMut};
use domain::{
base::{opt::AllOptData, Dname, Message, MessageBuilder, ParsedDname, Record, ToDname},
base::{opt::AllOptData, Dname, Message, MessageBuilder, ParsedDname, Record, Rtype, ToDname},
rdata::{
AllRecordData, Cname, Dname as DnameRecord, Mb, Md, Mf, Minfo, Mr, Mx, Ns, Nsec, Ptr,
Rrsig, Soa, Srv, Tsig,
Expand Down Expand Up @@ -128,7 +128,13 @@ impl IntoIterator for OptRecordsIter {
}
}

pub fn modify_opt(msg: &Message<Bytes>, opt: OptRecordsIter) -> Result<Message<Bytes>> {
impl OptRecordsIter {
pub fn iter(&self) -> std::slice::Iter<'_, AllOptData<Bytes>> {
self.0.iter()
}
}

pub fn modify_opt(msg: &Message<Bytes>, opt: Option<OptRecordsIter>) -> Result<Message<Bytes>> {
let mut builder = MessageBuilder::from_target(BytesMut::with_capacity(crate::MAX_LEN))?;
// Copy header
*builder.header_mut() = msg.header();
Expand All @@ -154,21 +160,51 @@ pub fn modify_opt(msg: &Message<Bytes>, opt: OptRecordsIter) -> Result<Message<B
}
}

// Per RFC 6891, there can only be one OPT pseudo-section within the additional section.
// Therefore, if we got a malformatted message with multiple OPT records, we only keep one of them.

// whether we have already replaced one OPT record
let mut flag = false;

// Copy other additional records
let mut builder = builder.additional();
for item in msg.additional()? {
if let Some(record) = item?.into_record::<AllRecordData<_, _>>()? {
builder.push(record)?;
match (record.rtype(), flag) {
// First time seeing an OPT record, replace it with what we build
(Rtype::Opt, false) => {
if let Some(ref opt) = opt {
// Build OPT record
builder.opt(|builder| {
for option in opt.iter() {
builder.push(option)?
}
Ok(())
})?;
}

flag = true;
}
// Multiple OPT record, do nothing.
(Rtype::Opt, true) => {}
// Other records, copy as usual
_ => builder.push(record)?,
}
}
}

// Build OPT record
builder.opt(|builder| {
for option in opt.into_iter() {
builder.push(&option)?
// If the original message doesn't contain any OPT record, we create them based on our needs
if !flag {
if let Some(ref opt) = opt {
// Build OPT record
builder.opt(|builder| {
for option in opt.iter() {
builder.push(option)?
}
Ok(())
})?;
}
Ok(())
})?;
}

Ok(builder.into_message())
}
Expand Down
76 changes: 63 additions & 13 deletions droute/src/router/script/message/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -494,10 +494,17 @@ pub mod rhai_mod {
}

pub mod getters {
use super::super::helpers::{dns_record_from_ref, DnsRecordsIter};
use crate::{IntoEvalAltResultError, IntoEvalAltResultStr};
use self::record::rdata::opt::get_options;

use super::{
super::helpers::{dns_record_from_ref, DnsRecordsIter},
convertions::to_opt,
};
use crate::{
router::script::message::OptRecordsIter, IntoEvalAltResultError, IntoEvalAltResultStr,
};
use bytes::Bytes;
use domain::base::{Dname, Header, Message, Question, ToDname};
use domain::base::{opt::OptRecord, Dname, Header, Message, Question, Rtype, ToDname};
use rhai::EvalAltResult;

#[rhai_fn(get = "header", pure)]
Expand Down Expand Up @@ -599,6 +606,30 @@ pub mod rhai_mod {
}
}

#[rhai_fn(get = "opt_section", pure, return_raw)]
pub fn get_opt_section(
msg: &mut Message<Bytes>,
) -> Result<Option<OptRecord<Bytes>>, Box<EvalAltResult>> {
for mut record in get_additional(msg)?.into_iter() {
if record.rtype() == Rtype::Opt {
return Ok(Some(to_opt(&mut record)?));
}
}
Ok(None)
}

#[rhai_fn(get = "options", pure, return_raw)]
pub fn get_options_msg(
msg: &mut Message<Bytes>,
) -> Result<OptRecordsIter, Box<EvalAltResult>> {
for mut record in get_additional(msg)?.into_iter() {
if record.rtype() == Rtype::Opt {
return get_options(&mut to_opt(&mut record)?);
}
}
Ok(OptRecordsIter(Vec::new()))
}

#[rhai_fn(get = "answer", pure, return_raw)]
pub fn get_answer(msg: &mut Message<Bytes>) -> Result<DnsRecordsIter, Box<EvalAltResult>> {
create_record_iter_impl!(answer, msg)
Expand Down Expand Up @@ -976,54 +1007,58 @@ pub mod rhai_mod {

pub mod opt {
use crate::{
router::script::message::helpers::{modify_opt, OptRecordsIter},
router::script::message::{
helpers::{modify_opt, OptRecordsIter},
rhai_mod::getters::get_options_msg,
},
IntoEvalAltResultError,
};
use bytes::Bytes;
use domain::base::{opt::AllOptData, Message};
use rhai::EvalAltResult;

// Create an empty OPT section, this doesn't alter the message
pub fn create_opt_section() -> OptRecordsIter {
OptRecordsIter(Vec::new())
}

#[rhai_fn(return_raw)]
pub fn update_opt(
msg: &mut Message<Bytes>,
opt: OptRecordsIter,
) -> Result<(), Box<EvalAltResult>> {
*msg = modify_opt(msg, opt).into_evalrst_err()?;
*msg = modify_opt(msg, Some(opt)).into_evalrst_err()?;
Ok(())
}

#[rhai_fn(return_raw)]
pub fn clear_opt(msg: &mut Message<Bytes>) -> Result<(), Box<EvalAltResult>> {
*msg = modify_opt(msg, None).into_evalrst_err()?;
Ok(())
}

#[rhai_fn(return_raw)]
pub fn push_opt(
msg: &mut Message<Bytes>,
mut opt: OptRecordsIter,
data: AllOptData<Bytes>,
) -> Result<(), Box<EvalAltResult>> {
let mut opt = get_options_msg(msg)?;
opt.0.push(data);
update_opt(msg, opt)
}

#[rhai_fn(return_raw)]
pub fn insert_opt(
msg: &mut Message<Bytes>,
mut opt: OptRecordsIter,
index: i32,
data: AllOptData<Bytes>,
) -> Result<(), Box<EvalAltResult>> {
let mut opt = get_options_msg(msg)?;
opt.0.insert(index as usize, data);
update_opt(msg, opt)
}

#[rhai_fn(return_raw)]
pub fn remove_opt(
msg: &mut Message<Bytes>,
mut opt: OptRecordsIter,
index: i32,
) -> Result<(), Box<EvalAltResult>> {
let mut opt = get_options_msg(msg)?;
opt.0.remove(index as usize);
update_opt(msg, opt)
}
Expand Down Expand Up @@ -1090,6 +1125,11 @@ pub mod rhai_mod {
Ok(())
}

#[rhai_fn(return_raw)]
pub fn clear_answer(msg: &mut Message<Bytes>) -> Result<(), Box<EvalAltResult>> {
update_answer(msg, DnsRecordsIter(Vec::new()))
}

#[rhai_fn(return_raw)]
pub fn insert_answer(
msg: &mut Message<Bytes>,
Expand Down Expand Up @@ -1145,6 +1185,11 @@ pub mod rhai_mod {
Ok(())
}

#[rhai_fn(return_raw)]
pub fn clear_authority(msg: &mut Message<Bytes>) -> Result<(), Box<EvalAltResult>> {
update_authority(msg, DnsRecordsIter(Vec::new()))
}

#[rhai_fn(return_raw)]
pub fn insert_authority(
msg: &mut Message<Bytes>,
Expand Down Expand Up @@ -1200,6 +1245,11 @@ pub mod rhai_mod {
Ok(())
}

#[rhai_fn(return_raw)]
pub fn clear_additional(msg: &mut Message<Bytes>) -> Result<(), Box<EvalAltResult>> {
update_additional(msg, DnsRecordsIter(Vec::new()))
}

#[rhai_fn(return_raw)]
pub fn insert_additional(
msg: &mut Message<Bytes>,
Expand Down

0 comments on commit b4d1c21

Please sign in to comment.