scalesocket 0.2.4

A collaborative websocket server and autoscaler
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

use {std::collections::HashMap, std::net::SocketAddr, urlencoding::encode};

use crate::types::ConnID;

#[derive(Clone, Debug, Default)]
pub struct Env {
    pub cgi: CGIEnv,
    pub query: HashMap<String, String>,
}

impl Env {
    pub fn set_room(&mut self, room: &str) {
        self.cgi.room = Some(room.to_string());
    }
}

#[derive(Clone, Debug)]
#[cfg_attr(not(test), derive(Default))]
pub struct CGIEnv {
    /// URL-encoded search or parameter string
    query_string: String,
    /// network address of the client sending the request
    remote_addr: String,
    /// room name (non standard)
    room: Option<String>,
}

#[cfg(test)]
impl Default for CGIEnv {
    fn default() -> Self {
        Self {
            query_string: String::default(),
            remote_addr: String::default(),
            room: Some("".to_string()),
        }
    }
}

impl CGIEnv {
    pub fn from_filter(query_string: Option<String>, remote_addr: Option<SocketAddr>) -> Self {
        let query_string = query_string.unwrap_or_default();
        let remote_addr = remote_addr.map(|a| a.to_string()).unwrap_or_default();

        Self {
            query_string,
            remote_addr,
            ..Default::default()
        }
    }
}

impl From<CGIEnv> for HashMap<String, String> {
    fn from(env: CGIEnv) -> Self {
        HashMap::from([
            // NOTE: implicit uppercase
            ("QUERY_STRING".to_string(), env.query_string),
            ("REMOTE_ADDR".to_string(), env.remote_addr),
            ("ROOM".to_string(), env.room.expect("room to be defined")),
        ])
    }
}

pub fn replace_template_env(template: &str, conn: ConnID, env: &Env) -> String {
    let template = template.replace("#ID", &conn.to_string());

    let cgi_vars = env.cgi.clone().into();
    let query_vars = env.query.clone().keys_upper();

    let result = replace_template(template, cgi_vars, "#", false);
    replace_template(result, query_vars, "#QUERY_", true)
}

fn replace_template(
    template: String,
    replace_values: HashMap<String, String>,
    prefix: &str,
    urlencode: bool,
) -> String {
    let mut result = template;
    for (key, value) in replace_values {
        let variable = &[prefix, &key].concat();
        if urlencode {
            result = result.replace(variable, &encode(&value));
        } else {
            result = result.replace(variable, &value);
        };
    }
    result
}
trait HashMapExt {
    fn keys_upper(self) -> HashMap<String, String>;
}

impl HashMapExt for HashMap<String, String> {
    fn keys_upper(self) -> HashMap<String, String> {
        self.into_iter()
            .map(|(k, v)| (k.to_uppercase(), v))
            .collect()
    }
}

#[cfg(test)]
mod tests {

    use std::collections::HashMap;

    use super::{CGIEnv, Env, replace_template_env};

    fn create_query() -> HashMap<String, String> {
        HashMap::from([("foo".to_string(), "bar baz".to_string())])
    }

    fn create_cgi() -> CGIEnv {
        CGIEnv {
            query_string: "foo=".to_string(),
            remote_addr: "127.0.0.1:1234".to_string(),
            room: "room".to_string().into(),
        }
    }

    #[tokio::test]
    async fn test_replace_template_env() {
        let env = Env {
            cgi: create_cgi(),
            query: create_query(),
        };
        let result = replace_template_env("test #ID #REMOTE_ADDR #QUERY_FOO", 1, &env);

        assert_eq!(result, "test 1 127.0.0.1:1234 bar%20baz");
    }

    #[tokio::test]
    async fn test_replace_template_env_omits_query_overrides() {
        let query = HashMap::from([
            ("remote_addr".to_string(), "overridden".to_string()),
            ("string".to_string(), "overridden".to_string()),
            ("hack".to_string(), "#SOMETHING".to_string()),
        ]);
        let env = Env {
            cgi: create_cgi(),
            query,
        };
        let result = replace_template_env("test #REMOTE_ADDR #QUERY_STRING #QUERY_HACK", 1, &env);

        assert_eq!(result, "test 127.0.0.1:1234 foo= %23SOMETHING");
    }
}