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
use std::marker::PhantomData;

use mlua::{FromLua, Lua, Result, String as LuaString, Table, TableExt, TablePairs, ToLua, Value};

/// The "Http" class contain all the HTTP manipulation functions.
#[derive(Clone)]
pub struct Http<'lua>(Table<'lua>);

#[derive(Clone)]
pub struct Headers<'lua>(Table<'lua>);

impl<'lua> Http<'lua> {
    /// Returns a `Headers` table containing all the request headers.
    pub fn req_get_headers(&self) -> Result<Headers<'lua>> {
        self.0.call_method("req_get_headers", ())
    }

    /// Returns a `Headers` table containing all the response headers.
    pub fn res_get_headers(&self) -> Result<Headers<'lua>> {
        self.0.call_method("res_get_headers", ())
    }

    /// Appends an HTTP header field `name` with `value` in the request.
    pub fn req_add_header<V: ToLua<'lua>>(&self, name: &str, value: V) -> Result<()> {
        self.0.call_method("req_add_header", (name, value))
    }

    /// Appends an HTTP header field `name` with `value` in the response.
    pub fn res_add_header<V: ToLua<'lua>>(&self, name: &str, value: V) -> Result<()> {
        self.0.call_method("res_add_header", (name, value))
    }

    /// Removes all HTTP header fields in the request by `name`.
    pub fn req_del_header(&self, name: &str) -> Result<()> {
        self.0.call_method("req_del_header", name)
    }

    /// Removes all HTTP header fields in the response by `name`.
    pub fn res_del_header(&self, name: &str) -> Result<()> {
        self.0.call_method("res_del_header", name)
    }

    /// Replaces all occurrence of HTTP request header `name`, by only one containing the `value`.
    pub fn req_set_header<V: ToLua<'lua>>(&self, name: &str, value: V) -> Result<()> {
        self.0.call_method("req_set_header", (name, value))
    }

    /// Replaces all occurrence of HTTP response header `name`, by only one containing the `value`.
    pub fn res_set_header<V: ToLua<'lua>>(&self, name: &str, value: V) -> Result<()> {
        self.0.call_method("res_set_header", (name, value))
    }

    /// Matches the regular expression in all occurrences of HTTP request header `name` according to `regex`,
    /// and replaces them with the `replace` argument.
    /// The replacement value can contain back references like 1, 2, ...
    pub fn req_rep_header(&self, name: &str, regex: &str, replace: &str) -> Result<()> {
        self.0.call_method("req_rep_header", (name, regex, replace))
    }

    /// Matches the regular expression in all occurrences of HTTP response header `name` according to `regex`,
    /// and replaces them with the `replace` argument.
    /// The replacement value can contain back references like 1, 2, ...
    pub fn res_rep_header(&self, name: &str, regex: &str, replace: &str) -> Result<()> {
        self.0.call_method("res_rep_header", (name, regex, replace))
    }

    /// Rewrites the request method with the `method`.
    pub fn req_set_method(&self, method: &str) -> Result<()> {
        self.0.call_method("req_set_method", method)
    }

    /// Rewrites the request path with the `path`.
    pub fn req_set_path(&self, path: &str) -> Result<()> {
        self.0.call_method("req_set_path", path)
    }

    /// Rewrites the request’s query string which appears after the first question mark (`?`)
    /// with the `query`.
    pub fn req_set_query(&self, query: &str) -> Result<()> {
        self.0.call_method("req_set_query", query)
    }

    /// Rewrites the request URI with the `uri`.
    pub fn req_set_uri(&self, uri: &str) -> Result<()> {
        self.0.call_method("req_set_uri", uri)
    }

    /// Rewrites the response status code with the `code`.
    /// If no custom reason is provided, it will be generated from the status.
    pub fn res_set_status(&self, status: u16, reason: Option<&str>) -> Result<()> {
        self.0.call_method("res_set_status", (status, reason))
    }
}

impl<'lua> Headers<'lua> {
    pub fn pairs<V: FromLua<'lua>>(self) -> HeaderPairs<'lua, V> {
        HeaderPairs {
            pairs: self.0.pairs(),
            phantom: PhantomData,
        }
    }

    /// Returns all header fields by `name`.
    pub fn get<V: FromLua<'lua>>(&self, name: &str) -> Result<Vec<V>> {
        let name = name.to_ascii_lowercase();
        let mut result = Vec::new();
        if let Some(values) = self.0.get::<_, Option<Table>>(name)? {
            let mut pairs = values.pairs::<i32, V>().collect::<Result<Vec<_>>>()?;
            pairs.sort_by_key(|x| x.0);
            result = pairs.into_iter().map(|(_, v)| v).collect();
        }
        Ok(result)
    }

    /// Returns first header field by `name`.
    pub fn get_first<V: FromLua<'lua>>(&self, name: &str) -> Result<Option<V>> {
        let name = name.to_ascii_lowercase();
        if let Some(values) = self.0.get::<_, Option<Table>>(name)? {
            return values.get(0); // Indexes starts from "0"
        }
        Ok(None)
    }
}

impl<'lua> FromLua<'lua> for Http<'lua> {
    fn from_lua(value: Value<'lua>, lua: &'lua Lua) -> Result<Self> {
        Ok(Http(Table::from_lua(value, lua)?))
    }
}

impl<'lua> FromLua<'lua> for Headers<'lua> {
    fn from_lua(value: Value<'lua>, lua: &'lua Lua) -> Result<Self> {
        Ok(Headers(Table::from_lua(value, lua)?))
    }
}

pub struct HeaderPairs<'lua, V: FromLua<'lua>> {
    pairs: TablePairs<'lua, LuaString<'lua>, Table<'lua>>,
    phantom: PhantomData<V>,
}

impl<'lua, V: FromLua<'lua>> Iterator for HeaderPairs<'lua, V> {
    type Item = Result<(String, Vec<V>)>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.pairs.next() {
            Some(Ok(item)) => {
                let name = String::from_utf8_lossy(item.0.as_bytes()).into_owned();
                let pairs = item.1.pairs::<i32, V>().collect::<Result<Vec<_>>>();
                match pairs {
                    Ok(mut pairs) => {
                        pairs.sort_by_key(|x| x.0);
                        Some(Ok((name, pairs.into_iter().map(|(_, v)| v).collect())))
                    }
                    Err(e) => Some(Err(e)),
                }
            }
            Some(Err(e)) => Some(Err(e)),
            None => None,
        }
    }
}