From 0ac5b243c723f079dc88325c5be4be8dfd9d324f Mon Sep 17 00:00:00 2001
From: Andrew Chambers <ac@acha.ninja>
Date: Wed, 4 Dec 2019 11:29:11 +1300
Subject: [PATCH] Add os/cryptorand.

---
 src/core/os.c       | 70 +++++++++++++++++++++++++++++++++++++++++++--
 src/include/janet.h |  6 +++-
 test/suite7.janet   | 22 ++++++++++++++
 tools/format.sh     |  2 +-
 4 files changed, 95 insertions(+), 5 deletions(-)

diff --git a/src/core/os.c b/src/core/os.c
index 65479f7f..dc617996 100644
--- a/src/core/os.c
+++ b/src/core/os.c
@@ -25,8 +25,6 @@
 #include "util.h"
 #endif
 
-#include <stdlib.h>
-
 #ifndef JANET_REDUCED_OS
 
 #include <time.h>
@@ -36,6 +34,8 @@
 #include <string.h>
 #include <sys/stat.h>
 
+#define RETRY_EINTR(RC, CALL) do { (RC) = CALL; } while((RC) < 0 && errno == EINTR)
+
 #ifdef JANET_WINDOWS
 #include <windows.h>
 #include <direct.h>
@@ -473,12 +473,13 @@ static Janet os_sleep(int32_t argc, Janet *argv) {
 #ifdef JANET_WINDOWS
     Sleep((DWORD)(delay * 1000));
 #else
+    int rc;
     struct timespec ts;
     ts.tv_sec = (time_t) delay;
     ts.tv_nsec = (delay <= UINT32_MAX)
                  ? (long)((delay - ((uint32_t)delay)) * 1000000000)
                  : 0;
-    nanosleep(&ts, NULL);
+    RETRY_EINTR(rc, nanosleep(&ts, &ts));
 #endif
     return janet_wrap_nil();
 }
@@ -497,6 +498,64 @@ static Janet os_cwd(int32_t argc, Janet *argv) {
     return janet_cstringv(ptr);
 }
 
+static Janet os_cryptorand(int32_t argc, Janet *argv) {
+    JanetBuffer *buffer;
+    const char *genericerr = "unable to get sufficient random data";
+    janet_arity(argc, 1, 2);
+    int32_t offset;
+    int32_t n = janet_getinteger(argv, 0);
+    if (n < 0) janet_panic("expected positive integer");
+    if (argc == 2) {
+        buffer = janet_getbuffer(argv, 1);
+        offset = buffer->count;
+    } else {
+        offset = 0;
+        buffer = janet_buffer(n);
+    }
+    /* We could optimize here by adding setcount_uninit */
+    janet_buffer_setcount(buffer, offset + n);
+
+#ifdef JANET_WINDOWS
+    for (int32_t i = offset; i < buffer->count; i += sizeof(unsigned int)) {
+        unsigned int v;
+        if (rand_s(&v))
+            janet_panic(genericerr);
+        for (int32_t j = 0; (j < sizeof(unsigned int)) && (i + j < buffer->count); j++) {
+            buffer->data[i + j] = v & 0xff;
+            v = v >> 8;
+        }
+    }
+#elif defined(__linux__) || defined(__APPLE__)
+    /* We should be able to call getrandom on linux, but it doesn't seem
+       to be uniformly supported on linux distros. Macos may support
+       arc4random_buf, but it needs investigation.
+
+       In both cases, use this fallback path for now... */
+    int rc;
+    int randfd;
+    RETRY_EINTR(randfd, open("/dev/urandom", O_RDONLY));
+    if (randfd < 0)
+        janet_panic(genericerr);
+    while (n > 0) {
+        ssize_t nread;
+        RETRY_EINTR(nread, read(randfd, buffer->data + offset, n));
+        if (nread <= 0) {
+            RETRY_EINTR(rc, close(randfd));
+            janet_panic(genericerr);
+        }
+        offset += nread;
+        n -= nread;
+    }
+    RETRY_EINTR(rc, close(randfd));
+#elif defined(__FreeBSD__) || defined(__OpenBSD__) || defined(__NetBSD__)
+    (void) errmsg;
+    arc4random_buf(buffer->data + offset, n);
+#else
+    janet_panic("cryptorand currently unsupported on this platform");
+#endif
+    return janet_wrap_buffer(buffer);
+}
+
 static Janet os_date(int32_t argc, Janet *argv) {
     janet_arity(argc, 0, 2);
     (void) argv;
@@ -981,6 +1040,11 @@ static const JanetReg os_cfuns[] = {
         JDOC("(os/cwd)\n\n"
              "Returns the current working directory.")
     },
+    {
+        "os/cryptorand", os_cryptorand,
+        JDOC("(os/cryptorand n &opt buf)\n\n"
+             "Get or append n bytes of good quality random data provided by the os. Returns a new buffer or buf.")
+    },
     {
         "os/date", os_date,
         JDOC("(os/date &opt time local)\n\n"
diff --git a/src/include/janet.h b/src/include/janet.h
index 02219c8d..e0bd5027 100644
--- a/src/include/janet.h
+++ b/src/include/janet.h
@@ -228,9 +228,13 @@ typedef struct {
 
 /***** START SECTION TYPES *****/
 
+#ifdef JANET_WINDOWS
+// Must be defined before including stdlib.h
+#define _CRT_RAND_S
+#endif
+#include <stdlib.h>
 #include <stdint.h>
 #include <string.h>
-#include <stdlib.h>
 #include <stdarg.h>
 #include <setjmp.h>
 #include <stddef.h>
diff --git a/test/suite7.janet b/test/suite7.janet
index 430951a1..3b1f6957 100644
--- a/test/suite7.janet
+++ b/test/suite7.janet
@@ -238,4 +238,26 @@
 # Issue #183 - just parse it :)
 1e-4000000000000000000000
 
+# Ensure randomness puts n of pred into our buffer eventually
+(defn cryptorand-check
+  [n pred]
+  (def max-attempts 10000)
+  (var attempts 0)
+  (while (not= attempts max-attempts)
+    (def cryptobuf (os/cryptorand 10))
+    (when (= n (count pred cryptobuf))
+      (break))
+    (++ attempts))
+  (not= attempts max-attempts))
+
+(def v (math/rng-int (math/rng (os/time)) 100))
+(assert (cryptorand-check 0 |(= $ v)) "cryptorand skips value sometimes")
+(assert (cryptorand-check 1 |(= $ v)) "cryptorand has value sometimes")
+
+(do 
+  (def buf (buffer/new-filled 1))
+  (os/cryptorand 1 buf)
+  (assert (= (in buf 0) 0) "cryptorand doesn't overwrite buffer")
+  (assert (= (length buf) 2) "cryptorand appends to buffer"))
+
 (end-suite)
diff --git a/tools/format.sh b/tools/format.sh
index 8d2ba436..f5b4a88c 100755
--- a/tools/format.sh
+++ b/tools/format.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 
 # Format all code with astyle