1 files changed, 98 insertions, 40 deletions

diff --git a/src/noise.c b/src/noise.c
index 83f2e4c..09a4c47 100644
--- a/src/noise.c
+++ b/src/noise.c
@@ -19,11 +19,22 @@
 volatile uint8_t host_packet_buf[MAX_HOST_PACKET_SIZE];
 volatile uint8_t host_packet_length = 0;
 
-static uint8_t local_key[CURVE25519_KEY_LEN];
-static NoiseCipherState *tx_cipher = NULL, *rx_cipher = NULL;
 
+void noise_state_init(struct NoiseState *st, uint8_t *remote_key_reference) {
+    st->handshake_state = HANDSHAKE_UNINITIALIZED;
+    st->handshake = NULL;
+    st->tx_cipher = NULL;
+    st->rx_cipher = NULL;
+    st->remote_key_reference = remote_key_reference;
+    st->failed_handshakes = 0;
+}
+
+int reset_protocol_handshake(struct NoiseState *st) {
+    uninit_handshake(st, HANDSHAKE_UNINITIALIZED);
+    return start_protocol_handshake(st);
+}
 
-NoiseHandshakeState *start_protocol_handshake() {
+int start_protocol_handshake(struct NoiseState *st) {
     /* TODO Noise-C is nice for prototyping, but we should really get rid of it for mostly three reasons:
      *   * We don't need cipher/protocol agility, and by baking the final protocol into the firmware we can save a lot
      *     of flash space by not including all the primitives we don't need as well as noise's dynamic protocol
@@ -40,18 +51,20 @@ NoiseHandshakeState *start_protocol_handshake() {
     HANDLE_NOISE_ERROR(noise_handshakestate_new_by_name(&handshake, "Noise_XX_25519_ChaChaPoly_BLAKE2s", NOISE_ROLE_RESPONDER), "instantiating handshake pattern");
 
     NoiseDHState *dh = noise_handshakestate_get_local_keypair_dh(handshake);
-    HANDLE_NOISE_ERROR(noise_dhstate_set_keypair_private(dh, local_key, sizeof(local_key)), "loading local private keys");
+    HANDLE_NOISE_ERROR(noise_dhstate_set_keypair_private(dh, st->local_key, sizeof(st->local_key)), "loading local private keys");
 
     HANDLE_NOISE_ERROR(noise_handshakestate_start(handshake), "starting handshake");
 
-    return handshake;
+    st->handshake = handshake;
+    st->handshake_state = HANDSHAKE_IN_PROGRESS;
+    return 0;
 
 errout:
     noise_handshakestate_free(handshake);
-    return 0;
+    return -1;
 }
 
-int generate_identity_key() {
+int generate_identity_key(struct NoiseState *st) {
     NoiseDHState *dh;
     int err;
 
@@ -59,84 +72,129 @@ int generate_identity_key() {
     HANDLE_NOISE_ERROR(noise_dhstate_generate_keypair(dh), "generating key pair");
 
     uint8_t unused[CURVE25519_KEY_LEN]; /* the noise api is a bit bad here. */
-    memset(local_key, 0, sizeof(local_key));
+    memset(st->local_key, 0, sizeof(st->local_key));
 
-    HANDLE_NOISE_ERROR(noise_dhstate_get_keypair(dh, local_key, sizeof(local_key), unused, sizeof(unused)), "saving key pair");
+    HANDLE_NOISE_ERROR(noise_dhstate_get_keypair(dh, st->local_key, sizeof(st->local_key), unused, sizeof(unused)), "saving key pair");
 
     return 0;
-
 errout:
     if (dh)
         noise_dhstate_free(dh);
     return -1;
 }
 
-NoiseHandshakeState *try_continue_noise_handshake(NoiseHandshakeState *handshake) {
+void uninit_handshake(struct NoiseState *st, enum handshake_state new_state) {
+    if (st->handshake)
+        noise_handshakestate_free(st->handshake);
+    st->handshake_state = new_state;
+    st->handshake = NULL;
+}
+
+enum handshake_state try_continue_noise_handshake(struct NoiseState *st, uint8_t *buf, size_t len, int *buf_consumed) {
     int err;
-    uint8_t message[MAX_HOST_PACKET_SIZE];
+    struct {
+        struct control_packet header;
+        uint8_t payload[MAX_HOST_PACKET_SIZE];
+    } pkt;
     NoiseBuffer noise_msg;
+
+    if (!st->handshake || st->handshake_state != HANDSHAKE_IN_PROGRESS) {
+        LOG_PRINTF("Error: Invalid handshake state\n");
+        goto errout;
+    }
+
     /* Run the protocol handshake */
-    switch (noise_handshakestate_get_action(handshake)) {
+    switch (noise_handshakestate_get_action(st->handshake)) {
     case NOISE_ACTION_WRITE_MESSAGE:
-        /* Write the next handshake message with a zero-length payload */
-        noise_buffer_set_output(noise_msg, message, sizeof(message));
-        HANDLE_NOISE_ERROR(noise_handshakestate_write_message(handshake, &noise_msg, NULL), "writing handshake message");
-        send_packet(usart2_out, message, noise_msg.size);
+        /* Write the next handshake message with a zero-length noise payload */
+        pkt.header.type = HOST_HANDSHAKE;
+        noise_buffer_set_output(noise_msg, &pkt.payload, sizeof(pkt.payload));
+        HANDLE_NOISE_ERROR(noise_handshakestate_write_message(st->handshake, &noise_msg, NULL), "writing handshake message");
+        send_packet(usart2_out, (uint8_t *)&pkt, noise_msg.size + sizeof(pkt.header));
         break;
 
     case NOISE_ACTION_READ_MESSAGE:
-        if (host_packet_length > 0) {
+        if (buf) {
             /* Read the next handshake message and discard the payload */
-            noise_buffer_set_input(noise_msg, (uint8_t *)host_packet_buf, host_packet_length);
-            HANDLE_NOISE_ERROR(noise_handshakestate_read_message(handshake, &noise_msg, NULL), "reading handshake message");
-            host_packet_length = 0; /* Acknowledge to USART ISR the buffer has been handled */
+            *buf_consumed = 1;
+            noise_buffer_set_input(noise_msg, buf, len);
+            HANDLE_NOISE_ERROR(noise_handshakestate_read_message(st->handshake, &noise_msg, NULL), "reading handshake message");
         }
         break;
 
     case NOISE_ACTION_SPLIT:
-        HANDLE_NOISE_ERROR(noise_handshakestate_split(handshake, &tx_cipher, &rx_cipher), "splitting handshake state");
+        HANDLE_NOISE_ERROR(noise_handshakestate_split(st->handshake, &st->tx_cipher, &st->rx_cipher), "splitting handshake state");
         LOG_PRINTF("Noise protocol handshake completed successfully, handshake hash:\n");
 
-        uint8_t buf[BLAKE2S_HASH_SIZE];
-        if (noise_handshakestate_get_handshake_hash(handshake, buf, sizeof(buf)) != NOISE_ERROR_NONE) {
+        if (noise_handshakestate_get_handshake_hash(st->handshake, st->handshake_hash, sizeof(st->handshake_hash)) != NOISE_ERROR_NONE) {
             LOG_PRINTF("Error fetching noise handshake state\n");
         } else {
             LOG_PRINTF("    ");
-            for (size_t i=0; i<sizeof(buf); i++)
-                LOG_PRINTF("%02x ", buf[i]);
+            for (size_t i=0; i<sizeof(st->handshake_hash); i++)
+                LOG_PRINTF("%02x ", st->handshake_hash[i]);
             LOG_PRINTF("\n");
         }
 
-        noise_handshakestate_free(handshake);
-        return NULL;
+        
+        NoiseDHState *remote_dh = noise_handshakestate_get_remote_public_key_dh(st->handshake);
+        if (!remote_dh) {
+            LOG_PRINTF("Error: Host has not identified itself\n");
+            goto errout;
+        }
+
+
+        HANDLE_NOISE_ERROR(noise_dhstate_get_public_key(remote_dh, st->remote_key, sizeof(st->remote_key)), "getting remote pubkey");
+
+        if (!memcmp(st->remote_key, st->remote_key_reference, sizeof(st->remote_key))) { /* keys match */
+            uninit_handshake(st, HANDSHAKE_DONE_KNOWN_HOST);
+            st->failed_handshakes = 0;
+        } else { /* keys don't match */
+            uninit_handshake(st, HANDSHAKE_DONE_UNKNOWN_HOST);
+            st->failed_handshakes++;
+        }
+        break;
 
     default:
-        LOG_PRINTF("Noise protocol handshake failed\n");
         goto errout;
     }
 
-    return handshake;
-
+    return st->handshake_state;
 errout:
-    noise_handshakestate_free(handshake);
-    return NULL;
+    uninit_handshake(st, HANDSHAKE_UNINITIALIZED);
+    st->failed_handshakes++;
+    LOG_PRINTF("Noise protocol handshake failed, %d failed attempts\n", st->failed_handshakes);
+    return st->handshake_state;
+}
+
+void persist_remote_key(struct NoiseState *st) {
+    memcpy(st->remote_key_reference, st->remote_key, sizeof(st->remote_key));
+    st->handshake_state = HANDSHAKE_DONE_KNOWN_HOST;
 }
 
-int send_encrypted_message(uint8_t *msg, size_t len) {
+int send_encrypted_message(struct NoiseState *st, uint8_t *msg, size_t len) {
     int err;
     NoiseBuffer noise_buf;
-    uint8_t raw_buf[MAX_HOST_PACKET_SIZE];
+    struct {
+        struct control_packet header;
+        uint8_t payload[MAX_HOST_PACKET_SIZE];
+    } pkt;
 
-    if (!tx_cipher) {
+    if (!st->tx_cipher) {
         LOG_PRINTF("Cannot send encrypted packet: Data ciphers not yet initialized\n");
         return -1;
     }
 
-    memcpy(raw_buf, msg, len); /* This is necessary because noises API doesn't support separate in and out buffers. D'oh! */
-    noise_buffer_set_inout(noise_buf, raw_buf, len, sizeof(raw_buf));
+    if (len > sizeof(pkt.payload)) {
+        LOG_PRINTF("Packet too long\n");
+        return -3;
+    }
+
+    pkt.header.type = HOST_DATA;
+    memcpy(pkt.payload, msg, len); /* This is necessary because noises API doesn't support separate in and out buffers. D'oh! */
+    noise_buffer_set_inout(noise_buf, pkt.payload, len, sizeof(pkt.payload));
 
-    HANDLE_NOISE_ERROR(noise_cipherstate_encrypt(tx_cipher, &noise_buf), "encrypting data");
-    send_packet(usart2_out, raw_buf, noise_buf.size);
+    HANDLE_NOISE_ERROR(noise_cipherstate_encrypt(st->tx_cipher, &noise_buf), "encrypting data");
+    send_packet(usart2_out, (uint8_t *)&pkt, noise_buf.size + sizeof(pkt.header));
 
     return 0;
 errout: