From 80feb4e70aac25d51cfc723d462b18f90b9325f8 Mon Sep 17 00:00:00 2001
From: lly <lly@espressif.com>
Date: Wed, 25 Mar 2020 12:30:53 +0800
Subject: [PATCH] ble_mesh: Update next alloc address when node info is added

---
 .../esp_ble_mesh/mesh_core/provisioner_main.c | 16 ++++++++++++++-
 .../esp_ble_mesh/mesh_core/provisioner_prov.c | 20 +++++++++++++++++++
 .../esp_ble_mesh/mesh_core/provisioner_prov.h | 12 +++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/components/bt/esp_ble_mesh/mesh_core/provisioner_main.c b/components/bt/esp_ble_mesh/mesh_core/provisioner_main.c
index e0d5fe654..a7177dd55 100644
--- a/components/bt/esp_ble_mesh/mesh_core/provisioner_main.c
+++ b/components/bt/esp_ble_mesh/mesh_core/provisioner_main.c
@@ -1512,6 +1512,8 @@ int bt_mesh_print_local_composition_data(void)
 #if CONFIG_BLE_MESH_TEST_AUTO_ENTER_NETWORK
 int bt_mesh_provisioner_store_node_info(struct bt_mesh_node *node)
 {
+    int err = 0;
+
     if (!node) {
         BT_ERR("%s, Invalid parameter", __func__);
         return -EINVAL;
@@ -1532,12 +1534,24 @@ int bt_mesh_provisioner_store_node_info(struct bt_mesh_node *node)
         return -EINVAL;
     }
 
+    if (node->unicast_addr + node->element_num - 1 > 0x7FFF) {
+        BT_ERR("%s, Not enough unicast address for the node", __func__);
+        return -EIO;
+    }
+
     if (bt_mesh_provisioner_net_key_get(node->net_idx) == NULL) {
         BT_ERR("%s, Invalid NetKey Index 0x%03x", __func__, node->net_idx);
         return -EINVAL;
     }
 
-    return provisioner_store_node(node, false, true, NULL);
+    err = provisioner_store_node(node, true, NULL);
+    if (err) {
+        BT_ERR("%s, Failed to store node info", __func__);
+        return err;
+    }
+
+    bt_mesh_test_provisioner_update_alloc_addr(node->unicast_addr, node->element_num);
+    return 0;
 }
 #endif /* CONFIG_BLE_MESH_TEST_AUTO_ENTER_NETWORK */
 
diff --git a/components/bt/esp_ble_mesh/mesh_core/provisioner_prov.c b/components/bt/esp_ble_mesh/mesh_core/provisioner_prov.c
index c061c682d..0245b3012 100644
--- a/components/bt/esp_ble_mesh/mesh_core/provisioner_prov.c
+++ b/components/bt/esp_ble_mesh/mesh_core/provisioner_prov.c
@@ -1245,6 +1245,26 @@ int bt_mesh_provisioner_set_primary_elem_addr(u16_t addr)
     return 0;
 }
 
+#if CONFIG_BLE_MESH_TEST_AUTO_ENTER_NETWORK
+int bt_mesh_test_provisioner_update_alloc_addr(u16_t unicast_addr, u16_t element_num)
+{
+    u16_t max_addr = FAST_PROV_ENABLE() ? prov_ctx.fast_prov.unicast_addr_max : PROV_MAX_ADDR_TO_ASSIGN;
+
+    if (unicast_addr + element_num > max_addr) {
+        BT_WARN("%s, Not enough unicast address to allocate", __func__);
+        prov_ctx.curr_alloc_addr = BLE_MESH_ADDR_UNASSIGNED;
+    } else {
+        prov_ctx.curr_alloc_addr = unicast_addr + element_num;
+    }
+
+    if (IS_ENABLED(CONFIG_BLE_MESH_SETTINGS)) {
+        bt_mesh_store_prov_info(prov_ctx.primary_addr, prov_ctx.curr_alloc_addr);
+    }
+
+    return 0;
+}
+#endif /* CONFIG_BLE_MESH_TEST_AUTO_ENTER_NETWORK */
+
 /* The following APIs are for fast provisioning */
 
 void bt_mesh_provisioner_fast_prov_enable(bool enable)
diff --git a/components/bt/esp_ble_mesh/mesh_core/provisioner_prov.h b/components/bt/esp_ble_mesh/mesh_core/provisioner_prov.h
index 7642e2c8e..ab5876b40 100644
--- a/components/bt/esp_ble_mesh/mesh_core/provisioner_prov.h
+++ b/components/bt/esp_ble_mesh/mesh_core/provisioner_prov.h
@@ -331,6 +331,18 @@ u16_t bt_mesh_provisioner_get_primary_elem_addr(void);
  */
 int bt_mesh_provisioner_set_primary_elem_addr(u16_t addr);
 
+/**
+ * @brief This function is used to update next allocated address by Provisioner.
+ *
+ * @note  This function is used for mesh internal test.
+ *
+ * @param[in] unicast_addr: unicast address of the node
+ * @param[in] element_num:  element count of the node
+ *
+ * @return Zero - success, otherwise - fail
+ */
+int bt_mesh_test_provisioner_update_alloc_addr(u16_t unicast_addr, u16_t element_num);
+
 /**
  * @brief This function is called to input number/string out-put by unprovisioned device.
  *