1use crate::error::{Error, ExtendBufferTooShortError, Result};
6use std::io::Write;
7
8pub trait ProtocolMessage: Sized {
10 const MIN_SIZE: usize;
12 fn write_bytes<W: Write>(&self, out: &mut W) -> Result<usize>;
14 fn byte_size(&self) -> usize;
17 fn try_from_bytes(bytes: &[u8]) -> Result<(Self, usize)>;
21}
22
23#[derive(Debug, Clone, PartialEq, Eq, Hash)]
31pub struct EncodableString(String);
32
33impl ProtocolMessage for EncodableString {
34 const MIN_SIZE: usize = 1;
35 fn write_bytes<W: Write>(&self, out: &mut W) -> Result<usize> {
36 let len: u8 =
37 self.0.as_bytes().len().try_into().expect("EncodableString wasn't encodable!");
38 out.write_all(&[len])?;
39 out.write_all(self.0.as_bytes())?;
40 Ok(usize::from(len) + 1)
41 }
42
43 fn byte_size(&self) -> usize {
44 self.0.as_bytes().len() + 1
45 }
46
47 fn try_from_bytes(bytes: &[u8]) -> Result<(Self, usize)> {
48 if bytes.is_empty() {
49 Err(Error::BufferTooShort(1))
50 } else if bytes.len() - 1 < bytes[0] as usize {
51 Err(Error::BufferTooShort(bytes[0] as usize + 1))
52 } else {
53 let len = bytes[0] as usize;
54 let bytes = &bytes[1..][..len];
55 Ok((
56 std::str::from_utf8(bytes)
57 .map_err(|_| Error::BadUTF8(String::from_utf8_lossy(bytes).to_string()))?
58 .to_owned()
59 .try_into()
60 .expect("String wasn't decodable right after encoding!"),
61 len + 1,
62 ))
63 }
64 }
65}
66
67impl TryFrom<String> for EncodableString {
68 type Error = Error;
69 fn try_from(src: String) -> Result<EncodableString> {
70 let _: u8 =
71 src.as_bytes().len().try_into().map_err(|_| Error::StringTooBig(src.clone()))?;
72 Ok(EncodableString(src))
73 }
74}
75
76impl std::ops::Deref for EncodableString {
77 type Target = String;
78 fn deref(&self) -> &String {
79 &self.0
80 }
81}
82
83impl<T> PartialEq<T> for EncodableString
84where
85 String: PartialEq<T>,
86{
87 fn eq(&self, other: &T) -> bool {
88 PartialEq::eq(&self.0, other)
89 }
90}
91
92impl std::fmt::Display for EncodableString {
93 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94 std::fmt::Display::fmt(&self.0, f)
95 }
96}
97
98#[derive(Debug)]
102pub struct Identify {
103 pub circuit_version: u8,
104 pub protocol: EncodableString,
105}
106
107impl Identify {
108 pub fn new(protocol: EncodableString) -> Self {
110 Identify { circuit_version: crate::CIRCUIT_VERSION, protocol }
111 }
112}
113
114impl ProtocolMessage for Identify {
115 const MIN_SIZE: usize = 1 + EncodableString::MIN_SIZE;
116 fn byte_size(&self) -> usize {
117 self.protocol.byte_size() + 1
118 }
119
120 fn write_bytes<W: Write>(&self, out: &mut W) -> Result<usize> {
121 let mut bytes = 0;
122 out.write_all(&[self.circuit_version])?;
123 bytes += 1;
124 bytes += self.protocol.write_bytes(out)?;
125 Ok(bytes)
126 }
127
128 fn try_from_bytes(bytes: &[u8]) -> Result<(Self, usize)> {
129 if bytes.len() < 2 {
130 return Err(Error::BufferTooShort(2));
131 }
132
133 let circuit_version = bytes[0];
134 let (protocol, proto_len) =
135 EncodableString::try_from_bytes(&bytes[1..]).extend_buffer_too_short(1)?;
136
137 Ok((Identify { circuit_version, protocol }, 1 + proto_len))
138 }
139}
140
141#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
145pub struct Quality(u8);
146
147impl Quality {
148 pub const SELF: Quality = Quality(0);
150 pub const IN_PROCESS: Quality = Quality(1);
152 pub const LOCAL_SOCKET: Quality = Quality(2);
154 pub const USB: Quality = Quality(5);
156 pub const NETWORK: Quality = Quality(20);
158 pub const WORST: Quality = Quality(u8::MAX - 1);
160 pub const UNKNOWN: Quality = Quality::WORST;
162
163 pub fn combine(self, other: Self) -> Self {
166 Quality(std::cmp::min(self.0.saturating_add(other.0), u8::MAX - 1))
167 }
168}
169
170impl TryFrom<u8> for Quality {
171 type Error = ();
172 fn try_from(value: u8) -> Result<Self, Self::Error> {
173 if value != u8::MAX {
174 Ok(Quality(value))
175 } else {
176 Err(())
177 }
178 }
179}
180
181#[derive(Debug)]
184pub enum NodeState {
185 Online(EncodableString, Quality),
187 Offline(EncodableString),
189}
190
191impl NodeState {
192 pub fn write_bytes_vec(&self, out: &mut Vec<u8>) -> usize {
194 self.write_bytes(out).expect("Write to vector should't fail but did!")
195 }
196}
197
198impl ProtocolMessage for NodeState {
199 const MIN_SIZE: usize = 1 + EncodableString::MIN_SIZE;
200 fn byte_size(&self) -> usize {
201 let s = match self {
202 NodeState::Online(s, _) => s,
203 NodeState::Offline(s) => s,
204 };
205
206 s.byte_size() + 1
207 }
208
209 fn write_bytes<W: Write>(&self, out: &mut W) -> Result<usize> {
210 let (st, speed) = match self {
211 NodeState::Online(s, quality) => {
212 debug_assert!(quality.0 != u8::MAX);
213 (s, quality.0)
214 }
215 NodeState::Offline(s) => (s, u8::MAX),
216 };
217 let mut bytes = 0;
218 out.write_all(&[speed])?;
219 bytes += 1;
220 bytes += st.write_bytes(out)?;
221 Ok(bytes)
222 }
223
224 fn try_from_bytes(bytes: &[u8]) -> Result<(Self, usize)> {
225 if bytes.len() < 2 {
226 return Err(Error::BufferTooShort(2));
227 }
228
229 let quality = bytes[0];
230 let (node, node_len) =
231 EncodableString::try_from_bytes(&bytes[1..]).extend_buffer_too_short(1)?;
232
233 let state = if let Ok(quality) = quality.try_into() {
234 NodeState::Online(node, quality)
235 } else {
236 NodeState::Offline(node)
237 };
238
239 Ok((state, 1 + node_len))
240 }
241}