Author: Armin Rigo <ar...@tunes.org>
Branch: 
Changeset: r89089:82ef21124af8
Date: 2016-12-16 10:42 +0100
http://bitbucket.org/pypy/pypy/changeset/82ef21124af8/

Log:    Linux: try to implement os.urandom() on top of the syscall
        getrandom()

diff --git a/pypy/module/posix/interp_posix.py 
b/pypy/module/posix/interp_posix.py
--- a/pypy/module/posix/interp_posix.py
+++ b/pypy/module/posix/interp_posix.py
@@ -1331,8 +1331,9 @@
     Return a string of n random bytes suitable for cryptographic use.
     """
     context = get(space).random_context
+    signal_checker = space.getexecutioncontext().checksignals
     try:
-        return space.wrap(rurandom.urandom(context, n))
+        return space.wrap(rurandom.urandom(context, n, signal_checker))
     except OSError as e:
         raise wrap_oserror(space, e)
 
diff --git a/rpython/rlib/rurandom.py b/rpython/rlib/rurandom.py
--- a/rpython/rlib/rurandom.py
+++ b/rpython/rlib/rurandom.py
@@ -7,12 +7,12 @@
 
 from rpython.rtyper.lltypesystem import lltype, rffi
 from rpython.rlib.objectmodel import not_rpython
+from rpython.translator.tool.cbuild import ExternalCompilationInfo
+from rpython.rtyper.tool import rffi_platform
 
 
 if sys.platform == 'win32':
     from rpython.rlib import rwin32
-    from rpython.translator.tool.cbuild import ExternalCompilationInfo
-    from rpython.rtyper.tool import rffi_platform
 
     eci = ExternalCompilationInfo(
         includes = ['windows.h', 'wincrypt.h'],
@@ -56,7 +56,7 @@
         return lltype.malloc(rffi.CArray(HCRYPTPROV), 1,
                              immortal=True, zero=True)
 
-    def urandom(context, n):
+    def urandom(context, n, signal_checker=None):
         provider = context[0]
         if not provider:
             # This handle is never explicitly released. The operating
@@ -80,11 +80,71 @@
     def init_urandom():
         return None
 
-    def urandom(context, n):
+    SYS_getrandom = None
+
+    if sys.platform.startswith('linux'):
+        eci = ExternalCompilationInfo(includes=['sys/syscall.h'])
+        class CConfig:
+            _compilation_info_ = eci
+            SYS_getrandom = rffi_platform.DefinedConstantInteger(
+                'SYS_getrandom')
+        globals().update(rffi_platform.configure(CConfig))
+
+    if SYS_getrandom is not None:
+        from rpython.rlib.rposix import get_saved_errno, handle_posix_error
+        import errno
+
+        eci = eci.merge(ExternalCompilationInfo(includes=['linux/random.h']))
+        class CConfig:
+            _compilation_info_ = eci
+            GRND_NONBLOCK = rffi_platform.ConstantInteger('GRND_NONBLOCK')
+        globals().update(rffi_platform.configure(CConfig))
+
+        # On Linux, use the syscall() function because the GNU libc doesn't
+        # expose the Linux getrandom() syscall yet.
+        syscall = rffi.llexternal(
+            'syscall',
+            [lltype.Signed, rffi.CCHARP, rffi.LONG, rffi.INT],
+            lltype.Signed,
+            compilation_info=eci,
+            save_err=rffi.RFFI_SAVE_ERRNO)
+
+        class Works:
+            status = True
+        getrandom_works = Works()
+
+        def _getrandom(n, result, signal_checker):
+            if not getrandom_works.status:
+                return n
+            while n > 0:
+                with rffi.scoped_alloc_buffer(n) as buf:
+                    got = syscall(SYS_getrandom, buf.raw, n, GRND_NONBLOCK)
+                    if got >= 0:
+                        s = buf.str(got)
+                        result.append(s)
+                        n -= len(s)
+                        continue
+                err = get_saved_errno()
+                if (err == errno.ENOSYS or err == errno.EPERM or
+                        err == errno.EAGAIN):   # see CPython 3.5
+                    getrandom_works.status = False
+                    return n
+                if err == errno.EINTR:
+                    if signal_checker is not None:
+                        signal_checker()
+                    continue
+                handle_posix_error("getrandom", got)
+                raise AssertionError("unreachable")
+            return n
+
+    def urandom(context, n, signal_checker=None):
         "Read n bytes from /dev/urandom."
-        result = ''
-        if n == 0:
-            return result
+        result = []
+        if SYS_getrandom is not None:
+            n = _getrandom(n, result, signal_checker)
+        if n <= 0:
+            return ''.join(result)
+
         # XXX should somehow cache the file descriptor.  It's a mess.
         # CPython has a 99% solution and hopes for the remaining 1%
         # not to occur.  For now, we just don't cache the file
@@ -98,8 +158,8 @@
                     if e.errno != errno.EINTR:
                         raise
                     data = ''
-                result += data
+                result.append(data)
                 n -= len(data)
         finally:
             os.close(fd)
-        return result
+        return ''.join(result)
diff --git a/rpython/rlib/test/test_rurandom.py 
b/rpython/rlib/test/test_rurandom.py
new file mode 100644
--- /dev/null
+++ b/rpython/rlib/test/test_rurandom.py
@@ -0,0 +1,12 @@
+from rpython.rlib import rurandom
+
+def test_rurandom():
+    context = rurandom.init_urandom()
+    s = rurandom.urandom(context, 5000)
+    assert type(s) is str and len(s) == 5000
+    for x in [1, 11, 111, 222]:
+        assert s.count(chr(x)) >= 1
+
+def test_rurandom_no_syscall(monkeypatch):
+    monkeypatch.setattr(rurandom, 'SYS_getrandom', None)
+    test_rurandom()
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit

Reply via email to