taskserver_protocol/
response.rs

1use itertools::Itertools;
2use std::{borrow::BorrowMut, collections::HashMap};
3
4/// All Response codes as defined in
5/// <https://taskwarrior.org/docs/design/protocol.html>
6pub 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
109/// Taskserver Response
110///
111/// see [ResponseBuildExt](crate::response::ResponseBuildExt) for details on modifying it
112pub struct Response {
113    /// the Response code the server is responding with
114    pub code: ResponseCode,
115    /// contains all headers. the client, code, status and protocol version headers are added when into_buf is called
116    pub extra_headers: HashMap<String, String>,
117    /// The payload the server is going to respond with
118    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    /// Creates new Response, with success Response code and empty payload
158    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    /// creates a buffer adapter to allow for writing out the response in one call
172    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) // length of key and value including ': ' and \n
187            .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
203/// Trait to allow modifying the response by either passing it by value or by mutable reference
204///
205/// ```
206/// use taskserver_protocol::response::*;
207/// let mut response = Response::new();
208/// let mut borrowed = &mut response;
209/// borrowed.code(ResponseCode::NoChange);
210///
211/// //we can still access the value (it was passed by ref)
212/// response.header("test", "foo");
213/// ```
214///
215pub trait ResponseBuildExt {
216    /// sets the code
217    fn code(self, code: ResponseCode) -> Self;
218    /// sets the given header to the provided name
219    fn header<I1: Into<String>, I2: Into<String>>(self, name: I1, value: I2) -> Self;
220    /// sets the response payload
221    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}