1 files changed, 38 insertions, 45 deletions

diff --git a/src/lj_alloc.c b/src/lj_alloc.c
index 70ca1e3b..bf2ae847 100644
--- a/src/lj_alloc.c
+++ b/src/lj_alloc.c
@@ -31,6 +31,7 @@
 #include "lj_def.h"
 #include "lj_arch.h"
 #include "lj_alloc.h"
+#include "lj_prng.h"
 
 #ifndef LUAJIT_USE_SYSMALLOC
 
@@ -140,7 +141,7 @@ static void init_mmap(void)
 #define INIT_MMAP()     init_mmap()
 
 /* Win64 32 bit MMAP via NtAllocateVirtualMemory. */
-static void *CALL_MMAP(size_t size)
+static void *mmap_plain(size_t size)
 {
   DWORD olderr = GetLastError();
   void *ptr = NULL;
@@ -164,7 +165,7 @@ static void *direct_mmap(size_t size)
 #else
 
 /* Win32 MMAP via VirtualAlloc */
-static void *CALL_MMAP(size_t size)
+static void *mmap_plain(size_t size)
 {
   DWORD olderr = GetLastError();
   void *ptr = LJ_WIN_VALLOC(0, size, MEM_RESERVE|MEM_COMMIT, PAGE_READWRITE);
@@ -184,7 +185,8 @@ static void *direct_mmap(size_t size)
 
 #endif
 
-#define DIRECT_MMAP(size)       direct_mmap(size)
+#define CALL_MMAP(prng, size)   mmap_plain(size)
+#define DIRECT_MMAP(prng, size) direct_mmap(size)
 
 /* This function supports releasing coalesed segments */
 static int CALL_MUNMAP(void *ptr, size_t size)
@@ -228,30 +230,10 @@ static int CALL_MUNMAP(void *ptr, size_t size)
 
 #define LJ_ALLOC_MMAP_PROBE_LOWER       ((uintptr_t)0x4000)
 
-/* No point in a giant ifdef mess. Just try to open /dev/urandom.
-** It doesn't really matter if this fails, since we get some ASLR bits from
-** every unsuitable allocation, too. And we prefer linear allocation, anyway.
-*/
-#include <fcntl.h>
-#include <unistd.h>
-
-static uintptr_t mmap_probe_seed(void)
-{
-  uintptr_t val;
-  int fd = open("/dev/urandom", O_RDONLY);
-  if (fd != -1) {
-    int ok = ((size_t)read(fd, &val, sizeof(val)) == sizeof(val));
-    (void)close(fd);
-    if (ok) return val;
-  }
-  return 1;  /* Punt. */
-}
-
-static void *mmap_probe(size_t size)
+static void *mmap_probe(PRNGState *rs, size_t size)
 {
   /* Hint for next allocation. Doesn't need to be thread-safe. */
   static uintptr_t hint_addr = 0;
-  static uintptr_t hint_prng = 0;
   int olderr = errno;
   int retry;
   for (retry = 0; retry < LJ_ALLOC_MMAP_PROBE_MAX; retry++) {
@@ -283,15 +265,8 @@ static void *mmap_probe(size_t size)
       }
     }
     /* Finally, try pseudo-random probing. */
-    if (LJ_UNLIKELY(hint_prng == 0)) {
-      hint_prng = mmap_probe_seed();
-    }
-    /* The unsuitable address we got has some ASLR PRNG bits. */
-    hint_addr ^= addr & ~((uintptr_t)(LJ_PAGESIZE-1));
-    do {  /* The PRNG itself is very weak, but see above. */
-      hint_prng = hint_prng * 1103515245 + 12345;
-      hint_addr ^= hint_prng * (uintptr_t)LJ_PAGESIZE;
-      hint_addr &= (((uintptr_t)1 << LJ_ALLOC_MBITS)-1);
+    do {
+      hint_addr = lj_prng_u64(rs) & (((uintptr_t)1<<LJ_ALLOC_MBITS)-LJ_PAGESIZE);
     } while (hint_addr < LJ_ALLOC_MMAP_PROBE_LOWER);
   }
   errno = olderr;
@@ -308,12 +283,16 @@ static void *mmap_probe(size_t size)
 #define LJ_ALLOC_MMAP32_START   ((uintptr_t)0)
 #endif
 
+#if LJ_ALLOC_MMAP_PROBE
+static void *mmap_map32(PRNGState *rs, size_t size)
+#else
 static void *mmap_map32(size_t size)
+#endif
 {
 #if LJ_ALLOC_MMAP_PROBE
   static int fallback = 0;
   if (fallback)
-    return mmap_probe(size);
+    return mmap_probe(rs, size);
 #endif
   {
     int olderr = errno;
@@ -323,7 +302,7 @@ static void *mmap_map32(size_t size)
 #if LJ_ALLOC_MMAP_PROBE
     if (ptr == MFAIL) {
       fallback = 1;
-      return mmap_probe(size);
+      return mmap_probe(rs, size);
     }
 #endif
     return ptr;
@@ -333,17 +312,22 @@ static void *mmap_map32(size_t size)
 #endif
 
 #if LJ_ALLOC_MMAP32
-#define CALL_MMAP(size)         mmap_map32(size)
+#if LJ_ALLOC_MMAP_PROBE
+#define CALL_MMAP(prng, size)   mmap_map32(prng, size)
+#else
+#define CALL_MMAP(prng, size)   mmap_map32(size)
+#endif
 #elif LJ_ALLOC_MMAP_PROBE
-#define CALL_MMAP(size)         mmap_probe(size)
+#define CALL_MMAP(prng, size)   mmap_probe(prng, size)
 #else
-static void *CALL_MMAP(size_t size)
+static void *mmap_plain(size_t size)
 {
   int olderr = errno;
   void *ptr = mmap(NULL, size, MMAP_PROT, MMAP_FLAGS, -1, 0);
   errno = olderr;
   return ptr;
 }
+#define CALL_MMAP(prng, size)   mmap_plain(size)
 #endif
 
 #if LJ_64 && !LJ_GC64 && ((defined(__FreeBSD__) && __FreeBSD__ < 10) || defined(__FreeBSD_kernel__)) && !LJ_TARGET_PS4
@@ -396,7 +380,7 @@ static void *CALL_MREMAP_(void *ptr, size_t osz, size_t nsz, int flags)
 #endif
 
 #ifndef DIRECT_MMAP
-#define DIRECT_MMAP(s)          CALL_MMAP(s)
+#define DIRECT_MMAP(prng, s)    CALL_MMAP(prng, s)
 #endif
 
 #ifndef CALL_MREMAP
@@ -555,6 +539,7 @@ struct malloc_state {
   mchunkptr  smallbins[(NSMALLBINS+1)*2];
   tbinptr    treebins[NTREEBINS];
   msegment   seg;
+  PRNGState  *prng;
 };
 
 typedef struct malloc_state *mstate;
@@ -837,11 +822,11 @@ static int has_segment_link(mstate m, msegmentptr ss)
 
 /* -----------------------  Direct-mmapping chunks ----------------------- */
 
-static void *direct_alloc(size_t nb)
+static void *direct_alloc(mstate m, size_t nb)
 {
   size_t mmsize = mmap_align(nb + SIX_SIZE_T_SIZES + CHUNK_ALIGN_MASK);
   if (LJ_LIKELY(mmsize > nb)) {     /* Check for wrap around 0 */
-    char *mm = (char *)(DIRECT_MMAP(mmsize));
+    char *mm = (char *)(DIRECT_MMAP(m->prng, mmsize));
     if (mm != CMFAIL) {
       size_t offset = align_offset(chunk2mem(mm));
       size_t psize = mmsize - offset - DIRECT_FOOT_PAD;
@@ -853,6 +838,7 @@ static void *direct_alloc(size_t nb)
       return chunk2mem(p);
     }
   }
+  UNUSED(m);
   return NULL;
 }
 
@@ -1001,7 +987,7 @@ static void *alloc_sys(mstate m, size_t nb)
 
   /* Directly map large chunks */
   if (LJ_UNLIKELY(nb >= DEFAULT_MMAP_THRESHOLD)) {
-    void *mem = direct_alloc(nb);
+    void *mem = direct_alloc(m, nb);
     if (mem != 0)
       return mem;
   }
@@ -1010,7 +996,7 @@ static void *alloc_sys(mstate m, size_t nb)
     size_t req = nb + TOP_FOOT_SIZE + SIZE_T_ONE;
     size_t rsize = granularity_align(req);
     if (LJ_LIKELY(rsize > nb)) { /* Fail if wraps around zero */
-      char *mp = (char *)(CALL_MMAP(rsize));
+      char *mp = (char *)(CALL_MMAP(m->prng, rsize));
       if (mp != CMFAIL) {
         tbase = mp;
         tsize = rsize;
@@ -1237,12 +1223,13 @@ static void *tmalloc_small(mstate m, size_t nb)
 
 /* ----------------------------------------------------------------------- */
 
-void *lj_alloc_create(void)
+void *lj_alloc_create(PRNGState *rs)
 {
   size_t tsize = DEFAULT_GRANULARITY;
   char *tbase;
   INIT_MMAP();
-  tbase = (char *)(CALL_MMAP(tsize));
+  UNUSED(rs);
+  tbase = (char *)(CALL_MMAP(rs, tsize));
   if (tbase != CMFAIL) {
     size_t msize = pad_request(sizeof(struct malloc_state));
     mchunkptr mn;
@@ -1261,6 +1248,12 @@ void *lj_alloc_create(void)
   return NULL;
 }
 
+void lj_alloc_setprng(void *msp, PRNGState *rs)
+{
+  mstate ms = (mstate)msp;
+  ms->prng = rs;
+}
+
 void lj_alloc_destroy(void *msp)
 {
   mstate ms = (mstate)msp;