1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
use std::cell::RefCell;
use std::fmt::{self, Debug, Display, Formatter};
use std::io::{Read, Write};
use std::net::TcpStream;
use std::sync::{Arc, Mutex};

use crate::consts;
use crate::header::{HeaderType, Headers};
use crate::http::status::Status;
use crate::{
    error::Result, header::headers_to_string, internal::handle::Writeable, Content, Header,
    SetCookie,
};

/// Http Response
#[derive(Debug)]
pub struct Response {
    /// Response status code
    pub status: Status,

    /// Response Data.
    /// Can be either a Static `Vec<u8>` or a Stream (impl [`Read`])
    pub data: ResponseBody,

    /// List of response headers.
    /// This does not contain the default headers.
    pub headers: Headers,

    /// Response reason phrase.
    /// If this is None, the reason phrase will be automatically generated based on the status code.
    pub reason: Option<String>,

    /// Response Flags:
    /// - Close: Set the Connection header to close and will close the connection after the response is sent.
    /// - End: End the connection without sending a response
    pub flag: ResponseFlag,
}

#[derive(Debug, PartialEq, Eq)]
pub enum ResponseFlag {
    /// No Flag
    None,
    /// Close the socket
    Close,
    /// End the connection without sending a response
    End,
}

/// Response Data.
/// Can be either a Static Vec<u8> or a Stream (impl [`Read`]).
/// Static responses are sent in one go, while streams are sent in chunks (chunked transfer encoding).
pub enum ResponseBody {
    Static(Vec<u8>),
    Stream(Writeable),
}

impl Response {
    /// Create a new Blank Response
    ///
    /// Default data is as follows
    /// - Status: 200
    /// - Data: OK
    /// - Headers: Vec::new()
    /// ## Example
    /// ```rust
    /// # use afire::{Response, Header};
    /// Response::new();
    /// ```
    pub fn new() -> Self {
        Self {
            status: Status::Ok,
            data: vec![79, 75].into(),
            headers: Default::default(),
            reason: None,
            flag: ResponseFlag::None,
        }
    }

    /// Creates a new Default Response with the End flag set.
    pub fn end() -> Self {
        Self {
            flag: ResponseFlag::End,
            ..Default::default()
        }
    }

    /// Add a status code to a Response.
    /// This accepts [`Status`] as well as a [`u16`].
    /// ## Example
    /// ```rust
    /// # use afire::{Response, Header, Status};
    /// // Create Response
    /// Response::new().status(Status::Ok);
    /// ```
    pub fn status(self, code: impl Into<Status>) -> Self {
        Self {
            status: code.into(),
            ..self
        }
    }

    /// Manually set the Reason Phrase.
    /// If this is not set, it will be inferred from the status code.
    /// Non standard status codes will have a reason phrase of "OK".
    /// ```rust
    /// # use afire::{Response, Header, Status};
    /// // Create Response
    /// let response = Response::new()
    ///     .status(Status::Ok)
    ///     .reason("Hello");
    /// ```
    pub fn reason(self, reason: impl AsRef<str>) -> Self {
        Self {
            reason: Some(reason.as_ref().to_owned()),
            ..self
        }
    }

    /// Add text as data to a Response.
    /// Will accept any type that implements Display, such as [`String`], [`str`], [`i32`], serde_json::Value, etc.
    /// This response type is considered static and will be sent in one go, not chunked.
    /// ## Example
    /// ```rust
    /// # use afire::Response;
    /// // Create Response
    /// let response = Response::new()
    ///    .text("Hello from afire!");
    /// ```
    pub fn text(self, text: impl Display) -> Self {
        Self {
            data: text.to_string().as_bytes().to_vec().into(),
            ..self
        }
    }

