Handle UDP/TCP

1 files changed, 19 insertions, 8 deletions

diff --git a/src/consul_actor.rs b/src/consul_actor.rs
index 94ed9ab..241a643 100644
--- a/src/consul_actor.rs
+++ b/src/consul_actor.rs
@@ -11,7 +11,8 @@ use crate::consul;
 
 #[derive(Serialize, Deserialize, Debug)]
 pub enum DiplonatParameter {
-  port(Vec<u16>)
+  tcp_port(Vec<u16>),
+  udp_port(Vec<u16>)
 }
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -20,12 +21,12 @@ pub enum DiplonatConsul {
 }
 
 pub struct ConsulActor {
-  pub rx_open_ports: watch::Receiver<messages::OpenPorts>,
+  pub rx_open_ports: watch::Receiver<messages::PublicExposedPorts>,
 
   consul: consul::Consul,
   node: String,
   retries: u32,
-  tx_open_ports: watch::Sender<messages::OpenPorts>
+  tx_open_ports: watch::Sender<messages::PublicExposedPorts>
 }
 
 fn retry_to_time(retries: u32, max_time: Duration) -> Duration {
@@ -34,8 +35,12 @@ fn retry_to_time(retries: u32, max_time: Duration) -> Duration {
   return Duration::from_secs(cmp::min(max_time.as_secs(), 1.2f64.powf(retries as f64) as u64))
 }
 
-fn from_catalog_to_open_ports(catalog: &consul::CatalogNode) -> messages::OpenPorts {
-  let mut op = messages::OpenPorts { ports: Vec::new() };
+fn from_catalog_to_open_ports(catalog: &consul::CatalogNode) -> messages::PublicExposedPorts {
+  let mut op = messages::PublicExposedPorts { 
+    tcp_ports: Vec::new(),
+    udp_ports: Vec::new()
+  };
+
   for (_, service_info) in &catalog.Services {
     for tag in &service_info.Tags {
       let diplo_conf: error::Result<DiplonatConsul> = from_str(tag);
@@ -43,8 +48,10 @@ fn from_catalog_to_open_ports(catalog: &consul::CatalogNode) -> messages::OpenPo
         Ok(conf) => {
           let DiplonatConsul::diplonat(c) = conf;
           for parameter in &c {
-            let DiplonatParameter::port(p) = parameter;
-            op.ports.extend(p);
+            match parameter {
+              DiplonatParameter::tcp_port(p) => op.tcp_ports.extend(p),
+              DiplonatParameter::udp_port(p) => op.udp_ports.extend(p),
+            }
           }
         }
         Err(e) => debug!("Failed to parse entry {}. {}", tag, e),
@@ -56,7 +63,11 @@ fn from_catalog_to_open_ports(catalog: &consul::CatalogNode) -> messages::OpenPo
 
 impl ConsulActor {
   pub fn new(url: &str, node: &str) -> Self {
-    let (tx, rx) = watch::channel(messages::OpenPorts{ports: Vec::new() });
+    let (tx, rx) = watch::channel(messages::PublicExposedPorts{
+      tcp_ports: Vec::new(),
+      udp_ports: Vec::new()
+    });
+
     return Self {
       consul: consul::Consul::new(url),
       rx_open_ports: rx,