Merge pull request #145 from ngrok/ngrok/ryan/support-traffic-policy

add support for TrafficPolicy

muxado/examples/heartbeat.rs

@@ -47,6 +47,7 @@ async fn main() -> Result<(), anyhow::Error> {
     Ok(())
 }
 
+#[allow(dead_code)]
 struct HHandler;
 
 #[async_trait::async_trait]

muxado/examples/server.rs

@@ -18,6 +18,7 @@ use tracing::{
 };
 
 /// Subscriber that claims that it's always enabled, but does nothing.
+#[allow(dead_code)]
 struct AlwaysSubscriber;
 impl tracing::Subscriber for AlwaysSubscriber {
     #[inline]

ngrok/examples/axum.rs

@@ -62,7 +62,7 @@ async fn start_tunnel() -> anyhow::Result<HttpTunnel> {
         //         .allow_domain("<domain>")
         //         .scope("<scope>"),
         // )
-        // .policy(create_policy())?
+        // .traffic_policy(POLICY_JSON)
         // .proxy_proto(ProxyProto::None)
         // .remove_request_header("X-Req-Nope")
         // .remove_response_header("X-Res-Nope")
@@ -81,26 +81,48 @@ async fn start_tunnel() -> anyhow::Result<HttpTunnel> {
 }
 
 #[allow(dead_code)]
-fn create_policy() -> Result<Policy, InvalidPolicy> {
-    Ok(Policy::new()
-        .add_inbound(
-            Rule::new("deny_put")
-                .add_expression("req.Method == 'PUT'")
-                .add_action(Action::new("deny", None)?),
-        )
-        .add_outbound(
-            Rule::new("200_response")
-                .add_expression("res.StatusCode == '200'")
-                .add_action(Action::new(
-                    "custom-response",
-                    Some(
-                        r###"{
-                    "status_code": 200,
-                    "content_type": "text/html",
-                    "content": "Custom 200 response."
-                }"###,
-                    ),
-                )?),
-        )
-        .to_owned())
-}
+const POLICY_JSON: &str = r###"{
+    "inbound":[
+        {
+            "name":"deny_put",
+            "expressions":["req.Method == 'PUT'"],
+            "actions":[{"Type":"deny"}]
+        }],
+    "outbound":[
+        {
+            "name":"change success response",
+            "expressions":["res.StatusCode == '200'"],
+            "actions":[{
+                "type":"custom-response",
+                "config":{
+                    "status_code":201, 
+                    "content": "Custom 200 response.", 
+                    "headers": {
+                        "content_type": "text/html"
+                    }
+                }
+            }]
+        }]
+}"###;
+
+#[allow(dead_code)]
+const POLICY_YAML: &str = r###"
+---
+inbound:
+    - name: "deny_put"
+      expressions:
+      - "req.Method == 'PUT'"
+      actions:
+      - type: "deny"
+outbound:
+    - name: "change success response"
+      expressions:
+      - "res.StatusCode == '200'"
+      actions:
+      - type: "custom-response"
+        config:
+          status_code: 201
+          content: "Custom 200 response."
+          headers:
+            content_type: "text/html"
+"###;

ngrok/src/config/common.rs

