From 768aa8a4c36d41af48780409a3e565adac59d51e Mon Sep 17 00:00:00 2001
From: Simon Maddocks <git@51m0n.com>
Date: Thu, 27 Jan 2022 00:43:24 +0000
Subject: [PATCH] Add key setting

---
 Romulus-M/encrypt_wrapper.c | 23 ++++++++++++-----------
 Romulus-M/encrypt_wrapper.h |  4 ++--
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/Romulus-M/encrypt_wrapper.c b/Romulus-M/encrypt_wrapper.c
index 9674022..8f8580c 100644
--- a/Romulus-M/encrypt_wrapper.c
+++ b/Romulus-M/encrypt_wrapper.c
@@ -11,13 +11,12 @@
 #include "api.h"
 
 
-int encrypt(packet *packet, unsigned char* outBuf, __uint16_t *len){
+int encrypt(packet *packet, unsigned char* outBuf, __uint16_t *len, unsigned char* key){
     if(sizeof(u_int16_t) + sizeof(u_int16_t) + packet->Size > *len){
         return 1;
     }
 
     int packetSize = packet->Size; 
-
     // enc
     // ciphertext - out, MAX 16 bytes larger than plaintext - defined as CRYPTO_ABYTES
     unsigned long long clen = packetSize + CRYPTO_ABYTES;
@@ -35,9 +34,7 @@ int encrypt(packet *packet, unsigned char* outBuf, __uint16_t *len){
     unsigned char npub[CRYPTO_NPUBBYTES];
     getrandom(&npub[0], sizeof(npub), 0);
 
-    // key, CRYPTO_KEYBYTES size 16
-    unsigned char k[CRYPTO_KEYBYTES] = {0};
-    int ret = romulus_m_encrypt(c, &clen, m, mlen, ad, adlen, 0, &npub[0], &k[0]);
+    int ret = romulus_m_encrypt(c, &clen, m, mlen, ad, adlen, 0, &npub[0], key);
       
     if(clen > *len + CRYPTO_NPUBBYTES || ret != 0){
         free(c);
@@ -72,8 +69,11 @@ int encrypt(packet *packet, unsigned char* outBuf, __uint16_t *len){
     return 0;
 }
 
-int decrypt(packetEx *packet, unsigned char* inBuf, __uint16_t len){
+int decrypt(packetEx *packet, unsigned char* inBuf, __uint16_t len, unsigned char* key){
     int loc = 0;
+    if(inBuf == 0 ){
+        return 1;
+    }
     memcpy(&packet->DataType, &inBuf[loc], sizeof(u_int16_t));
     loc += sizeof(u_int16_t);
     memcpy(&packet->Id, &inBuf[loc], sizeof(u_int32_t));
@@ -85,6 +85,10 @@ int decrypt(packetEx *packet, unsigned char* inBuf, __uint16_t len){
     packet->Size = __bswap_16(packet->Size);
     packet->Id = __bswap_32(packet->Id);
 
+    if(packet->Size < CRYPTO_NPUBBYTES){
+        return 1;
+    }
+
     // Data
     // Nonce|Data
     // dec
@@ -104,17 +108,14 @@ int decrypt(packetEx *packet, unsigned char* inBuf, __uint16_t len){
     int adlen = sizeof(packet->DataType);
     unsigned char* ad = (unsigned char*)&packet->DataType;
 
-    // key, CRYPTO_KEYBYTES size 16
-    unsigned char k[CRYPTO_KEYBYTES] = {0};
-
-    int ret = romulus_m_decrypt(m, &mlen, 0, c, clen, ad, adlen, npub, &k[0]);   
+    int ret = romulus_m_decrypt(m, &mlen, 0, c, clen, ad, adlen, npub, key);   
     if(mlen <= 1000 && ret == 0){
         memcpy(packet->Data, m, mlen);
         packet->Size = (u_int16_t)mlen;
     }
     else{
         free(m);
-        return 1;
+        return -1;
     }
     free(m);
     return 0;
diff --git a/Romulus-M/encrypt_wrapper.h b/Romulus-M/encrypt_wrapper.h
index 6f0f8ac..6c637d3 100644
--- a/Romulus-M/encrypt_wrapper.h
+++ b/Romulus-M/encrypt_wrapper.h
@@ -14,7 +14,7 @@ typedef struct {
     unsigned char* Data;
 }packetEx;
 
-EXPORT int encrypt(packet *packet, unsigned char* outBuf, __uint16_t *len);
+EXPORT int encrypt(packet *packet, unsigned char* outBuf, __uint16_t *len, unsigned char* key);
 
 
-EXPORT int decrypt(packetEx *packet, unsigned char* inBuf, __uint16_t len);
\ No newline at end of file
+EXPORT int decrypt(packetEx *packet, unsigned char* inBuf, __uint16_t len, unsigned char* key);
\ No newline at end of file