1 files changed, 16 insertions, 8 deletions
diff --git a/fw/fe310/eos/net/rng.c b/fw/fe310/eos/net/rng.c
index 7d05a81..414a6fa 100644
--- a/fw/fe310/eos/net/rng.c
+++ b/fw/fe310/eos/net/rng.c
@@ -2,26 +2,34 @@
 #include <stdint.h>
 #include <string.h>
 
+#include "eos.h"
+#include "event.h"
 #include "dev/net.h"
 
 int getentropy(unsigned char *b, size_t sz) {
     unsigned char type;
-    unsigned char *buffer;
+    EOSMessage msg;
     uint16_t len;
     int rv;
 
-    buffer = eos_net_alloc();
-
     type = EOS_NET_MTYPE_RNG;
     len = sizeof(uint16_t);
-    buffer[0] = sz >> 8;
-    buffer[1] = sz;
 
-    rv = eos_net_xchg(&type, buffer, &len);
+    eos_net_alloc(&msg);
+
+    if ((msg.size < len) || (msg.size < sz)) {
+        eos_net_free(&msg, 1);
+        return -1;
+    }
+
+    msg.buffer[0] = sz >> 8;
+    msg.buffer[1] = sz;
+
+    rv = eos_net_xchg(&type, &msg, &len);
     if (rv || (len != sz)) rv = -1;
 
-    if (!rv) memcpy(b, buffer, sz);
-    eos_net_free(buffer, 1);
+    if (!rv) memcpy(b, msg.buffer, sz);
+    eos_net_free(&msg, 1);
 
     return rv;
 }