@@ -206,9 +206,11 @@ pub(crate) struct CommonOpts {
     pub(crate) forwards_proto: Option<String>,
     // Whether to disable certificate verification for this tunnel.
     verify_upstream_tls: Option<bool>,
+    // DEPRECATED: use traffic_policy instead.
+    pub(crate) policy: Option<Policy>,
     // Policy that defines rules that should be applied to incoming or outgoing
     // connections to the edge.
-    pub(crate) policy: Option<Policy>,
+    pub(crate) traffic_policy: Option<String>,
 }
 
 impl CommonOpts {

ngrok/src/config/http.rs

@@ -1,6 +1,7 @@
 use std::{
     borrow::Borrow,
     collections::HashMap,
+    convert::From,
     str::FromStr,
 };
 
@@ -188,7 +189,13 @@ impl TunnelConfig for HttpOptions {
                 .websocket_tcp_conversion
                 .then_some(WebsocketTcpConverter {}),
             user_agent_filter: self.user_agent_filter(),
-            policy: self.common_opts.policy.clone().map(From::from),
+            traffic_policy: if self.common_opts.traffic_policy.is_some() {
+                self.common_opts.traffic_policy.clone().map(From::from)
+            } else if self.common_opts.policy.is_some() {
+                self.common_opts.policy.clone().map(From::from)
+            } else {
+                None
+            },
             ..Default::default()
         };
 
@@ -440,7 +447,7 @@ impl HttpTunnelBuilder {
         self
     }
 
-    /// Set the policy for this edge.
+    /// DEPRECATED: use traffic_policy instead.
     pub fn policy<S>(&mut self, s: S) -> Result<&mut Self, S::Error>
     where
         S: TryInto<Policy>,
@@ -449,6 +456,12 @@ impl HttpTunnelBuilder {
         Ok(self)
     }
 
+    /// Set policy for this edge.
+    pub fn traffic_policy(&mut self, policy_str: impl Into<String>) -> &mut Self {
+        self.options.common_opts.traffic_policy = Some(policy_str.into());
+        self
+    }
+
     pub(crate) async fn for_forwarding_to(&mut self, to_url: &Url) -> &mut Self {
         self.options.common_opts.for_forwarding_to(to_url);
         if let Some(host) = to_url.host_str().filter(|_| self.options.rewrite_host) {

ngrok/src/config/policies.rs

@@ -163,6 +163,12 @@ impl Action {
     }
 }
 
+impl From<Policy> for proto::PolicyWrapper {
+    fn from(value: Policy) -> Self {
+        proto::PolicyWrapper::Policy(value.into())
+    }
+}
+
 // transform into the wire protocol format
 impl From<Policy> for proto::Policy {
     fn from(o: Policy) -> Self {

ngrok/src/config/tcp.rs

@@ -1,4 +1,7 @@
-use std::collections::HashMap;
+use std::{
+    collections::HashMap,
+    convert::From,
+};
 
 use url::Url;
 
@@ -74,7 +77,13 @@ impl TunnelConfig for TcpOptions {
 
         tcp_endpoint.ip_restriction = self.common_opts.ip_restriction();
 
-        tcp_endpoint.policy = self.common_opts.policy.clone().map(From::from);
+        tcp_endpoint.traffic_policy = if self.common_opts.traffic_policy.is_some() {
+            self.common_opts.traffic_policy.clone().map(From::from)
+        } else if self.common_opts.policy.is_some() {
+            self.common_opts.policy.clone().map(From::from)
+        } else {
+            None
+        };
 
         Some(BindOpts::Tcp(tcp_endpoint))
     }
@@ -152,7 +161,7 @@ impl TcpTunnelBuilder {
         self
     }
 
-    /// Set the policy for this edge.
+    /// DEPRECATED: use traffic_policy instead.
     pub fn policy<S>(&mut self, s: S) -> Result<&mut Self, S::Error>
     where
         S: TryInto<Policy>,
@@ -161,6 +170,12 @@ impl TcpTunnelBuilder {
         Ok(self)
     }
 
+    /// Set policy for this edge.
+    pub fn traffic_policy(&mut self, policy_str: impl Into<String>) -> &mut Self {
+        self.options.common_opts.traffic_policy = Some(policy_str.into());
+        self
+    }
+
     pub(crate) async fn for_forwarding_to(&mut self, to_url: &Url) -> &mut Self {
         self.options.common_opts.for_forwarding_to(to_url);
         self

ngrok/src/config/tls.rs

@@ -94,7 +94,13 @@ impl TunnelConfig for TlsOptions {
         tls_endpoint.mutual_tls_at_edge =
             (!self.mutual_tlsca.is_empty()).then_some(self.mutual_tlsca.as_slice().into());
         tls_endpoint.tls_termination = tls_termination;
-        tls_endpoint.policy = self.common_opts.policy.clone().map(From::from);
+        tls_endpoint.traffic_policy = if self.common_opts.traffic_policy.is_some() {
+            self.common_opts.traffic_policy.clone().map(From::from)
+        } else if self.common_opts.policy.is_some() {
+            self.common_opts.policy.clone().map(From::from)
+        } else {
+            None
+        };
 
         Some(BindOpts::Tls(tls_endpoint))
     }
@@ -192,7 +198,7 @@ impl TlsTunnelBuilder {
         self
     }
 
-    /// Set the policy for this edge.
+    /// DEPRECATED: use traffic_policy instead.
     pub fn policy<S>(&mut self, s: S) -> Result<&mut Self, S::Error>
     where
         S: TryInto<Policy>,
@@ -201,6 +207,12 @@ impl TlsTunnelBuilder {
         Ok(self)
     }
 
+    /// Set policy for this edge.
+    pub fn traffic_policy(&mut self, policy_str: impl Into<String>) -> &mut Self {
+        self.options.common_opts.traffic_policy = Some(policy_str.into());
+        self
+    }
+
     pub(crate) async fn for_forwarding_to(&mut self, to_url: &Url) -> &mut Self {
         self.options.common_opts.for_forwarding_to(to_url);
         self