    /// Add raw bytes as data to a Response.
    /// This response type is considered static and will be sent in one go, not chunked.
    /// ## Example
    /// ```rust
    /// # use afire::Response;
    /// // Create Response
    /// let response = Response::new()
    ///   .bytes(&[79, 75]); // Bytes for "OK"
    /// ```
    pub fn bytes(self, bytes: &[u8]) -> Self {
        Self {
            data: bytes.to_vec().into(),
            ..self
        }
    }

    /// Add a stream as data to a Response.
    /// This response type is considered dynamic and will be streamed to the client in chunks using `Transfer-Encoding: chunked`.
    /// ## Example
    /// ```rust,no_run
    /// # use afire::{Response, Method, Server};
    /// # use std::fs::File;
    /// const PATH: &str = "path/to/file.txt";
    /// let mut server = Server::<()>::new("localhost", 8080);
    ///
    /// server.route(Method::GET, "/download-stream", |_| {
    ///     let stream = File::open(PATH).unwrap();
    ///     Response::new().stream(stream)
    /// });
    /// ```
    pub fn stream(self, stream: impl Read + Send + 'static) -> Self {
        Self {
            data: ResponseBody::Stream(Box::new(RefCell::new(stream))),
            ..self
        }
    }

    /// Add a Header to a Response.
    /// Will accept any type that implements `AsRef<str>`, so [`String`], [`str`], [`&str`], etc.
    /// ## Example
    /// ```rust
    /// # use afire::{Response, Header};
    /// // Create Response
    /// let response = Response::new()
    ///    .header("Test-Header", "Test-Value");
    /// ```
    pub fn header(mut self, key: impl Into<HeaderType>, value: impl AsRef<str>) -> Self {
        self.headers.push(Header::new(key, value));
        self
    }

    /// Add a list of Headers to a Response.
    /// Only accepts a slice of [`Header`]s.
    /// ## Example
    /// ```rust
    /// # use afire::{Response, Header};
    /// // Create Response
    /// let response = Response::new()
    ///     .headers(&[
    ///         Header::new("Content-Type", "text/html"),
    ///         Header::new("Test-Header", "Test-Value")
    ///     ]);
    /// ```
    pub fn headers(mut self, headers: &[Header]) -> Self {
        self.headers.append(&mut headers.to_vec());
        self
    }

    /// Will set the `Connection: close` header on the Response.
    /// Then it will close the connection after the Response has been sent.
    /// ## Example
    /// ```rust
    /// # use afire::{Response};
    /// // Create Response
    /// let response = Response::new()
    ///   .text("goodbye!")
    ///   .close();
    /// ```
    pub fn close(self) -> Self {
        Self {
            flag: ResponseFlag::Close,
            ..self
        }
    }

    /// Add a cookie to a response.
    /// The [`SetCookie`] will be converted to a [`Header`] and added to the Response.
    /// ## Example
    /// ```
    /// # use afire::{Response, SetCookie};
    /// // Create Response and add cookie
    /// let response = Response::new()
    ///     .cookie(SetCookie::new("name", "value"))
    ///     .cookie(SetCookie::new("name2", "value2"));
    /// ```
    pub fn cookie(mut self, cookie: SetCookie) -> Self {
        self.headers
            .push(Header::new("Set-Cookie", cookie.to_string()));
        self
    }

    /// Add a list of cookies to a response.
    /// ## Example
    /// ```
    /// # use afire::{Response, SetCookie};
    /// // Create Response and add cookie
    /// let response = Response::new()
    ///     .cookies(&[
    ///         SetCookie::new("name", "value"),
    ///         SetCookie::new("name2", "value2")
    ///     ]);
    /// ```
    pub fn cookies(self, cookie: &[SetCookie]) -> Self {
        let mut new = Vec::new();

        for c in cookie {
            new.push(Header::new("Set-Cookie", c.to_string()));
        }

        self.headers(&new)
    }

    /// Set a Content Type on a Response with a [`Content`] enum.
    /// This will add a `Content-Type` header to the Response.
    /// ## Example
    /// ```
    /// # use afire::{Response, Content};
    /// // Create Response and type
    /// let response = Response::new()
    ///     .content(Content::HTML);
    /// ```
    pub fn content(mut self, content_type: Content) -> Self {
        self.headers.push(content_type.into());
        self
    }

