1use itertools::Itertools;
2use std::{borrow::BorrowMut, collections::HashMap};
3
4pub enum ResponseCode {
7 Success,
8 NoChange,
9 Redirect,
10 MalformedData,
11 UnsupportedEncoding,
12 TemporarilyUnavailable,
13 ShuttingDown,
14 AccessDenied,
15 AccountSuspended,
16 AccountTerminated,
17 SyntaxErrorInRequest,
18 SyntaxErrorIllegalParameters,
19 NotImplemented,
20 CommandParameterNotImplemented,
21 RequestTooBig,
22}
23
24impl ResponseCode {
25 fn get_code(&self) -> u16 {
26 match self {
27 ResponseCode::Success => 200,
28 ResponseCode::NoChange => 201,
29 ResponseCode::Redirect => 301,
30 ResponseCode::MalformedData => 400,
31 ResponseCode::UnsupportedEncoding => 401,
32 ResponseCode::TemporarilyUnavailable => 420,
33 ResponseCode::ShuttingDown => 421,
34 ResponseCode::AccessDenied => 430,
35 ResponseCode::AccountSuspended => 431,
36 ResponseCode::AccountTerminated => 432,
37 ResponseCode::SyntaxErrorInRequest => 500,
38 ResponseCode::SyntaxErrorIllegalParameters => 501,
39 ResponseCode::NotImplemented => 502,
40 ResponseCode::CommandParameterNotImplemented => 503,
41 ResponseCode::RequestTooBig => 504,
42 }
43 }
44 fn get_status_text(&self) -> &str {
45 match self {
46 ResponseCode::Success => "ok",
47 ResponseCode::NoChange => "no change",
48 ResponseCode::Redirect => "redirect to specified port",
49 ResponseCode::MalformedData => "malformed data",
50 ResponseCode::UnsupportedEncoding => "unsupported encoding",
51 ResponseCode::TemporarilyUnavailable => "max number of concurrent connections reached",
52 ResponseCode::ShuttingDown => "Server shutting down on operator request",
53 ResponseCode::AccessDenied => "access denied",
54 ResponseCode::AccountSuspended => "account suspended",
55 ResponseCode::AccountTerminated => "account terminated",
56 ResponseCode::SyntaxErrorInRequest => "syntax error in request",
57 ResponseCode::SyntaxErrorIllegalParameters => "illegal request parameters:",
58 ResponseCode::NotImplemented => "not implemented",
59 ResponseCode::CommandParameterNotImplemented => "command parameter not implemented",
60 ResponseCode::RequestTooBig => "request too big",
61 }
62 }
63}
64
65struct SerializedResponse {
66 serialized_header: String,
67 payload: Vec<String>,
68}
69
70#[derive(Clone)]
71enum RespIterValue {
72 MessageSizeBuf([u8; 4]),
73 BorrowedSlice(&'static [u8]),
74 String(String),
75}
76
77impl AsRef<[u8]> for RespIterValue {
78 fn as_ref(&self) -> &[u8] {
79 match self {
80 RespIterValue::MessageSizeBuf(buf) => buf,
81 RespIterValue::BorrowedSlice(slice) => *slice,
82 RespIterValue::String(string) => string.as_bytes(),
83 }
84 }
85}
86
87impl SerializedResponse {
88 fn into_buf(self) -> impl bytes::Buf {
89 let header_size: usize = self.serialized_header.len();
90 let payload_size: usize =
91 self.payload.len() + self.payload.iter().map(|s| s.len()).sum::<usize>();
92 let total_size: u32 = (header_size + payload_size + 4) as u32;
93
94 let data_iter = std::iter::once(self.serialized_header)
95 .chain(self.payload.into_iter())
96 .map(RespIterValue::String);
97 let newlines_iter = std::iter::repeat(RespIterValue::BorrowedSlice("\n".as_bytes()));
98 let data_iter = data_iter.interleave_shortest(newlines_iter);
99
100 ResponseBuf {
101 remaining_bytes: total_size as usize,
102 current_index: 0,
103 current_buf: RespIterValue::MessageSizeBuf(total_size.to_be_bytes()),
104 remaining_entries: data_iter,
105 }
106 }
107}
108
109pub struct Response {
113 pub code: ResponseCode,
115 pub extra_headers: HashMap<String, String>,
117 pub payload: Vec<String>,
119}
120
121struct ResponseBuf<I, T> {
122 remaining_bytes: usize,
123 current_index: usize,
124 current_buf: T,
125 remaining_entries: I,
126}
127
128impl<T: AsRef<[u8]>, I: Iterator<Item = T>> bytes::Buf for ResponseBuf<I, T> {
129 fn remaining(&self) -> usize {
130 self.remaining_bytes
131 }
132 fn bytes(&self) -> &[u8] {
133 &self.current_buf.as_ref()
134 }
135 fn advance(&mut self, mut cnt: usize) {
136 if cnt > self.remaining_bytes {
137 panic!("trying to advance beyond the buffer size");
138 }
139 self.remaining_bytes -= cnt;
140 while cnt > 0 {
141 let current_buf_len_remaining = self.current_buf.as_ref().len() - self.current_index;
142 if current_buf_len_remaining > cnt {
143 self.current_index += cnt;
144 return;
145 }
146 cnt -= current_buf_len_remaining;
147 self.current_index = 0;
148 self.current_buf = self
149 .remaining_entries
150 .next()
151 .expect("the underlying iterator did not provide enough bytes to fill the buffer");
152 }
153 }
154}
155
156impl Response {
157 pub fn new() -> Response {
159 let mut headers = HashMap::new();
160 headers.insert(
161 "client".into(),
162 concat!("taskd_rs ", env!("CARGO_PKG_VERSION")).into(),
163 );
164 Response {
165 code: ResponseCode::Success,
166 extra_headers: headers,
167 payload: Vec::new(),
168 }
169 }
170
171 pub fn into_buf(self) -> impl bytes::Buf {
173 self.serialize().into_buf()
174 }
175
176 fn serialize(mut self) -> SerializedResponse {
177 self.extra_headers
178 .insert("code".into(), format!("{}", self.code.get_code()));
179 self.extra_headers
180 .insert("status".into(), self.code.get_status_text().into());
181 self.extra_headers.insert("type".into(), "response".into());
182 self.extra_headers.insert("protocol".into(), "v1".into());
183 let len = self
184 .extra_headers
185 .iter()
186 .map(|(k, v)| k.len() + v.len() + 2 + 1) .sum::<usize>();
188 let mut serialized_header = String::with_capacity(len);
189 for (name, value) in self.extra_headers {
190 serialized_header.push_str(&name);
191 serialized_header.push_str(": ");
192 serialized_header.push_str(&value);
193 serialized_header.push('\n');
194 }
195
196 SerializedResponse {
197 serialized_header,
198 payload: self.payload,
199 }
200 }
201}
202
203pub trait ResponseBuildExt {
216 fn code(self, code: ResponseCode) -> Self;
218 fn header<I1: Into<String>, I2: Into<String>>(self, name: I1, value: I2) -> Self;
220 fn payload(self, payload: Vec<String>) -> Self;
222}
223
224impl<T: BorrowMut<Response>> ResponseBuildExt for T {
225 fn code(mut self, code: ResponseCode) -> Self {
226 self.borrow_mut().code = code;
227 self
228 }
229 fn header<I1: Into<String>, I2: Into<String>>(mut self, name: I1, value: I2) -> Self {
230 self.borrow_mut()
231 .extra_headers
232 .insert(name.into(), value.into());
233 self
234 }
235 fn payload(mut self, payload: Vec<String>) -> Self {
236 self.borrow_mut().payload = payload;
237 self
238 }
239}