Merge pull request #609 from splasky/cipher_add_mac

jserv · web-flow · commit 3a4942f421bb · 2020-05-18T17:14:11.000+08:00
feat(endpoint): Add HMAC support for AES encryption
diff --git a/endpoint/endpoint.c b/endpoint/endpoint.c
@@ -10,6 +10,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <time.h>
 #include "http_parser.h"
 #include "utils/cipher.h"
 #include "utils/connectivity/conn_http.h"
@@ -47,14 +48,19 @@ status_t send_transaction_information(int value, const char* message, const char
   }
 
   size_t msg_len = 0;
+  struct timespec t;
+  clock_gettime(CLOCK_MONOTONIC, &t);
+
   ta_cipher_ctx encrypt_ctx = {.plaintext = raw_msg,
                                .plaintext_len = ret,
                                .ciphertext = ciphertext,
                                .ciphertext_len = MAX_MSG_LEN,
                                .device_id = device_id,
                                .iv = {0},
+                               .hmac = {0},
                                .key = private_key,
-                               .keybits = TA_AES_KEY_BITS};
+                               .keybits = TA_AES_KEY_BITS,
+                               .timestamp = t.tv_sec};
   memcpy(encrypt_ctx.iv, iv, AES_IV_SIZE);
   ret = aes_encrypt(&encrypt_ctx);
   memcpy(iv, encrypt_ctx.iv, AES_IV_SIZE);
@@ -64,7 +70,8 @@ status_t send_transaction_information(int value, const char* message, const char
     fprintf(stderr, "%s\n", "encrypt msg error");
     return ret;
   }
-  serialize_msg(iv, encrypt_ctx.ciphertext_len, (char*)encrypt_ctx.ciphertext, msg, &msg_len);
+  serialize_msg(iv, encrypt_ctx.ciphertext_len, (char*)encrypt_ctx.ciphertext, encrypt_ctx.timestamp, encrypt_ctx.hmac,
+                msg, &msg_len);
   bytes_to_trytes((const unsigned char*)msg, msg_len, tryte_msg);
 
   memset(req_body, 0, sizeof(char) * MAX_MSG_LEN);
diff --git a/tests/unit-test/test_cipher.c b/tests/unit-test/test_cipher.c
@@ -10,6 +10,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <time.h>
 #include "tests/test_define.h"
 #include "utils/cipher.h"
 
@@ -45,15 +46,19 @@ void tearDown(void) {}
 void test_cipher1(void) {
   uint8_t ciphertext[MAXLINE] = {0};
   uint8_t plaintext[MAXLINE] = {0};
+  struct timespec t;
+  clock_gettime(CLOCK_MONOTONIC, &t);
 
   ta_cipher_ctx encrypt_ctx = {.plaintext = (uint8_t*)test_payload1,
                                .plaintext_len = test_paylen1,
                                .ciphertext = ciphertext,
                                .ciphertext_len = MAXLINE,
                                .device_id = device_id,
                                .iv = {0},
+                               .hmac = {0},
                                .key = key,
-                               .keybits = TA_AES_KEY_BITS};
+                               .keybits = TA_AES_KEY_BITS,
+                               .timestamp = t.tv_sec};
   // Don't forget to set the initialization vector
   memcpy(encrypt_ctx.iv, iv_global, AES_IV_SIZE);
   TEST_ASSERT_EQUAL_INT32(SC_OK, aes_encrypt(&encrypt_ctx));
@@ -64,10 +69,14 @@ void test_cipher1(void) {
                                .ciphertext_len = encrypt_ctx.ciphertext_len,
                                .device_id = device_id,
                                .iv = {0},
+                               .hmac = {0},
                                .key = key,
-                               .keybits = TA_AES_KEY_BITS};
+                               .keybits = TA_AES_KEY_BITS,
+                               .timestamp = encrypt_ctx.timestamp};
   // Don't forget to set the initialization vector
   memcpy(decrypt_ctx.iv, encrypt_ctx.iv, AES_IV_SIZE);