    /// Lets you modify the Response with a function before it is sent to the client.
    /// This can be used to have middleware that modifies the Response on specific routes.
    pub fn modifier(self, modifier: impl Fn(Response) -> Response) -> Self {
        modifier(self)
    }

    // TODO: Make crate local
    /// Writes a Response to a TcpStream.
    /// Will take care of adding default headers and closing the connection if needed.
    pub fn write(
        &mut self,
        stream: Arc<Mutex<TcpStream>>,
        default_headers: &[Header],
    ) -> Result<()> {
        // Add default headers to response
        // Only the ones that aren't already in the response
        for i in default_headers {
            if !self.headers.has(&i.name) {
                self.headers.push(i.clone());
            }
        }

        let static_body = self.data.is_static();

        // Add content-length header to response if we are sending a static body
        if static_body && !self.headers.has(HeaderType::ContentLength) {
            self.headers.push(self.data.content_len());
        }

        // Add Connection: close if response is set to close
        if self.flag == ResponseFlag::Close && !self.headers.has(HeaderType::Connection) {
            self.headers.push(Header::new("Connection", "close"));
        }

        if !static_body && !self.headers.has(HeaderType::TransferEncoding) {
            self.headers
                .push(Header::new("Transfer-Encoding", "chunked"));
        }

        // Convert the response to a string
        let response = format!(
            "HTTP/1.1 {} {}\r\n{}\r\n\r\n",
            self.status.code(),
            self.reason
                .to_owned()
                .unwrap_or_else(|| self.status.reason_phrase().to_owned()),
            headers_to_string(&self.headers)
        );

        let mut stream = stream.lock().unwrap();
        stream.write_all(response.as_bytes())?;
        self.data.write(&mut stream)?;

        Ok(())
    }
}

impl Default for Response {
    fn default() -> Response {
        Response::new()
    }
}

impl ResponseBody {
    pub fn empty() -> Self {
        ResponseBody::Static(Vec::new())
    }

    /// Checks if the ResponseBody is static.
    fn is_static(&self) -> bool {
        matches!(self, ResponseBody::Static(_))
    }

    /// Gets the content length header of a static ResponseBody.
    /// If the ResponseBody is not static it will panic.
    fn content_len(&self) -> Header {
        let len = match self {
            ResponseBody::Static(data) => data.len(),
            _ => unreachable!("Can't get content length of a stream"),
        };
        Header::new("Content-Length", len.to_string())
    }

    /// Writes a ResponseBody to a TcpStream.
    /// Either in one go if it is static or in chunks if it is a stream.
    fn write(&mut self, stream: &mut TcpStream) -> Result<()> {
        match self {
            ResponseBody::Static(data) => stream.write_all(data)?,
            ResponseBody::Stream(data) => {
                let data = data.get_mut();
                loop {
                    let mut chunk = vec![0; consts::CHUNK_SIZE];
                    let read = data.read(&mut chunk)?;
                    if read == 0 {
                        break;
                    }

                    let mut section = format!("{read:X}\r\n").as_bytes().to_vec();
                    section.extend(&chunk[..read]);
                    section.extend(b"\r\n");

                    stream.write_all(&section)?;
                }

                stream.write_all(b"0\r\n\r\n")?;
            }
        };

        Ok(())
    }
}

impl From<Vec<u8>> for ResponseBody {
    fn from(x: Vec<u8>) -> Self {
        ResponseBody::Static(x)
    }
}

impl From<Writeable> for ResponseBody {
    fn from(x: Writeable) -> Self {
        ResponseBody::Stream(x)
    }
}

impl Debug for ResponseBody {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Self::Static(arg) => f.debug_tuple("Static").field(arg).finish(),
            Self::Stream(_arg) => f.debug_tuple("Stream").finish(),
        }
    }
}