From 7274028dfc8084c2ead804e7998177c2a847bd20 Mon Sep 17 00:00:00 2001 From: joalopez Date: Sun, 29 Dec 2024 12:15:23 -0300 Subject: [PATCH] added error handling and tests --- src/message/rdata/opt_rdata.rs | 130 ++++++++++++++++++++++++++++++++- 1 file changed, 129 insertions(+), 1 deletion(-) diff --git a/src/message/rdata/opt_rdata.rs b/src/message/rdata/opt_rdata.rs index b9a57304..edef4e1d 100644 --- a/src/message/rdata/opt_rdata.rs +++ b/src/message/rdata/opt_rdata.rs @@ -77,6 +77,12 @@ impl FromBytes> for OptRdata { let option_data = OptionData::from_with_opt_type(option_data, option_code); + if let Err(_) = option_data { + return Err("Format Error"); + } + + let option_data = option_data.unwrap(); + let option = OptOption::new(option_code, option_length, option_data); opt_rdata.option.push(option); @@ -116,7 +122,7 @@ impl fmt::Display for OptRdata { for option in &self.option { let option_code = option.get_option_code(); let option_length = option.get_option_len(); - let option_data = option.get_opt_data().to_bytes(); + let option_data = option.get_opt_data(); result.push_str(&format!("OPTION-CODE: {}\n", option_code)); result.push_str(&format!("OPTION-LENGTH: {}\n", option_length)); result.push_str(&format!("OPTION-DATA: {:?} \n", option_data)); @@ -132,6 +138,10 @@ impl fmt::Display for OptRdata { #[cfg(test)] mod opt_rdata_test{ use crate::message::rdata::opt_rdata::option_data::OptionData; + use crate::message::rdata::opt_rdata::option_code::OptionCode; + use crate::message::rdata::opt_rdata::ede_optdata::EdeOptData; + use crate::message::rdata::opt_rdata::ede_code::EdeCode; + use super::*; #[test] @@ -184,4 +194,122 @@ mod opt_rdata_test{ assert_eq!(opt_rdata.get_option(), option); } + + #[test] + fn test_opt_rdata_ede_stale_answer() { + let mut opt_rdata = OptRdata::new(); + + // Create EDE data + let code = EdeCode::StaleAns; + let msg = "Stale Answer".to_string(); + let ede_data = EdeOptData::new(code, msg.clone()); + + // Wrap it in OptionData::EDE + let option_data = OptionData::EDE(ede_data); + // Determine the length for the OPT option + let option_data_bytes = option_data.to_bytes(); + let option_len = option_data_bytes.len() as u16; + + // Build an OptOption with OptionCode::EDE + let option = OptOption::new(OptionCode::EDE, option_len, option_data); + opt_rdata.option.push(option); + + // Round-trip: to_bytes -> from_bytes + let serialized = opt_rdata.to_bytes(); + let deserialized = OptRdata::from_bytes(&serialized, &serialized).unwrap(); + + // Check we got one option + let retrieved_options = deserialized.get_option(); + assert_eq!(retrieved_options.len(), 1); + + let retrieved_option = &retrieved_options[0]; + assert_eq!(retrieved_option.get_option_code(), OptionCode::EDE); + assert_eq!(retrieved_option.get_option_len(), option_len); + + // Now confirm the EDE contents + match retrieved_option.get_opt_data() { + OptionData::EDE(ede) => { + assert_eq!(ede.get_err_code(), EdeCode::StaleAns); + assert_eq!(ede.get_err_message(), msg); + } + _ => panic!("Expected OptionData::EDE, got something else!"), + } + } + + /// 2) Test an EDE "DNSSEC Bogus" code + #[test] + fn test_opt_rdata_ede_dnssec_bogus() { + let mut opt_rdata = OptRdata::new(); + + let code = EdeCode::DnssecBogus; + let msg = "DNSSEC Bogus".to_string(); + let ede_data = EdeOptData::new(code, msg.clone()); + + let option_data = OptionData::EDE(ede_data); + let option_data_bytes = option_data.to_bytes(); + let option_len = option_data_bytes.len() as u16; + + let option = OptOption::new(OptionCode::EDE, option_len, option_data); + opt_rdata.option.push(option); + + let serialized = opt_rdata.to_bytes(); + let deserialized = OptRdata::from_bytes(&serialized, &serialized).unwrap(); + + let retrieved_options = deserialized.get_option(); + assert_eq!(retrieved_options.len(), 1); + + let retrieved_option = &retrieved_options[0]; + assert_eq!(retrieved_option.get_option_code(), OptionCode::EDE); + assert_eq!(retrieved_option.get_option_len(), option_len); + + match retrieved_option.get_opt_data() { + OptionData::EDE(ede) => { + assert_eq!(ede.get_err_code(), EdeCode::DnssecBogus); + assert_eq!(ede.get_err_message(), msg); + } + _ => panic!("Expected OptionData::EDE, got something else!"), + } + } + + /// 3) Test an EDE with an Unknown code, plus demonstrate using setters + #[test] + fn test_opt_rdata_ede_unknown() { + let mut opt_rdata = OptRdata::new(); + + // Start with an unknown code + let code = EdeCode::Unknown(999); + let msg = "Some unknown EDE error".to_string(); + let mut ede_data = EdeOptData::new(code, msg); + + // Use your EdeOptData's setters (if you have them) + ede_data.set_err_code(EdeCode::Unknown(1000)); + ede_data.set_err_message("Modified unknown EDE".to_string()); + + let option_data = OptionData::EDE(ede_data); + let option_data_bytes = option_data.to_bytes(); + let option_len = option_data_bytes.len() as u16; + + let option = OptOption::new(OptionCode::EDE, option_len, option_data); + opt_rdata.option.push(option); + + // Serialize + let serialized = opt_rdata.to_bytes(); + // Deserialize + let deserialized = OptRdata::from_bytes(&serialized, &serialized).unwrap(); + + let retrieved_options = deserialized.get_option(); + assert_eq!(retrieved_options.len(), 1); + + let retrieved_option = &retrieved_options[0]; + assert_eq!(retrieved_option.get_option_code(), OptionCode::EDE); + assert_eq!(retrieved_option.get_option_len(), option_len); + + match retrieved_option.get_opt_data() { + OptionData::EDE(ede) => { + assert_eq!(ede.get_err_code(), EdeCode::Unknown(1000)); + assert_eq!(ede.get_err_message(), "Modified unknown EDE"); + } + _ => panic!("Expected OptionData::EDE, got something else!"), + } + } } \ No newline at end of file