use crate::errors::DiameterResult;
use crate::modeling::avp::avp::{Avp, AvpFlags, AvpValue};
use crate::modeling::message::application_id::ApplicationId;
use crate::modeling::message::command_code::CommandCode;
use crate::modeling::message::command_flags::CommandFlag;
use crate::modeling::message::dictionary::Dictionary;
use std::io::{Read, Write};
use std::sync::Arc;
#[derive(Debug)]
pub struct DiameterMessage {
header: DiameterHeader,
avps: Vec<Avp>,
}
#[derive(Debug)]
pub struct DiameterHeader {
version: u8,
message_length: u32, command_flag: u8,
command_code: CommandCode, application_id: ApplicationId,
hop_by_hop: u32,
end_to_end: u32,
}
impl DiameterMessage {
pub fn new(
command_flag: CommandFlag,
command_code: CommandCode,
application_id: ApplicationId,
hop_by_hop: u32,
end_to_end: u32,
) -> Self {
Self {
header: DiameterHeader {
version: 1,
message_length: 20,
command_flag: command_flag.value(),
command_code,
application_id,
hop_by_hop,
end_to_end,
},
avps: vec![],
}
}
pub fn get_version(&self) -> u8 {
self.header.version
}
pub fn get_length(&self) -> u32 {
self.header.message_length
}
pub fn get_command_flag(&self) -> u8 {
self.header.command_flag
}
pub fn get_command_code(&self) -> &CommandCode {
&self.header.command_code
}
pub fn get_application_id(&self) -> &ApplicationId {
&self.header.application_id
}
pub fn get_hop_by_hop(&self) -> u32 {
self.header.hop_by_hop
}
pub fn get_end_to_end(&self) -> u32 {
self.header.end_to_end
}
pub fn get_avps(&self) -> &Vec<Avp> {
&self.avps
}
pub fn get_avp(&self, code: u32) -> Option<&Avp> {
self.avps.iter().find(|avp| avp.get_code() == code)
}
pub fn add(&mut self, avp: Avp) {
self.header.message_length += avp.get_length() + avp.get_padding();
self.avps.push(avp);
}
pub fn add_avp<T: Into<AvpValue>>(
&mut self,
code: u32,
flags: AvpFlags,
vendor_id: Option<u32>,
value: T,
) {
let avp: Avp = Avp::new(code, flags, vendor_id, value);
self.add(avp);
}
pub fn encode_to<W: Write>(&mut self, writer: &mut W) -> DiameterResult<()> {
writer.write(&self.header.version.to_be_bytes())?;
writer.write(&self.header.message_length.to_be_bytes()[1..])?;
writer.write(&self.header.command_flag.to_be_bytes())?;
writer.write(&self.header.command_code.get_code().to_be_bytes()[1..])?;
writer.write(&self.header.application_id.value().to_be_bytes())?;
writer.write(&self.header.hop_by_hop.to_be_bytes())?;
writer.write(&self.header.end_to_end.to_be_bytes())?;
for avp in self.avps.iter_mut() {
avp.encode_to(writer)?;
}
Ok(())
}
pub fn decode_from<R: Read>(
reader: &mut R,
dict: Arc<Dictionary>,
) -> DiameterResult<DiameterMessage> {
let mut b = [0u8; 20];
reader.read_exact(&mut b)?;
let version = b[0];
let mut message_length = u32::from_be_bytes([0, b[1], b[2], b[3]]);
let command_flag = b[4];
let command_code = u32::from_be_bytes([b[4], b[5], b[6], b[7]]);
let application_id = u32::from_be_bytes([b[8], b[9], b[10], b[11]]);
let hop_by_hop = u32::from_be_bytes([b[12], b[13], b[14], b[15]]);
let end_to_end = u32::from_be_bytes([b[16], b[17], b[18], b[19]]);
let header = DiameterHeader {
version,
message_length,
command_flag,
application_id: ApplicationId::try_from(application_id)?,
command_code: CommandCode::try_from(command_code)?,
hop_by_hop,
end_to_end,
};
let mut message = DiameterMessage {
header,
avps: vec![],
};
message_length -= 20;
while message_length > 0 {
let avp = Avp::decode_from(reader, Arc::clone(&dict))?;
message_length = message_length - avp.get_length() - avp.get_padding();
message.add(avp);
}
Ok(message)
}
}