ngrok/config/
policies.rs

1use std::{
2    fs::read_to_string,
3    io,
4};
5
6use serde::{
7    Deserialize,
8    Serialize,
9};
10use thiserror::Error;
11
12use crate::internals::proto;
13
14/// A policy that defines rules that should be applied to incoming or outgoing
15/// connections to the edge.
16#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)]
17#[serde(default)]
18pub struct Policy {
19    inbound: Vec<Rule>,
20    outbound: Vec<Rule>,
21}
22
23/// A policy rule that should be applied
24#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)]
25#[serde(default)]
26pub struct Rule {
27    name: String,
28    expressions: Vec<String>,
29    actions: Vec<Action>,
30}
31
32/// An action that should be taken if the rule matches
33#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)]
34#[serde(default)]
35pub struct Action {
36    #[serde(rename = "type")]
37    type_: String,
38    config: Option<serde_json::Value>,
39}
40
41/// Errors in creating or serializing Policies
42#[derive(Debug, Error)]
43pub enum InvalidPolicy {
44    /// Error representing an invalid string for a Policy
45    #[error("failure to parse or generate policy")]
46    SerializationError(#[from] serde_json::Error),
47    /// An error loading a Policy from a file
48    #[error("failure to read policy file '{}'", .1)]
49    FileReadError(#[source] io::Error, String),
50}
51
52impl Policy {
53    /// Create a new empty [Policy] struct
54    pub fn new() -> Self {
55        Policy {
56            ..Default::default()
57        }
58    }
59
60    /// Create a new [Policy] from a json string
61    fn from_json(json: impl AsRef<str>) -> Result<Self, InvalidPolicy> {
62        serde_json::from_str(json.as_ref()).map_err(InvalidPolicy::SerializationError)
63    }
64
65    /// Create a new [Policy] from a json file
66    pub fn from_file(json_file_path: impl AsRef<str>) -> Result<Self, InvalidPolicy> {
67        Policy::from_json(
68            read_to_string(json_file_path.as_ref()).map_err(|e| {
69                InvalidPolicy::FileReadError(e, json_file_path.as_ref().to_string())
70            })?,
71        )
72    }
73
74    /// Convert [Policy] to json string
75    pub fn to_json(&self) -> Result<String, InvalidPolicy> {
76        serde_json::to_string(&self).map_err(InvalidPolicy::SerializationError)
77    }
78
79    /// Add an inbound policy
80    pub fn add_inbound(&mut self, rule: impl Into<Rule>) -> &mut Self {
81        self.inbound.push(rule.into());
82        self
83    }
84
85    /// Add an outbound policy
86    pub fn add_outbound(&mut self, rule: impl Into<Rule>) -> &mut Self {
87        self.outbound.push(rule.into());
88        self
89    }
90}
91
92impl TryFrom<&Policy> for Policy {
93    type Error = InvalidPolicy;
94
95    fn try_from(other: &Policy) -> Result<Policy, Self::Error> {
96        Ok(other.clone())
97    }
98}
99
100impl TryFrom<Result<Policy, InvalidPolicy>> for Policy {
101    type Error = InvalidPolicy;
102
103    fn try_from(other: Result<Policy, InvalidPolicy>) -> Result<Policy, Self::Error> {
104        other
105    }
106}
107
108impl TryFrom<&str> for Policy {
109    type Error = InvalidPolicy;
110
111    fn try_from(other: &str) -> Result<Policy, Self::Error> {
112        Policy::from_json(other)
113    }
114}
115
116impl Rule {
117    /// Create a new [Rule]
118    pub fn new(name: impl Into<String>) -> Self {
119        Rule {
120            name: name.into(),
121            ..Default::default()
122        }
123    }
124
125    /// Convert [Rule] to json string
126    pub fn to_json(&self) -> Result<String, InvalidPolicy> {
127        serde_json::to_string(&self).map_err(InvalidPolicy::SerializationError)
128    }
129
130    /// Add an expression
131    pub fn add_expression(&mut self, expression: impl Into<String>) -> &mut Self {
132        self.expressions.push(expression.into());
133        self
134    }
135
136    /// Add an action
137    pub fn add_action(&mut self, action: Action) -> &mut Self {
138        self.actions.push(action);
139        self
140    }
141}
142
143impl From<&mut Rule> for Rule {
144    fn from(other: &mut Rule) -> Self {
145        other.to_owned()
146    }
147}
148
149impl Action {
150    /// Create a new [Action]
151    pub fn new(type_: impl Into<String>, config: Option<&str>) -> Result<Self, InvalidPolicy> {
152        Ok(Action {
153            type_: type_.into(),
154            config: config
155                .map(|c| serde_json::from_str(c).map_err(InvalidPolicy::SerializationError))
156                .transpose()?,
157        })
158    }
159
160    /// Convert [Action] to json string
161    pub fn to_json(&self) -> Result<String, InvalidPolicy> {
162        serde_json::to_string(&self).map_err(InvalidPolicy::SerializationError)
163    }
164}
165
166impl From<Policy> for proto::PolicyWrapper {
167    fn from(value: Policy) -> Self {
168        proto::PolicyWrapper::Policy(value.into())
169    }
170}
171
172// transform into the wire protocol format
173impl From<Policy> for proto::Policy {
174    fn from(o: Policy) -> Self {
175        proto::Policy {
176            inbound: o.inbound.into_iter().map(|p| p.into()).collect(),
177            outbound: o.outbound.into_iter().map(|p| p.into()).collect(),
178        }
179    }
180}
181
182impl From<Rule> for proto::Rule {
183    fn from(p: Rule) -> Self {
184        proto::Rule {
185            name: p.name,
186            expressions: p.expressions,
187            actions: p.actions.into_iter().map(|a| a.into()).collect(),
188        }
189    }
190}
191
192impl From<Action> for proto::Action {
193    fn from(a: Action) -> Self {
194        proto::Action {
195            type_: a.type_,
196            config: a
197                .config
198                .map(|c| c.to_string().into_bytes())
199                .unwrap_or_default(),
200        }
201    }
202}
203
204#[cfg(test)]
205pub(crate) mod test {
206    use super::*;
207
208    pub(crate) const POLICY_JSON: &str = r###"
209        {"inbound": [
210            {
211                "name": "test_in",
212                "expressions": ["req.Method == 'PUT'"],
213                "actions": [{"type": "deny"}]
214            }
215        ],
216        "outbound": [
217            {
218                "name": "test_out",
219                "expressions": ["res.StatusCode == '200'"],
220                "actions": [{"type": "custom-response", "config": {"status_code":201}}]
221            }
222        ]}
223        "###;
224
225    #[test]
226    fn test_json_to_policy() {
227        let policy: Policy = Policy::from_json(POLICY_JSON).unwrap();
228        assert_eq!(1, policy.inbound.len());
229        assert_eq!(1, policy.outbound.len());
230        let inbound = &policy.inbound[0];
231        let outbound = &policy.outbound[0];
232
233        assert_eq!("test_in", inbound.name);
234        assert_eq!(1, inbound.expressions.len());
235        assert_eq!(1, inbound.actions.len());
236        assert_eq!("req.Method == 'PUT'", inbound.expressions[0]);
237        assert_eq!("deny", inbound.actions[0].type_);
238        assert_eq!(None, inbound.actions[0].config);
239
240        assert_eq!("test_out", outbound.name);
241        assert_eq!(1, outbound.expressions.len());
242        assert_eq!(1, outbound.actions.len());
243        assert_eq!("res.StatusCode == '200'", outbound.expressions[0]);
244        assert_eq!("custom-response", outbound.actions[0].type_);
245        assert_eq!(
246            "{\"status_code\":201}",
247            outbound.actions[0].config.as_ref().unwrap().to_string()
248        );
249    }
250
251    #[test]
252    fn test_empty_json_to_policy() {
253        let policy: Policy = Policy::from_json("{}").unwrap();
254        assert_eq!(0, policy.inbound.len());
255        assert_eq!(0, policy.outbound.len());
256    }
257
258    #[test]
259    fn test_policy_to_json() {
260        let policy = Policy::from_json(POLICY_JSON).unwrap();
261        let json = policy.to_json().unwrap();
262        let policy2 = Policy::from_json(json).unwrap();
263        assert_eq!(policy, policy2);
264    }
265
266    #[test]
267    fn test_policy_to_json_error() {
268        let error = Policy::from_json("asdf").err().unwrap();
269        assert!(matches!(error, InvalidPolicy::SerializationError { .. }));
270    }
271
272    #[test]
273    fn test_rule_to_json() {
274        let policy = Policy::from_json(POLICY_JSON).unwrap();
275        let rule = &policy.outbound[0];
276        let json = rule.to_json().unwrap();
277        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
278        let rule_map = parsed.as_object().unwrap();
279        assert_eq!("test_out", rule_map["name"]);
280
281        // expressions
282        let expressions = rule_map["expressions"].as_array().unwrap();
283        assert_eq!(1, expressions.len());
284        assert_eq!("res.StatusCode == '200'", expressions[0]);
285
286        // actions
287        let actions = rule_map["actions"].as_array().unwrap();
288        assert_eq!(1, actions.len());
289        assert_eq!("custom-response", actions[0]["type"]);
290        assert_eq!(201, actions[0]["config"]["status_code"]);
291    }
292
293    #[test]
294    fn test_action_to_json() {
295        let policy = Policy::from_json(POLICY_JSON).unwrap();
296        let action = &policy.outbound[0].actions[0];
297        let json = action.to_json().unwrap();
298        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
299        let action_map = parsed.as_object().unwrap();
300        assert_eq!("custom-response", action_map["type"]);
301        assert_eq!(201, action_map["config"]["status_code"]);
302    }
303
304    #[test]
305    fn test_builders() {
306        let policy = Policy::from_json(POLICY_JSON).unwrap();
307        let policy2 = Policy::new()
308            .add_inbound(
309                Rule::new("test_in")
310                    .add_expression("req.Method == 'PUT'")
311                    .add_action(Action::new("deny", None).unwrap()),
312            )
313            .add_outbound(
314                Rule::new("test_out")
315                    .add_expression("res.StatusCode == '200'")
316                    // .add_action(Action::new("deny", ""))
317                    .add_action(
318                        Action::new("custom-response", Some("{\"status_code\":201}")).unwrap(),
319                    ),
320            )
321            .to_owned();
322        assert_eq!(policy, policy2);
323    }
324
325    #[test]
326    fn test_load_file() {
327        let policy = Policy::from_json(POLICY_JSON).unwrap();
328        let policy2 = Policy::from_file("assets/policy.json").unwrap();
329        assert_eq!("test_in", policy2.inbound[0].name);
330        assert_eq!("test_out", policy2.outbound[0].name);
331        assert_eq!(policy, policy2);
332    }
333
334    #[test]
335    fn test_load_inbound_file() {
336        let policy = Policy::from_file("assets/policy-inbound.json").unwrap();
337        assert_eq!("test_in", policy.inbound[0].name);
338        assert_eq!(0, policy.outbound.len());
339    }
340
341    #[test]
342    fn test_load_file_error() {
343        let error = Policy::from_file("assets/absent.json").err().unwrap();
344        assert!(matches!(error, InvalidPolicy::FileReadError { .. }));
345    }
346}