ads_proto/proto/
state_flags.rs1use 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 pub fn resp_default() -> Self {
36 StateFlags::new(true, true, NetProto::Tcp)
37 }
38
39 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
61impl 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}