abol_core/
attribute.rs

1use std::{
2    net::{Ipv4Addr, Ipv6Addr},
3    time::{Duration, SystemTime, UNIX_EPOCH},
4};
5
6use bytes::{BufMut, Bytes, BytesMut};
7pub type AttributeType = u8;
8pub type AttributeValue = Bytes;
9
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct Avp {
12    pub attribute_type: AttributeType,
13    pub value: AttributeValue,
14}
15
16#[derive(Debug, Clone, PartialEq, Eq, Default)]
17pub struct Attributes(pub Vec<Avp>);
18
19#[derive(Debug, thiserror::Error, PartialEq, Eq)]
20pub enum AttributeParseError {
21    #[error("Buffer too short for attribute header")]
22    ShortBuffer,
23    #[error("Invalid attribute length: {0}")]
24    InvalidLength(u8),
25}
26
27pub fn parse_attributes(mut b: Bytes) -> Result<Attributes, AttributeParseError> {
28    let mut attrs = Vec::new();
29    while !b.is_empty() {
30        if b.len() < 2 {
31            return Err(AttributeParseError::ShortBuffer);
32        }
33        let attribute_type = b[0];
34        let length = b[1] as usize;
35        if length < 2 || length > b.len() {
36            return Err(AttributeParseError::InvalidLength(b[1]));
37        }
38        let mut attr_block = b.split_to(length);
39        let value = attr_block.split_off(2); // Remove type/length bytes
40
41        attrs.push(Avp {
42            attribute_type,
43            value,
44        });
45    }
46    Ok(Attributes(attrs))
47}
48
49impl Attributes {
50    pub fn add(&mut self, key: AttributeType, value: AttributeValue) {
51        self.0.push(Avp {
52            attribute_type: key,
53            value,
54        });
55    }
56
57    pub fn get(&self, key: AttributeType) -> Option<&AttributeValue> {
58        self.0
59            .iter()
60            .find(|avp| avp.attribute_type == key)
61            .map(|avp| &avp.value)
62    }
63
64    pub fn del(&mut self, key: AttributeType) {
65        self.0.retain(|avp| avp.attribute_type != key);
66    }
67
68    pub fn set(&mut self, key: AttributeType, value: AttributeValue) {
69        self.del(key);
70        self.add(key, value);
71    }
72    pub fn encode_to(&self, buf: &mut [u8]) -> Result<(), AttributeParseError> {
73        let mut offset: usize = 0;
74        for attr in &self.0 {
75            let value_len = attr.value.len();
76            if value_len > 253 {
77                return Err(AttributeParseError::InvalidLength((value_len + 2) as u8));
78            }
79            let avp_size = 2 + value_len;
80            let next_offset = offset + avp_size;
81            if next_offset > buf.len() {
82                return Err(AttributeParseError::ShortBuffer);
83            }
84            let avp_slice = &mut buf[offset..next_offset];
85            avp_slice[0] = attr.attribute_type;
86            avp_slice[1] = avp_size as u8;
87            avp_slice[2..].copy_from_slice(&attr.value);
88            offset = next_offset;
89        }
90        Ok(())
91    }
92
93    pub fn encoded_len(&self) -> Result<usize, AttributeParseError> {
94        let mut total_len = 0;
95        for attr in &self.0 {
96            let value_len = attr.value.len();
97            if value_len > 253 {
98                return Err(AttributeParseError::InvalidLength((value_len + 2) as u8));
99            }
100            total_len += 2 + value_len;
101        }
102        Ok(total_len)
103    }
104
105    pub fn get_vsa_attribute(&self, vendor_id: u32, vendor_type: u8) -> Option<&[u8]> {
106        for avp in &self.0 {
107            if avp.attribute_type == 26 && avp.value.len() >= 6 {
108                let v_id =
109                    u32::from_be_bytes([avp.value[0], avp.value[1], avp.value[2], avp.value[3]]);
110                let v_type = avp.value[4];
111
112                if v_id == vendor_id && v_type == vendor_type {
113                    return Some(&avp.value[6..]);
114                }
115            }
116        }
117        None
118    }
119    pub fn set_vsa_attribute(&mut self, vendor_id: u32, vendor_type: u8, value: Bytes) {
120        self.0.retain(|avp| {
121            if avp.attribute_type != 26 || avp.value.len() < 6 {
122                return true;
123            }
124
125            let v_id = u32::from_be_bytes([avp.value[0], avp.value[1], avp.value[2], avp.value[3]]);
126            let v_type = avp.value[4];
127
128            !(v_id == vendor_id && v_type == vendor_type)
129        });
130
131        let mut vsa_value = BytesMut::with_capacity(6 + value.len());
132
133        vsa_value.put_u32(vendor_id);
134        vsa_value.put_u8(vendor_type);
135        vsa_value.put_u8((2 + value.len()) as u8); // 2 is for Type + Length headers
136        vsa_value.put(value);
137
138        self.0.push(Avp {
139            attribute_type: 26,
140            value: vsa_value.freeze(), // Converts BytesMut to Bytes (O(1))
141        });
142    }
143}
144pub trait FromRadiusAttribute: Sized {
145    fn from_bytes(bytes: &[u8]) -> Option<Self>;
146}
147
148pub trait ToRadiusAttribute {
149    fn to_bytes(&self) -> Vec<u8>;
150}
151
152impl FromRadiusAttribute for u16 {
153    fn from_bytes(bytes: &[u8]) -> Option<Self> {
154        let b: [u8; 2] = bytes.try_into().ok()?;
155        Some(u16::from_be_bytes(b))
156    }
157}
158
159impl ToRadiusAttribute for u16 {
160    fn to_bytes(&self) -> Vec<u8> {
161        self.to_be_bytes().to_vec()
162    }
163}
164
165impl FromRadiusAttribute for u32 {
166    fn from_bytes(bytes: &[u8]) -> Option<Self> {
167        let b: [u8; 4] = bytes.try_into().ok()?;
168        Some(u32::from_be_bytes(b))
169    }
170}
171
172impl ToRadiusAttribute for u32 {
173    fn to_bytes(&self) -> Vec<u8> {
174        self.to_be_bytes().to_vec()
175    }
176}
177
178impl FromRadiusAttribute for u64 {
179    fn from_bytes(bytes: &[u8]) -> Option<Self> {
180        let b: [u8; 8] = bytes.try_into().ok()?;
181        Some(u64::from_be_bytes(b))
182    }
183}
184
185impl ToRadiusAttribute for u64 {
186    fn to_bytes(&self) -> Vec<u8> {
187        self.to_be_bytes().to_vec()
188    }
189}
190
191impl FromRadiusAttribute for Ipv4Addr {
192    fn from_bytes(bytes: &[u8]) -> Option<Self> {
193        let b: [u8; 4] = bytes.try_into().ok()?;
194        Some(Ipv4Addr::from(b))
195    }
196}
197
198impl ToRadiusAttribute for Ipv4Addr {
199    fn to_bytes(&self) -> Vec<u8> {
200        self.octets().to_vec()
201    }
202}
203
204impl FromRadiusAttribute for Ipv6Addr {
205    fn from_bytes(bytes: &[u8]) -> Option<Self> {
206        let b: [u8; 16] = bytes.try_into().ok()?;
207        Some(Ipv6Addr::from(b))
208    }
209}
210
211impl ToRadiusAttribute for Ipv6Addr {
212    fn to_bytes(&self) -> Vec<u8> {
213        self.octets().to_vec()
214    }
215}
216
217impl FromRadiusAttribute for String {
218    fn from_bytes(bytes: &[u8]) -> Option<Self> {
219        String::from_utf8(bytes.to_vec()).ok()
220    }
221}
222
223impl ToRadiusAttribute for String {
224    fn to_bytes(&self) -> Vec<u8> {
225        self.as_bytes().to_vec()
226    }
227}
228
229impl FromRadiusAttribute for Vec<u8> {
230    fn from_bytes(bytes: &[u8]) -> Option<Self> {
231        Some(bytes.to_vec())
232    }
233}
234
235impl ToRadiusAttribute for Vec<u8> {
236    fn to_bytes(&self) -> Vec<u8> {
237        self.clone()
238    }
239}
240
241impl FromRadiusAttribute for SystemTime {
242    fn from_bytes(bytes: &[u8]) -> Option<Self> {
243        let secs = u32::from_bytes(bytes)?;
244        Some(UNIX_EPOCH + Duration::from_secs(secs as u64))
245    }
246}
247
248impl ToRadiusAttribute for SystemTime {
249    fn to_bytes(&self) -> Vec<u8> {
250        let duration = self.duration_since(UNIX_EPOCH).unwrap_or_default();
251        (duration.as_secs() as u32).to_bytes()
252    }
253}
254
255/// Helper for Type-Length-Value (TLV) attributes
256pub struct Tlv {
257    pub tlv_type: u8,
258    pub value: Vec<u8>,
259}
260
261impl FromRadiusAttribute for Tlv {
262    fn from_bytes(bytes: &[u8]) -> Option<Self> {
263        if bytes.len() < 2 || bytes[1] as usize != bytes.len() {
264            return None;
265        }
266        Some(Tlv {
267            tlv_type: bytes[0],
268            value: bytes[2..].to_vec(),
269        })
270    }
271}
272
273impl ToRadiusAttribute for Tlv {
274    fn to_bytes(&self) -> Vec<u8> {
275        let mut v = Vec::with_capacity(2 + self.value.len());
276        v.push(self.tlv_type);
277        v.push((2 + self.value.len()) as u8);
278        v.extend_from_slice(&self.value);
279        v
280    }
281}