From 7ee8da5d3e81ea91d7ed49cdcda3f6b87ca39445 Mon Sep 17 00:00:00 2001
From: Simon Maddocks <git@51m0n.com>
Date: Sun, 6 Feb 2022 03:22:15 +0000
Subject: [PATCH] Catch undersize bytes and panics

---
 romulus_m.go | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/romulus_m.go b/romulus_m.go
index ae1886d..bccd91e 100644
--- a/romulus_m.go
+++ b/romulus_m.go
@@ -11,6 +11,10 @@ const noncesize int = 16
 
 func Encrypt(key []byte, plaintext []byte, additionalData []byte) ([]byte, error) {
 
+	if len(key) < keysize {
+		return nil, errors.New("Failed to encrypt")
+	}
+
 	cipher := make([]byte, len(plaintext)+abytes+noncesize)
 	var cipherlen uint64 = (uint64)(len(cipher))
 	nonce := genNonce()
@@ -27,15 +31,25 @@ func Encrypt(key []byte, plaintext []byte, additionalData []byte) ([]byte, error
 	)
 
 	if ret != 0 {
-		return nil, errors.New("Failed to decrypt")
+		return nil, errors.New("Failed to encrypt")
 	}
 	copy(cipher, nonce)
 	return cipher[:(int)(cipherlen)+noncesize], nil
 }
 
-func Decrypt(key []byte, ciphertext []byte, additionalData []byte) ([]byte, bool, error) {
+func Decrypt(key []byte, ciphertext []byte, additionalData []byte) (plaintext []byte, auth bool, err error) {
 
-	plaintext := make([]byte, len(ciphertext))
+	if len(ciphertext) <= noncesize || len(key) < keysize {
+		return nil, false, errors.New("Failed to decrypt")
+	}
+
+	defer func() {
+		recover()
+		err = errors.New("Recoverd from panic in decrypt")
+		auth = false
+	}()
+
+	plaintext = make([]byte, len(ciphertext))
 	var plaintextLen uint64 = (uint64)(len(plaintext))
 	ret := romulus_m_decrypt(
 		plaintext,
@@ -48,6 +62,7 @@ func Decrypt(key []byte, ciphertext []byte, additionalData []byte) ([]byte, bool
 		ciphertext[:noncesize],
 		key[:keysize],
 	)
+
 	if ret == -1 {
 		return nil, false, errors.New("Failed to authenticate")
 	} else if ret != 0 {