+  // Don't forget to set the hmac
+  memcpy(decrypt_ctx.hmac, encrypt_ctx.hmac, TA_AES_HMAC_SIZE);
   TEST_ASSERT_EQUAL_INT32(SC_OK, aes_decrypt(&decrypt_ctx));
 
   TEST_ASSERT_EQUAL_INT(test_paylen1, decrypt_ctx.plaintext_len);
@@ -77,15 +86,18 @@ void test_cipher1(void) {
 void test_cipher2(void) {
   uint8_t ciphertext[MAXLINE] = {0};
   uint8_t plaintext[MAXLINE] = {0};
-
+  struct timespec t;
+  clock_gettime(CLOCK_MONOTONIC, &t);
   ta_cipher_ctx encrypt_ctx = {.plaintext = test_payload2,
                                .plaintext_len = test_paylen2,
                                .ciphertext = ciphertext,
                                .ciphertext_len = MAXLINE,
                                .device_id = device_id,
                                .iv = {0},
+                               .hmac = {0},
                                .key = key,
-                               .keybits = TA_AES_KEY_BITS};
+                               .keybits = TA_AES_KEY_BITS,
+                               .timestamp = t.tv_sec};
   // Don't forget to set the initialization vector
   memcpy(encrypt_ctx.iv, iv_global, AES_IV_SIZE);
   TEST_ASSERT_EQUAL_INT32(SC_OK, aes_encrypt(&encrypt_ctx));
@@ -96,10 +108,14 @@ void test_cipher2(void) {
                                .ciphertext_len = encrypt_ctx.ciphertext_len,
                                .device_id = device_id,
                                .iv = {0},
+                               .hmac = {0},
                                .key = key,
-                               .keybits = TA_AES_KEY_BITS};
+                               .keybits = TA_AES_KEY_BITS,
+                               .timestamp = encrypt_ctx.timestamp};
   // Don't forget to set the initialization vector
   memcpy(decrypt_ctx.iv, encrypt_ctx.iv, AES_IV_SIZE);
+  // Don't forget to set the hmac
+  memcpy(decrypt_ctx.hmac, encrypt_ctx.hmac, TA_AES_HMAC_SIZE);
   TEST_ASSERT_EQUAL_INT32(SC_OK, aes_decrypt(&decrypt_ctx));
 
   TEST_ASSERT_EQUAL_UINT(test_paylen2, decrypt_ctx.plaintext_len);
diff --git a/tests/unit-test/test_text_serializer.c b/tests/unit-test/test_text_serializer.c
@@ -28,27 +28,31 @@ const uint8_t payload[] = {
     49,  162, 1,   218, 50,  65,  239, 170, 29,  207, 210, 133, 167, 129, 150, 35,  165, 148, 255, 252, 131, 31,
     251, 91,  130, 34,  222, 70,  36,  45,  140, 85,  207, 141, 48,  1,   206, 31,  171, 235, 238, 126, 113};
 const uint16_t payload_len = 263;
+const uint64_t test_timestamp = 1589274915;
+const uint8_t test_hmac[TA_AES_HMAC_SIZE] = {37, 39, 48, 30, 53, 63, 45, 61, 87, 60, 60, 90, 63, 63, 13, 94,
+                                             53, 93, 52, 7,  19, 38, 15, 31, 29, 32, 97, 2,  47, 67, 59, 62};
 
 void setUp(void) {}
 
 void tearDown(void) {}
 
 void test_serialize_deserialize(void) {
   uint8_t out[1024], iv_out[AES_BLOCK_SIZE], payload_out[1024];
-
   size_t payload_len_out, out_msg_len;
-  int rc1 = serialize_msg(iv, payload_len, (char*)payload, (char*)out, &out_msg_len);
-  TEST_ASSERT(rc1 == SC_OK);
-  int rc2 = deserialize_msg((char*)out, iv_out, &payload_len_out, (char*)payload_out);
-  TEST_ASSERT(rc2 == SC_OK);
+  uint64_t timestamp;
+  uint8_t hmac[TA_AES_HMAC_SIZE];
+
+  int rc1 = serialize_msg(iv, payload_len, (char*)payload, test_timestamp, test_hmac, (char*)out, &out_msg_len);
+  TEST_ASSERT_EQUAL_INT32(SC_OK, rc1);
+  int rc2 = deserialize_msg((char*)out, iv_out, &payload_len_out, (char*)payload_out, &timestamp, hmac);
+  TEST_ASSERT_EQUAL_INT32(SC_OK, rc2);
 
   out[1023] = 0;
   payload_out[payload_len] = 0;
-
   TEST_ASSERT_EQUAL_UINT8_ARRAY(iv, iv_out, AES_BLOCK_SIZE);
-
+  TEST_ASSERT_EQUAL_UINT64(test_timestamp, timestamp);
+  TEST_ASSERT_EQUAL_UINT8_ARRAY(test_hmac, hmac, TA_AES_HMAC_SIZE);
   TEST_ASSERT_EQUAL_UINT32(payload_len, payload_len_out);
-
   TEST_ASSERT_EQUAL_UINT8_ARRAY(payload, payload_out, payload_len);
 }
 
