ads_proto/proto/
state_flags.rs

1use crate::proto::proto_traits::{ReadFrom, WriteTo};
2use bitfield::Bit;
3use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
4use std::io::{self, Read, Write};
5
6#[derive(Debug)]
7pub enum NetProto {
8    Tcp,
9    Udp,
10}
11
12#[derive(Debug, Clone)]
13pub struct StateFlags {
14    value: u16,
15}
16
17impl StateFlags {
18    pub fn new(response: bool, ads_command: bool, net_proto: NetProto) -> Self {
19        let mut state_flags: u16 = 0;
20        state_flags.set_bit(0, response);
21        state_flags.set_bit(2, ads_command);
22
23        state_flags.set_bit(
24            6,
25            match net_proto {
26                NetProto::Tcp => false,
27                NetProto::Udp => true,
28            },
29        );
30
31        StateFlags { value: state_flags }
32    }
33
34    ///default for response (response=true, ads_command=true, net_proto=Tcp)
35    pub fn resp_default() -> Self {
36        StateFlags::new(true, true, NetProto::Tcp)
37    }
38
39    ///default for response (response=false, ads_command=true, net_proto=Tcp)
40    pub fn req_default() -> Self {
41        StateFlags::new(false, true, NetProto::Tcp)
42    }
43
44    pub fn value(&self) -> u16 {
45        self.value
46    }
47
48    pub fn is_tcp(&self) -> bool {
49        !self.value.bit(6)
50    }
51
52    pub fn is_response(&self) -> bool {
53        self.value.bit(0)
54    }
55
56    pub fn is_ads_command(&self) -> bool {
57        self.value.bit(2)
58    }
59}
60
61///from u16 may results in an invalid state flag
62impl From<u16> for StateFlags {
63    fn from(value: u16) -> StateFlags {
64        StateFlags { value }
65    }
66}
67
68impl WriteTo for StateFlags {
69    fn write_to<W: Write>(&self, mut wtr: W) -> io::Result<()> {
70        wtr.write_u16::<LittleEndian>(self.value())?;
71        Ok(())
72    }
73}
74
75impl ReadFrom for StateFlags {
76    fn read_from<R: Read>(read: &mut R) -> io::Result<Self> {
77        Ok(StateFlags::from(read.read_u16::<LittleEndian>()?))
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    #[test]
86    fn new_test() {
87        let state_flags = StateFlags::new(false, true, NetProto::Tcp);
88        assert_eq!(state_flags.value, 4);
89
90        let state_flags = StateFlags::new(true, true, NetProto::Tcp);
91        assert_eq!(state_flags.value, 5);
92
93        let state_flags = StateFlags::new(false, false, NetProto::Tcp);
94        assert_eq!(state_flags.value, 0);
95
96        let state_flags = StateFlags::new(true, false, NetProto::Tcp);
97        assert_eq!(state_flags.value, 1);
98
99        let state_flags = StateFlags::new(false, true, NetProto::Udp);
100        assert_eq!(state_flags.value, 68);
101
102        let state_flags = StateFlags::new(true, true, NetProto::Udp);
103        assert_eq!(state_flags.value, 69);
104
105        let state_flags = StateFlags::new(false, false, NetProto::Udp);
106        assert_eq!(state_flags.value, 64);
107
108        let state_flags = StateFlags::new(true, false, NetProto::Udp);
109        assert_eq!(state_flags.value, 65);
110    }
111
112    #[test]
113    fn resp_default_test() {
114        let state_flags = StateFlags::resp_default();
115        assert_eq!(state_flags.value, 5);
116    }
117
118    #[test]
119    fn req_default_test() {
120        let state_flags = StateFlags::req_default();
121        assert_eq!(state_flags.value, 4);
122    }
123
124    #[test]
125    fn write_to_test() {
126        let state_flags = StateFlags::resp_default();
127        let mut buffer: Vec<u8> = Vec::with_capacity(16);
128        state_flags.write_to(&mut buffer).unwrap();
129        assert_eq!(buffer, [5, 0]);
130
131        let state_flags = StateFlags::new(true, true, NetProto::Udp);
132        let mut buffer: Vec<u8> = Vec::with_capacity(16);
133        state_flags.write_to(&mut buffer).unwrap();
134        assert_eq!(buffer, [69, 0]);
135    }
136
137    #[test]
138    fn read_from_test() {
139        let buffer: Vec<u8> = vec![5, 0, 1, 99];
140        let state_flags = StateFlags::read_from(&mut buffer.as_slice()).unwrap();
141
142        assert_eq!(state_flags.value(), 5);
143    }
144
145    #[test]
146    fn is_tcp_test() {
147        let buffer: Vec<u8> = vec![5, 0, 1, 99, 4];
148        let state_flags = StateFlags::read_from(&mut buffer.as_slice()).unwrap();
149
150        assert_eq!(state_flags.is_tcp(), true);
151    }
152
153    #[test]
154    fn is_udp_test() {
155        let buffer: Vec<u8> = vec![69, 0, 1, 99, 4];
156        let state_flags = StateFlags::read_from(&mut buffer.as_slice()).unwrap();
157
158        assert_eq!(state_flags.is_tcp(), false);
159    }
160
161    #[test]
162    fn is_response_test() {
163        let buffer: Vec<u8> = vec![5, 0, 1, 99, 4];
164        let state_flags = StateFlags::read_from(&mut buffer.as_slice()).unwrap();
165
166        assert_eq!(state_flags.is_response(), true);
167    }
168
169    #[test]
170    fn is_ads_command_test() {
171        let buffer: Vec<u8> = vec![4, 0, 1, 99, 4];
172        let state_flags = StateFlags::read_from(&mut buffer.as_slice()).unwrap();
173
174        assert_eq!(state_flags.is_ads_command(), true);
175        assert_eq!(state_flags.is_response(), false);
176    }
177}