From be5b46d23fa426f9f5d17e0369f745d7502a386c Mon Sep 17 00:00:00 2001
From: Sviatoslav Boichuk <sviatoslav.boichuk@plvision.eu>
Date: Tue, 22 Oct 2024 13:53:32 +0300
Subject: [PATCH] Update NB on device join/leave event

---
 src/cgw_connection_processor.rs |  3 ++-
 src/cgw_connection_server.rs    | 25 +++++++++++++++++--
 src/cgw_nb_api_listener.rs      | 44 +++++++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/src/cgw_connection_processor.rs b/src/cgw_connection_processor.rs
index 3f84669..92332ba 100644
--- a/src/cgw_connection_processor.rs
+++ b/src/cgw_connection_processor.rs
@@ -218,7 +218,8 @@ impl CGWConnectionProcessor {
         // we can proceed.
         debug!("Sending ACK request for device serial: {}", self.serial);
         let (mbox_tx, mut mbox_rx) = unbounded_channel::<CGWConnectionProcessorReqMsg>();
-        let msg = CGWConnectionServerReqMsg::AddNewConnection(evt.serial, caps, mbox_tx);
+        let msg =
+            CGWConnectionServerReqMsg::AddNewConnection(evt.serial, self.addr.ip(), caps, mbox_tx);
         self.cgw_server
             .enqueue_mbox_message_to_cgw_server(msg)
             .await;
diff --git a/src/cgw_connection_server.rs b/src/cgw_connection_server.rs
index 8635ebe..2def66d 100644
--- a/src/cgw_connection_server.rs
+++ b/src/cgw_connection_server.rs
@@ -5,7 +5,8 @@ use crate::cgw_nb_api_listener::{
     cgw_construct_device_capabilities_changed_msg, cgw_construct_device_enqueue_response,
     cgw_construct_foreign_infra_connection_msg, cgw_construct_infra_group_create_response,
     cgw_construct_infra_group_delete_response, cgw_construct_infra_group_device_add_response,
-    cgw_construct_infra_group_device_del_response, cgw_construct_rebalance_group_response,
+    cgw_construct_infra_group_device_del_response, cgw_construct_infra_join_msg,
+    cgw_construct_infra_leave_msg, cgw_construct_rebalance_group_response,
     cgw_construct_unassigned_infra_connection_msg,
 };
 use crate::cgw_runtime::{cgw_get_runtime, CGWRuntimeType};
@@ -34,7 +35,11 @@ use crate::{
 use crate::cgw_errors::{Error, Result};
 
 use std::str::FromStr;
-use std::{collections::HashMap, net::SocketAddr, sync::Arc};
+use std::{
+    collections::HashMap,
+    net::{IpAddr, SocketAddr},
+    sync::Arc,
+};
 use tokio::{
     net::TcpStream,
     runtime::Runtime,
@@ -84,6 +89,7 @@ pub enum CGWConnectionServerReqMsg {
     // Connection-related messages
     AddNewConnection(
         MacAddress,
+        IpAddr,
         CGWDeviceCapabilities,
         UnboundedSender<CGWConnectionProcessorReqMsg>,
     ),
@@ -1458,6 +1464,7 @@ impl CGWConnectionServer {
 
                 if let CGWConnectionServerReqMsg::AddNewConnection(
                     device_mac,
+                    ip_addr,
                     caps,
                     conn_processor_mbox_tx,
                 ) = msg
@@ -1661,6 +1668,14 @@ impl CGWConnectionServer {
                             .await;
                     }
 
+                    if let Ok(resp) =
+                        cgw_construct_infra_join_msg(device_group_id, device_mac, ip_addr)
+                    {
+                        self.enqueue_mbox_message_from_cgw_to_nb_api(device_group_id, resp);
+                    } else {
+                        error!("Failed to construct device_join message!");
+                    }
+
                     connmap_w_lock.insert(device_mac, conn_processor_mbox_tx);
 
                     tokio::spawn(async move {
@@ -1723,6 +1738,12 @@ impl CGWConnectionServer {
                             .await;
                     }
 
+                    if let Ok(resp) = cgw_construct_infra_leave_msg(device_group_id, device_mac) {
+                        self.enqueue_mbox_message_from_cgw_to_nb_api(device_group_id, resp);
+                    } else {
+                        error!("Failed to construct device_leave message!");
+                    }
+
                     CGWMetrics::get_ref().change_counter(
                         CGWMetricsCounterType::ConnectionsNum,
                         CGWMetricsCounterOpType::Dec,
diff --git a/src/cgw_nb_api_listener.rs b/src/cgw_nb_api_listener.rs
index 30d0343..5272fc6 100644
--- a/src/cgw_nb_api_listener.rs
+++ b/src/cgw_nb_api_listener.rs
@@ -19,6 +19,7 @@ use rdkafka::{
 };
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
+use std::net::IpAddr;
 use std::sync::Arc;
 use tokio::{
     runtime::{Builder, Runtime},
@@ -140,6 +141,21 @@ pub struct APClientMigrateMessage {
     pub to_band: String,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct InfraJoinMessage {
+    pub r#type: &'static str,
+    pub infra_group_id: i32,
+    pub infra_group_infra: MacAddress,
+    infra_public_ip: IpAddr,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct InfraLeaveMessage {
+    pub r#type: &'static str,
+    pub infra_group_id: i32,
+    pub infra_group_infra: MacAddress,
+}
+
 pub fn cgw_construct_infra_group_create_response(
     infra_group_id: i32,
     infra_name: String,
@@ -356,6 +372,34 @@ pub fn cgw_construct_client_migrate_msg(
     Ok(serde_json::to_string(&client_migrate_msg)?)
 }
 
+pub fn cgw_construct_infra_join_msg(
+    infra_group_id: i32,
+    infra_group_infra: MacAddress,
+    infra_public_ip: IpAddr,
+) -> Result<String> {
+    let infra_join_msg = InfraJoinMessage {
+        r#type: "infra_join",
+        infra_group_id,
+        infra_group_infra,
+        infra_public_ip,
+    };
+
+    Ok(serde_json::to_string(&infra_join_msg)?)
+}
+
+pub fn cgw_construct_infra_leave_msg(
+    infra_group_id: i32,
+    infra_group_infra: MacAddress,
+) -> Result<String> {
+    let infra_leave_msg = InfraLeaveMessage {
+        r#type: "infra_leave",
+        infra_group_id,
+        infra_group_infra,
+    };
+
+    Ok(serde_json::to_string(&infra_leave_msg)?)
+}
+
 struct CustomContext;
 impl ClientContext for CustomContext {}