diff --git a/utils/BUILD b/utils/BUILD
@@ -68,7 +68,10 @@ cc_library(
     name = "text_serializer",
     srcs = ["text_serializer.c"],
     hdrs = ["text_serializer.h"],
-    deps = ["//common:ta_errors"],
+    deps = [
+        ":cipher",
+        "//common:ta_errors",
+    ],
 )
 
 cc_library(
diff --git a/utils/cipher.c b/utils/cipher.c
@@ -20,21 +20,42 @@
 status_t aes_decrypt(ta_cipher_ctx* cipher_ctx) {
   // FIXME: Add logger and some checks here
   mbedtls_aes_context ctx;
+  mbedtls_md_context_t sha_ctx;
   int status;
-  char* err;
+  char* err = NULL;
   uint8_t buf[AES_BLOCK_SIZE];
-
+  uint8_t digest[AES_BLOCK_SIZE * 2];
+  uint8_t nonce[IMSI_LEN + MAX_TIMESTAMP_LEN + 1] = {0};
   /* Create and initialise the context */
   mbedtls_aes_init(&ctx);
+  mbedtls_md_init(&sha_ctx);
+  if (mbedtls_md_setup(&sha_ctx, mbedtls_md_info_from_type(MBEDTLS_MD_SHA256), 1) != 0) {
+    err = "Failed to set up message-digest information";
+    status = SC_UTILS_CIPHER_ERROR;
+    goto exit;
+  }
   mbedtls_platform_zeroize(cipher_ctx->plaintext, sizeof(cipher_ctx->plaintext));
   mbedtls_platform_zeroize(buf, AES_BLOCK_SIZE);
+  mbedtls_platform_zeroize(digest, AES_BLOCK_SIZE * 2);
 
   /* set decryption key */
   if ((status = mbedtls_aes_setkey_dec(&ctx, cipher_ctx->key, TA_AES_KEY_BITS)) != EXIT_SUCCESS) {
-    err = "set aes key failed";
+    err = "Failed to set AES key";
+    status = SC_UTILS_CIPHER_ERROR;
     goto exit;
   }
 
+  // concatenate (Device_ID, timestamp)
+  snprintf((char*)nonce, IMSI_LEN + MAX_TIMESTAMP_LEN + 1, "%s-%ld", cipher_ctx->device_id, cipher_ctx->timestamp);
+  // hash base data
+  mbedtls_md_starts(&sha_ctx);
+  mbedtls_md_update(&sha_ctx, digest, AES_BLOCK_SIZE * 2);
+  mbedtls_md_update(&sha_ctx, nonce, IMSI_LEN + MAX_TIMESTAMP_LEN);
+  mbedtls_md_update(&sha_ctx, cipher_ctx->key, TA_AES_KEY_BITS / 8);
+  mbedtls_md_finish(&sha_ctx, digest);
+
+  mbedtls_md_hmac_starts(&sha_ctx, digest, TA_AES_HMAC_SIZE);
+
   // Provide the message to be decrypted, and obtain the plaintext output.
   const size_t ciphertext_len = cipher_ctx->ciphertext_len;
   uint8_t* ciphertext = cipher_ctx->ciphertext;
@@ -43,21 +64,30 @@ status_t aes_decrypt(ta_cipher_ctx* cipher_ctx) {
     memset(buf, 0, AES_BLOCK_SIZE);
     int n = (ciphertext_len - i > AES_BLOCK_SIZE) ? AES_BLOCK_SIZE : (int)(ciphertext_len - i);
     memcpy(buf, ciphertext + i, n);
+    mbedtls_md_hmac_update(&sha_ctx, buf, AES_BLOCK_SIZE);
     if ((status = mbedtls_aes_crypt_cbc(&ctx, MBEDTLS_AES_DECRYPT, AES_BLOCK_SIZE, cipher_ctx->iv, buf, buf)) != 0) {
-      err = "aes decrpyt failed";
+      err = "Failed to decrypt AES message";
+      status = SC_UTILS_CIPHER_ERROR;
       goto exit;
     }
     memcpy(plaintext, buf, AES_BLOCK_SIZE);
     plaintext += AES_BLOCK_SIZE;
   }
 
-  /* Clean up */
-  mbedtls_aes_free(&ctx);
-  return SC_OK;
+  // compare hmac
+  mbedtls_md_hmac_finish(&sha_ctx, digest);
+  if (memcmp(digest, cipher_ctx->hmac, TA_AES_HMAC_SIZE) != 0) {
+    err = "Failed to validate HMAC";
+    status = SC_UTILS_CIPHER_ERROR;
+    goto exit;
+  }
+  status = SC_OK;
 exit:
-  fprintf(stderr, "%s\n", err);
+  // FIXME: Use default logger instead
+  if (!err) fprintf(stderr, "%s\n", err);
   mbedtls_aes_free(&ctx);
-  return SC_UTILS_CIPHER_ERROR;
+  mbedtls_md_free(&sha_ctx);
+  return status;
 }
 
 status_t aes_encrypt(ta_cipher_ctx* cipher_ctx) {
@@ -79,29 +109,29 @@ status_t aes_encrypt(ta_cipher_ctx* cipher_ctx) {
   mbedtls_md_init(&sha_ctx);
   mbedtls_aes_init(&ctx);
   if (mbedtls_md_setup(&sha_ctx, mbedtls_md_info_from_type(MBEDTLS_MD_SHA256), 1) != 0) {
-    err = "mbedtls_md_setup error";
+    err = "Failed to set up message-digest information";
     goto exit;
   }
 
   // Check ciphertext has enough space
   size_t new_len = plaintext_len + (AES_BLOCK_SIZE - plaintext_len % 16);
   if (new_len > ciphertext_len) {
-    err = "ciphertext has not enough space";
+    err = "Failed to get enough space inside ciphertext buffer";
+    status = SC_UTILS_CIPHER_ERROR;
     goto exit;
   }
   cipher_ctx->ciphertext_len = new_len;
   mbedtls_platform_zeroize(tmp, sizeof(tmp));
   mbedtls_platform_zeroize(digest, sizeof(digest));
   mbedtls_platform_zeroize(ciphertext, sizeof(ciphertext));
 
-  // fetch timestamp
-  uint64_t timestamp = time(NULL);
   // concatenate (Device_ID, timestamp)
-  snprintf((char*)nonce, IMSI_LEN + MAX_TIMESTAMP_LEN + 1, "%s-%ld", cipher_ctx->device_id, timestamp);
+  snprintf((char*)nonce, IMSI_LEN + MAX_TIMESTAMP_LEN + 1, "%s-%ld", cipher_ctx->device_id, cipher_ctx->timestamp);
   // hash base data
   mbedtls_md_starts(&sha_ctx);
   mbedtls_md_update(&sha_ctx, digest, AES_BLOCK_SIZE * 2);
   mbedtls_md_update(&sha_ctx, nonce, IMSI_LEN + MAX_TIMESTAMP_LEN);
+  mbedtls_md_update(&sha_ctx, cipher_ctx->key, TA_AES_KEY_BITS / 8);
   mbedtls_md_finish(&sha_ctx, digest);
 
   for (int i = 0; i < AES_BLOCK_SIZE; ++i) {
@@ -111,7 +141,14 @@ status_t aes_encrypt(ta_cipher_ctx* cipher_ctx) {
 
   /* set encryption key */
   if ((status = mbedtls_aes_setkey_enc(&ctx, cipher_ctx->key, TA_AES_KEY_BITS)) != 0) {
-    err = "set aes key failed";
+    err = "Failed to set AES key";
+    status = SC_UTILS_CIPHER_ERROR;
+    goto exit;
+  }
+
+  if ((status = mbedtls_md_hmac_starts(&sha_ctx, digest, TA_AES_HMAC_SIZE)) != 0) {
+    err = "Failed to initialize HMAC context";
+    status = SC_UTILS_CIPHER_ERROR;
     goto exit;
   }
 
@@ -121,19 +158,22 @@ status_t aes_encrypt(ta_cipher_ctx* cipher_ctx) {
     int n = (plaintext_len - i > AES_BLOCK_SIZE) ? AES_BLOCK_SIZE : (int)(plaintext_len - i);
     memcpy(buf, plaintext + i, n);
     if ((status = mbedtls_aes_crypt_cbc(&ctx, MBEDTLS_AES_ENCRYPT, AES_BLOCK_SIZE, tmp, buf, buf)) != 0) {
-      err = "aes decrpyt failed";
+      err = "Failed to encrypt AES message";
+      status = SC_UTILS_CIPHER_ERROR;
       goto exit;
     }
+    mbedtls_md_hmac_update(&sha_ctx, buf, AES_BLOCK_SIZE);
     memcpy(ciphertext, buf, AES_BLOCK_SIZE);
     ciphertext += AES_BLOCK_SIZE;
   }
 
-  mbedtls_aes_free(&ctx);
-  mbedtls_md_free(&sha_ctx);
-  return SC_OK;
+  mbedtls_md_hmac_finish(&sha_ctx, digest);
+  memcpy(cipher_ctx->hmac, digest, TA_AES_HMAC_SIZE);
+  status = SC_OK;
 exit:
-  fprintf(stderr, "%s", err);
+  // FIXME: Use default logger instead
+  if (!err) fprintf(stderr, "%s\n", err);
   mbedtls_aes_free(&ctx);
   mbedtls_md_free(&sha_ctx);
-  return SC_UTILS_CIPHER_ERROR;
+  return status;
 }
diff --git a/utils/cipher.h b/utils/cipher.h
@@ -22,17 +22,20 @@ extern "C" {
 #define AES_IV_SIZE AES_BLOCK_SIZE
 #define IMSI_LEN 16
 #define TA_AES_KEY_BITS 256
+#define TA_AES_HMAC_SIZE AES_BLOCK_SIZE * 2
 
 /** context of aes cipher */
 typedef struct ta_cipher_ctx {
-  uint8_t* plaintext;      /**< Plaintext */
-  size_t plaintext_len;    /**< Plaintext length */
-  uint8_t* ciphertext;     /**<  Ciphrtext */
-  size_t ciphertext_len;   /**< Ciphertext length */
-  uint8_t iv[AES_IV_SIZE]; /**< Initialization vector, mbedtls_aes needs r/w iv[] */
-  const uint8_t* key;      /**< Encryption key */
-  const size_t keybits;    /**< Bits of key, valid options are 128,192,256 bits */
-  const char* device_id;   /**< Device id */
+  uint8_t* plaintext;             /**< Plaintext */
+  size_t plaintext_len;           /**< Plaintext length */
+  uint8_t* ciphertext;            /**< Ciphrtext */
+  size_t ciphertext_len;          /**< Ciphertext length */
+  uint8_t iv[AES_IV_SIZE];        /**< Initialization vector, mbedtls_aes needs r/w iv[] */
+  const uint8_t* key;             /**< Encryption key */
+  const size_t keybits;           /**< Bits of key, valid options are 128,192,256 bits */
+  const char* device_id;          /**< Device id */
+  uint8_t hmac[TA_AES_HMAC_SIZE]; /**< r/w buffer for hash-based message authentication code */
+  uint64_t timestamp;             /**< Timestamp for generate hmac */
 } ta_cipher_ctx;
 
 /**
diff --git a/utils/text_serializer.c b/utils/text_serializer.c
diff --git a/utils/text_serializer.h b/utils/text_serializer.h
