From d07ae9f2af6270114ca2069e811a611af657ef35 Mon Sep 17 00:00:00 2001
From: Ryan Grant <ryan@ngrok.com>
Date: Wed, 24 Jul 2024 12:15:54 -0700
Subject: [PATCH] plumb TrafficPolicy, remap Policy

---
 ngrok-java-native/src/lib.rs                  | 18 +++++++++--
 .../src/test/java/com/ngrok/DataTest.java     | 12 ++++++++
 .../main/java/com/ngrok/EndpointBuilder.java  | 30 +++++++++++++++----
 3 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/ngrok-java-native/src/lib.rs b/ngrok-java-native/src/lib.rs
index 0b584e7..8623a16 100644
--- a/ngrok-java-native/src/lib.rs
+++ b/ngrok-java-native/src/lib.rs
@@ -470,7 +470,11 @@ impl<'local> NativeSessionRsImpl<'local> {
         bldr.proxy_proto(ProxyProto::from(jeb.get_proxy_proto_version(self.env)));
 
         if let Some(policy) = jeb.get_policy(self.env).of_string(self.env) {
-            bldr.policy(policy.as_str()).map_err(io_exc)?;
+            bldr.traffic_policy(policy.as_str());
+        }
+
+        if let Some(traffic_policy) = jeb.get_traffic_policy(self.env).of_string(self.env) {
+            bldr.traffic_policy(traffic_policy);
         }
 
         // from TcpBuilder
@@ -514,7 +518,11 @@ impl<'local> NativeSessionRsImpl<'local> {
         bldr.proxy_proto(ProxyProto::from(jeb.get_proxy_proto_version(self.env)));
 
         if let Some(policy) = jeb.get_policy(self.env).of_string(self.env) {
-            bldr.policy(policy.as_str()).map_err(io_exc)?;
+            bldr.traffic_policy(policy.as_str());
+        }
+
+        if let Some(traffic_policy) = jeb.get_traffic_policy(self.env).of_string(self.env) {
+            bldr.traffic_policy(traffic_policy);
         }
 
         // from TlsBuilder
@@ -580,7 +588,11 @@ impl<'local> NativeSessionRsImpl<'local> {
         bldr.proxy_proto(ProxyProto::from(jeb.get_proxy_proto_version(self.env)));
 
         if let Some(policy) = jeb.get_policy(self.env).of_string(self.env) {
-            bldr.policy(policy.as_str()).map_err(io_exc)?;
+            bldr.traffic_policy(policy.as_str());
+        }
+
+        if let Some(traffic_policy) = jeb.get_traffic_policy(self.env).of_string(self.env) {
+            bldr.traffic_policy(traffic_policy);
         }
 
         // from HttpBuilder
diff --git a/ngrok-java-native/src/test/java/com/ngrok/DataTest.java b/ngrok-java-native/src/test/java/com/ngrok/DataTest.java
index 92bd0bd..1603a8a 100644
--- a/ngrok-java-native/src/test/java/com/ngrok/DataTest.java
+++ b/ngrok-java-native/src/test/java/com/ngrok/DataTest.java
@@ -60,4 +60,16 @@ public void testPolicy() throws Exception {
             Runtime.getLogger().log("info", "session", listener.getUrl());
         }
     }
+
+    @Test
+    public void testTrafficPolicy() throws Exception {
+        final ClassLoader classLoader = getClass().getClassLoader();
+        final String trafficPolicy = new String(classLoader.getResourceAsStream("policy.json").readAllBytes());
+
+        try (var session = Session.withAuthtokenFromEnv().connect();
+                var listener = session.httpEndpoint().metadata("java-endpoint").trafficPolicy(trafficPolicy).listen()) {
+            assertEquals("java-endpoint", listener.getMetadata());
+            Runtime.getLogger().log("info", "session", listener.getUrl());
+        }
+    }
 }
diff --git a/ngrok-java/src/main/java/com/ngrok/EndpointBuilder.java b/ngrok-java/src/main/java/com/ngrok/EndpointBuilder.java
index 4f27771..e70cdfe 100644
--- a/ngrok-java/src/main/java/com/ngrok/EndpointBuilder.java
+++ b/ngrok-java/src/main/java/com/ngrok/EndpointBuilder.java
@@ -14,7 +14,7 @@ public abstract class EndpointBuilder<T extends EndpointBuilder<T>> extends Meta
     private final List<String> allowCIDR = new ArrayList<>();
     private final List<String> denyCIDR = new ArrayList<>();
     private ProxyProto proxyProto = ProxyProto.None;
-    private Optional<String> policy = Optional.empty();
+    private Optional<String> trafficPolicy = Optional.empty();
 
     /**
      * Adds a CIDR to the list of allowed CIDRs for this endpoint.
@@ -60,13 +60,24 @@ public T proxyProto(ProxyProto proxyProto) {
     }
 
     /**
-     * Sets the policy for this endpoint.
+     * DEPRECATED: use trafficPolicy instead.
      *
      * @param policy the policy for the builder
      * @return An instance the builder represented by type T
      */
     public T policy(final String policy) {
-        this.policy = Optional.ofNullable(policy);
+        this.trafficPolicy = Optional.ofNullable(policy);
+        return (T) this;
+    }
+
+    /**
+     * Sets the policy for this endpoint.
+     *
+     * @param trafficPolicy the policy for the builder
+     * @return An instance the builder represented by type T
+     */
+    public T trafficPolicy(final String trafficPolicy) {
+        this.trafficPolicy = Optional.ofNullable(trafficPolicy);
         return (T) this;
     }
 
@@ -107,11 +118,20 @@ public long getProxyProtoVersion() {
     }
 
     /**
-     * Returns the policy for this endpoint.
+     * DEPRECATED: use getTrafficPolicy instead.
      *
      * @return the currently set policy
      */
     public Optional<String> getPolicy() {
-        return this.policy;
+        return this.trafficPolicy;
+    }
+
+    /**
+     * Returns the policy for this endpoint.
+     *
+     * @return the currently set policy
+     */
+    public Optional<String> getTrafficPolicy() {
+        return this.trafficPolicy;
     }
 }