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

use crate::errors::*;
use hyper::body::Body as HyperBody;
use hyper::http::Response as HyperResponse;
use hyper::{header::HeaderValue, HeaderMap};

// Extends HyperResponse
pub struct RhodResponse {
    res: Option<HyperResponse<HyperBody>>, // Is allways Some(..)
}

impl RhodResponse {
    pub fn new(res: HyperResponse<HyperBody>) -> RhodResponse {
        RhodResponse { res: Some(res) }
    }

    pub fn headers(&self) -> &HeaderMap<HeaderValue> {
        self.res.as_ref().unwrap().headers()
    }

    pub fn headers_mut(&mut self) -> &mut HeaderMap<HeaderValue> {
        self.res.as_mut().unwrap().headers_mut()
    }

    pub fn into_hyper_response(self) -> HyperResponse<HyperBody> {
        self.res.unwrap()
    }

    pub fn status_as_int(&self) -> u16 {
        self.res.as_ref().unwrap().status().as_u16()
    }

    pub async fn body(&mut self) -> RhodResult<Vec<u8>> {
        let r = self.res.take().unwrap();

        let (header, body) = r.into_parts();
        match hyper::body::to_bytes(body).await {
            Ok(b) => {
                let cloned = b.clone();
                self.res = Some(HyperResponse::from_parts(header, HyperBody::from(b)));
                Ok(cloned.to_vec())
            }
            Err(e) => {
                // If error, body cant be recovered.
                self.res = Some(HyperResponse::from_parts(header, HyperBody::empty()));

                Err(RhodError::from_string(
                    format!("Cant parse response body to bytes. {}", e),
                    RhodErrorLevel::Error,
                ))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_headers() {
        let mut res = RhodResponse::new(
            HyperResponse::builder()
                .header("Foo", "Bar")
                .body(HyperBody::empty())
                .unwrap(),
        );

        assert_eq!(res.headers().get("Foo").unwrap(), "Bar");
        assert!(res.headers().get("Accept").is_none());

        res.headers_mut()
            .insert("Accept", "text/html".parse().unwrap());
        assert_eq!(res.headers().get("Accept").unwrap(), "text/html");
    }

    #[test]
    fn test_status() {
        let res = RhodResponse::new(
            HyperResponse::builder()
                .status(404)
                .body(HyperBody::empty())
                .unwrap(),
        );

        assert_eq!(res.status_as_int(), 404);
    }

    #[tokio::test]
    async fn test_body() {
        let mut response = RhodResponse::new(
            HyperResponse::builder()
                .body(HyperBody::from("response bodyy %% #"))
                .unwrap(),
        );

        assert_eq!(
            response.body().await.unwrap(),
            "response bodyy %% #".as_bytes().to_vec()
        );

        let mut response =
            RhodResponse::new(HyperResponse::builder().body(HyperBody::empty()).unwrap());

        assert_eq!(response.body().await.unwrap(), [])
    }
}