feat(mam): Add PSK for channel encryption

Sender/Receiver can assign multiple PSK keys to de/encrypt
a MAM channel.

Author: howjmay

accelerator/core/BUILD

@@ -63,6 +63,8 @@ cc_library(
         "//accelerator/core",
         "//accelerator/core/request",
         "//accelerator/core/response",
+        "//utils:fill_nines",
+        "//utils:tryte_byte_conv",
         "@entangled//common/trinary:flex_trit",
         "@entangled//mam/api",
         "@entangled//utils/containers/hash:hash_array",

accelerator/core/mam_core.c

@@ -9,6 +9,8 @@
 #include "accelerator/core/mam_core.h"
 #include "common/model/transfer.h"
 #include "utils/containers/hash/hash_array.h"
+#include "utils/fill_nines.h"
+#include "utils/tryte_byte_conv.h"
 
 #define MAM_LOGGER "mam_core"
 
@@ -161,26 +163,53 @@ static mam_endpoint_t *mam_api_endpoint_get(mam_api_t const *const api, tryte_t
   return NULL;
 }
 
+/**
+ * @brief Free 'mam_encrypt_key_t' object
+ *
+ * @param mam_key[in] 'mam_encrypt_key_t' object to be freed
+ *
+ * @return status code
+ */
+static inline void mam_encrypt_key_free(mam_encrypt_key_t *mam_key) {
+  mam_psk_t_set_free(&mam_key->psks);
+  mam_ntru_pk_t_set_free(&mam_key->ntru_pks);
+  mam_ntru_sk_t_set_free(&mam_key->ntru_sks);
+}
+
 /**
  * @brief Add all the keys in the list of Pre-Shared Key and NTRU public key into corresponding key set.
  *
- * @param psks[in] Pre-Shared Key set
- * @param ntru_pks[in] NTRU public key set
+ * The PSK keys are converting the index into trytes as PSK ID.
+ *
+ * @param psks[out] Pre-Shared Key set
+ * @param ntru_pks[out] NTRU public key set
  * @param psk[in] List of Pre-Shared Key
  * @param ntru_pk[in] List of NTRU public key
  *
  * @return status code
  */
-static status_t ta_set_mam_key(mam_psk_t_set_t *const psks, mam_ntru_pk_t_set_t *const ntru_pks,
-                               UT_array const *const psk, UT_array const *const ntru_pk) {
+static status_t ta_set_mam_key(mam_encrypt_key_t *const mam_keys, UT_array const *const psk,
+                               UT_array const *const ntru_pk, UT_array const *const ntru_sk) {
+  status_t ret = SC_OK;
   char **p = NULL;
   if (psk) {
     mam_psk_t psk_obj;
+    uint16_t psk_id_cnt = 0;
     while ((p = (char **)utarray_next(psk, p))) {
+      tryte_t raw_psk_id[NUM_TRYTES_MAM_PSK_ID_SIZE + 1] = {}, psk_id[NUM_TRYTES_MAM_PSK_ID_SIZE + 1] = {};
+      bytes_to_trytes((unsigned char *)&psk_id_cnt, sizeof(psk_id_cnt) / sizeof(char), (char *)raw_psk_id);
+      ret = fill_nines((char *)psk_id, (char *)raw_psk_id, NUM_TRYTES_MAM_PSK_ID_SIZE);
+      if (ret) {
+        ta_log_error("%s\n", ta_error_to_string(ret));
+        return ret;
+      }
+
+      trytes_to_trits((tryte_t *)psk_id, psk_obj.id, NUM_TRYTES_MAM_PSK_ID_SIZE);
       trytes_to_trits((tryte_t *)*p, psk_obj.key, NUM_TRYTES_MAM_PSK_KEY_SIZE);
-      if (mam_psk_t_set_add(psks, &psk_obj) != RC_OK) {
-        ta_log_error("%s\n", "SC_MAM_FAILED_INIT");
-        return SC_MAM_FAILED_INIT;
+      if (mam_psk_t_set_add(&mam_keys->psks, &psk_obj) != RC_OK) {
+        ret = SC_MAM_FAILED_INIT;
+        ta_log_error("%s\n", ta_error_to_string(ret));
+        return ret;
       }
     }
   }
@@ -189,9 +218,22 @@ static status_t ta_set_mam_key(mam_psk_t_set_t *const psks, mam_ntru_pk_t_set_t
     mam_ntru_pk_t ntru_pk_obj;
     while ((p = (char **)utarray_next(ntru_pk, p))) {
       trytes_to_trits((tryte_t *)*p, ntru_pk_obj.key, NUM_TRYTES_MAM_NTRU_PK_SIZE);
-      if (mam_ntru_pk_t_set_add(ntru_pks, &ntru_pk_obj) != RC_OK) {
-        ta_log_error("%s\n", "SC_MAM_FAILED_INIT");
-        return SC_MAM_FAILED_INIT;
+      if (mam_ntru_pk_t_set_add(&mam_keys->ntru_pks, &ntru_pk_obj) != RC_OK) {
+        ret = SC_MAM_FAILED_INIT;
+        ta_log_error("%s\n", ta_error_to_string(ret));
+        return ret;
+      }
+    }
+  }
+
+  if (ntru_sk) {
+    mam_ntru_sk_t ntru_sk_obj;
+    while ((p = (char **)utarray_next(ntru_sk, p))) {
+      trytes_to_trits((tryte_t *)*p, ntru_sk_obj.secret_key, NUM_TRYTES_MAM_NTRU_SK_SIZE);
+      if (mam_ntru_sk_t_set_add(&mam_keys->ntru_sks, &ntru_sk_obj) != RC_OK) {
+        ret = SC_MAM_FAILED_INIT;
+        ta_log_error("%s\n", ta_error_to_string(ret));
+        return ret;
       }
     }
   }
@@ -528,7 +570,8 @@ status_t ta_send_mam_message(const ta_config_t *const info, const iota_config_t
   tryte_t chid[MAM_CHANNEL_ID_TRYTE_SIZE] = {}, msg_id[NUM_TRYTES_MAM_MSG_ID] = {};
   bundle_transactions_t *bundle = NULL;
   send_mam_data_mam_v1_t *data = (send_mam_data_mam_v1_t *)req->data;
-  mam_encrypt_key_t mam_key = {.psks = NULL, .ntru_pks = NULL};
+  send_mam_key_mam_v1_t *key = (send_mam_key_mam_v1_t *)req->key;
+  mam_encrypt_key_t mam_key = {.psks = NULL, .ntru_pks = NULL, .ntru_sks = NULL};
   bool msg_sent = false;
 
   // Creating MAM API
@@ -538,6 +581,12 @@ status_t ta_send_mam_message(const ta_config_t *const info, const iota_config_t
     goto done;
   }
 
+  ret = ta_set_mam_key(&mam_key, key->psk_array, key->ntru_array, NULL);
+  if (ret) {
+    ta_log_error("%s\n", ta_error_to_string(ret));
+    goto done;
+  }
+
   mam_send_operation_t mam_operation;
   while (!msg_sent) {
     bundle_transactions_renew(&bundle);
@@ -608,7 +657,7 @@ status_t ta_send_mam_message(const ta_config_t *const info, const iota_config_t
     }
   }
   bundle_transactions_free(&bundle);
-
+  mam_encrypt_key_free(&mam_key);
   return ret;
 }
 
@@ -619,6 +668,7 @@ status_t ta_recv_mam_message(const iota_config_t *const iconf, const iota_client
   bundle_array_t *bundle_array = NULL;
   bundle_array_new(&bundle_array);
   recv_mam_data_id_mam_v1_t *data_id = (recv_mam_data_id_mam_v1_t *)req->data_id;
+  recv_mam_key_mam_v1_t *key = (recv_mam_key_mam_v1_t *)req->key;
   if (mam_api_init(&mam, (tryte_t *)iconf->seed) != RC_OK) {
     ret = SC_MAM_FAILED_INIT;
     ta_log_error("%s\n", ta_error_to_string(ret));
@@ -652,6 +702,23 @@ status_t ta_recv_mam_message(const iota_config_t *const iconf, const iota_client
     }
   }
 
+  // Add decryption keys
+  mam_encrypt_key_t mam_key = {.psks = NULL, .ntru_pks = NULL, .ntru_sks = NULL};
+  ret = ta_set_mam_key(&mam_key, key->psk_array, NULL, key->ntru_array);
+  if (ret != SC_OK) {
+    ta_log_error("%s\n", ta_error_to_string(ret));
+    goto done;
+  }
+
+  mam_psk_t_set_entry_t *curr_psk_p = NULL;
+  mam_psk_t_set_entry_t *tmp_psk_p = NULL;
+  HASH_ITER(hh, mam_key.psks, curr_psk_p, tmp_psk_p) {
+    if (mam_api_add_psk(&mam, &curr_psk_p->value)) {
+      ta_log_error("%s\n", "Failed to add PSK keys");
+      goto done;
+    }
+  }
+
   // Copy the trusted_channel_pks, before fetching the information from MAM.
   mam_pk_t_set_t init_trusted_ch = NULL;
   mam_pk_t_set_entry_t *curr_entry = NULL;
@@ -694,5 +761,6 @@ status_t ta_recv_mam_message(const iota_config_t *const iconf, const iota_client
   }
   bundle_array_free(&bundle_array);
   mam_pk_t_set_free(&init_trusted_ch);
+  mam_encrypt_key_free(&mam_key);
   return ret;
 }

accelerator/core/mam_core.h

@@ -44,6 +44,7 @@ typedef enum mam_send_operation_e { ANNOUNCE_CHID, SEND_MESSAGE } mam_send_opera
 typedef struct mam_encrypt_key_s {
   mam_psk_t_set_t psks;
   mam_ntru_pk_t_set_t ntru_pks;
+  mam_ntru_sk_t_set_t ntru_sks;
 } mam_encrypt_key_t;
 
 /**

accelerator/core/request/ta_recv_mam.c

@@ -34,7 +34,8 @@ static void recv_mam_req_v1_free(ta_recv_mam_req_t** req) {
 
   if ((*req)->key) {
     recv_mam_key_mam_v1_t* key = (*req)->key;
-    free(key->enc_key);
+    utarray_free(key->psk_array);
+    utarray_free(key->ntru_array);
     free(key);
     (*req)->key = NULL;
   }
@@ -58,21 +59,39 @@ void recv_mam_req_free(ta_recv_mam_req_t** req) {
   *req = NULL;
 }
 
-status_t recv_mam_set_mam_v1_data_id(ta_recv_mam_req_t* req, char* bundle_hash, char* chid, char* msg_id) {
-  if (req == NULL || (!bundle_hash && !chid && !msg_id)) {
+status_t recv_mam_req_v1_init(ta_recv_mam_req_t* req) {
+  if (req == NULL) {
     return SC_TA_NULL;
   }
-  status_t ret = SC_OK;
 
   req->data_id = (recv_mam_data_id_mam_v1_t*)malloc(sizeof(recv_mam_data_id_mam_v1_t));
-  if (!req->data_id) {
+  if (req->data_id == NULL) {
+    return SC_TA_OOM;
+  }
+  req->key = (recv_mam_key_mam_v1_t*)malloc(sizeof(recv_mam_key_mam_v1_t));
+  if (req->key == NULL) {
     return SC_TA_OOM;
   }
 
   recv_mam_data_id_mam_v1_t* data_id = (recv_mam_data_id_mam_v1_t*)req->data_id;
   data_id->bundle_hash = NULL;
   data_id->chid = NULL;
   data_id->msg_id = NULL;
+
+  recv_mam_key_mam_v1_t* key = (recv_mam_key_mam_v1_t*)req->key;
+  utarray_new(key->psk_array, &ut_str_icd);
+  utarray_new(key->ntru_array, &ut_str_icd);
+
+  return SC_OK;
+}
+
+status_t recv_mam_set_mam_v1_data_id(ta_recv_mam_req_t* req, char* bundle_hash, char* chid, char* msg_id) {
+  if (req == NULL || (!bundle_hash && !chid && !msg_id)) {
+    return SC_TA_NULL;
+  }
+  status_t ret = SC_OK;
+
+  recv_mam_data_id_mam_v1_t* data_id = (recv_mam_data_id_mam_v1_t*)req->data_id;
   if (bundle_hash) {
     data_id->bundle_hash = (tryte_t*)strdup(bundle_hash);
     if (!data_id->bundle_hash) {
@@ -98,38 +117,8 @@ status_t recv_mam_set_mam_v1_data_id(ta_recv_mam_req_t* req, char* bundle_hash,
   return ret;
 
 error:
-  free(req->data_id);
-  req->data_id = NULL;
+  free(data_id->bundle_hash);
+  free(data_id->chid);
+  free(data_id->msg_id);
   return ret;
 }
-
-status_t recv_mam_set_mam_v1_key(ta_recv_mam_req_t* req, tryte_t* psk, tryte_t* ntru) {
-  if (!req) {
-    return SC_TA_NULL;
-  }
-
-  req->key = (recv_mam_key_mam_v1_t*)malloc(sizeof(recv_mam_key_mam_v1_t));
-  if (!req->key) {
-    return SC_TA_OOM;
-  }
-
-  recv_mam_key_mam_v1_t* key = (recv_mam_key_mam_v1_t*)req->key;
-  // We will set either psk or ntru, so initializing both fields will avoid errors in the future.
-  key->enc_key = NULL;
-
-  if (psk) {
-    key->enc_key = (tryte_t*)malloc(sizeof(tryte_t) * NUM_TRYTES_MAM_PSK_KEY_SIZE);
-    if (!key->enc_key) {
-      return SC_TA_OOM;
-    }
-    memcpy(key->enc_key, psk, sizeof(tryte_t) * NUM_TRYTES_MAM_PSK_KEY_SIZE);
-  } else if (ntru) {
-    key->enc_key = (tryte_t*)malloc(sizeof(tryte_t) * NUM_TRYTES_MAM_NTRU_PK_SIZE);
-    if (!key->enc_key) {
-      return SC_TA_OOM;
-    }
-    memcpy(key->enc_key, ntru, sizeof(tryte_t) * NUM_TRYTES_MAM_NTRU_PK_SIZE);
-  }
-
-  return SC_OK;
-}

accelerator/core/request/ta_recv_mam.h

@@ -44,7 +44,7 @@ ta_recv_mam_req_t* recv_mam_req_new();
 /**
  * Free memory of ta_recv_mam_req_t
  *
- * @param req Data type of ta_recv_mam_req_t
+ * @param req[out] Data type of ta_recv_mam_req_t
  */
 void recv_mam_req_free(ta_recv_mam_req_t** req);
 
@@ -60,18 +60,19 @@ typedef struct recv_mam_data_id_mam_v1_s {
 } recv_mam_data_id_mam_v1_t;
 
 typedef struct recv_mam_key_mam_v1_s {
-  /** Message encryption key. This field could be Pre-Shared Key (81 trytes) or NTRU public key (1024 trytes). Default:
-   * NULL. */
-  tryte_t* enc_key;
+  /** Optional. The pre-shared key to decrypt the message. Each psk is in length of 81 trytes. Default: NULL. */
+  UT_array* psk_array;
+  /** Optional. The NTRU Secret key to decrypt the message. Each psk is in length of 1024 trytes. Default: NULL. */
+  UT_array* ntru_array;
 } recv_mam_key_mam_v1_t;
 
 /**
  * Set the data ID for MAMv1
  *
- * @param[in] req Response data in type of ta_recv_mam_req_t object
- * @param[in] bundle_hash Bundle hash of the message
- * @param[in] chid Channel ID of the messages
- * @param[in] msg_id Message ID of the message
+ * @param req[in] Response data in type of ta_recv_mam_req_t object
+ * @param bundle_hash[in] Bundle hash of the message
+ * @param chid[in] Channel ID of the messages
+ * @param msg_id[in] Message ID of the message
  *
  * @return
  * - struct of ta_recv_mam_req_t on success
@@ -82,16 +83,27 @@ status_t recv_mam_set_mam_v1_data_id(ta_recv_mam_req_t* req, char* bundle_hash,
 /**
  * Set the key for MAMv1
  *
- * @param[in] req Response data in type of ta_recv_mam_req_t object
- * @param[in] psk Pre-Shared Key to decrypt message
- * @param[in] ntru NTRU public key to decrypt message
+ * @param req[in] Response data in type of ta_recv_mam_req_t object
+ * @param psk[in] Pre-Shared Key to decrypt message
+ * @param ntru[in] NTRU public key to decrypt message
  *
  * @return
  * - struct of ta_recv_mam_req_t on success
  * - NULL on error
  */
 status_t recv_mam_set_mam_v1_key(ta_recv_mam_req_t* req, tryte_t* psk, tryte_t* ntru);
 
+/**
+ * Initialize 'ta_recv_mam_req_t' object as MAMv1 object
+ *
+ * @param req[out] ta_recv_mam_req_t object to be initialized
+ *
+ * @return
+ * - struct of ta_recv_mam_req_t on success
+ * - NULL on error
+ */
+status_t recv_mam_req_v1_init(ta_recv_mam_req_t* req);
+
 #ifdef __cplusplus
 }
 #endif