From 99252ae2efcd9c8db9da74589a2fed67b8387c64 Mon Sep 17 00:00:00 2001
From: Yuta Saito <kateinoigakukun@gmail.com>
Date: Fri, 26 Apr 2024 11:12:48 +0900
Subject: [PATCH] Add RandomBufferGenerator protocol for WASI random_get

This protocol allows to inject a custom random number generator for
`wasi_snapshot_preview1::random_get` function. The default implementation
continues to use `swift_stdlib_random` function but users can provide
their own implementation. Additionally, types conforming to
`RandomNumberGenerator` can automatically conform to the new protocol.

This is useful when we want fully deterministic behavior like build
tools or tests.
---
 Sources/WASI/RandomBufferGenerator.swift      | 48 +++++++++++++++++++
 Sources/WASI/WASI.swift                       |  8 ++--
 .../RandomBufferGeneratorTests.swift          | 35 ++++++++++++++
 3 files changed, 88 insertions(+), 3 deletions(-)
 create mode 100644 Sources/WASI/RandomBufferGenerator.swift
 create mode 100644 Tests/WASITests/RandomBufferGeneratorTests.swift

diff --git a/Sources/WASI/RandomBufferGenerator.swift b/Sources/WASI/RandomBufferGenerator.swift
new file mode 100644
index 00000000..a0a7d420
--- /dev/null
+++ b/Sources/WASI/RandomBufferGenerator.swift
@@ -0,0 +1,48 @@
+import SwiftShims  // For swift_stdlib_random
+
+/// A type that provides random bytes.
+///
+/// This type is similar to `RandomNumberGenerator` in Swift standard library,
+/// but it provides a way to fill a buffer with random bytes instead of a single
+/// random number.
+public protocol RandomBufferGenerator {
+
+    /// Fills the buffer with random bytes.
+    ///
+    /// - Parameter buffer: The destination buffer to fill with random bytes.
+    mutating func fill(buffer: UnsafeMutableBufferPointer<UInt8>)
+}
+
+extension RandomBufferGenerator where Self: RandomNumberGenerator {
+    public mutating func fill(buffer: UnsafeMutableBufferPointer<UInt8>) {
+        // The buffer is filled with 8 bytes at once.
+        let count = buffer.count / 8
+        for i in 0..<count {
+            let random = self.next()
+            withUnsafeBytes(of: random) { randomBytes in
+                let startOffset = i * 8
+                let destination = UnsafeMutableBufferPointer(rebasing: buffer[startOffset..<(startOffset + 8)])
+                randomBytes.copyBytes(to: destination)
+            }
+        }
+
+        // If the buffer size is not a multiple of 8, fill the remaining bytes.
+        let remaining = buffer.count % 8
+        if remaining > 0 {
+            let random = self.next()
+            withUnsafeBytes(of: random) { randomBytes in
+                let startOffset = count * 8
+                let destination = UnsafeMutableBufferPointer(rebasing: buffer[startOffset..<(startOffset + remaining)])
+                randomBytes.copyBytes(to: destination)
+            }
+        }
+    }
+}
+
+extension SystemRandomNumberGenerator: RandomBufferGenerator {
+    public mutating func fill(buffer: UnsafeMutableBufferPointer<UInt8>) {
+        guard let baseAddress = buffer.baseAddress else { return }
+        // Directly call underlying C function of SystemRandomNumberGenerator
+        swift_stdlib_random(baseAddress, Int(buffer.count))
+    }
+}
diff --git a/Sources/WASI/WASI.swift b/Sources/WASI/WASI.swift
index 7a0ecd26..6ce801cd 100644
--- a/Sources/WASI/WASI.swift
+++ b/Sources/WASI/WASI.swift
@@ -1,5 +1,4 @@
 import Foundation
-import SwiftShims  // For swift_stdlib_random
 import SystemExtras
 import SystemPackage
 import WasmTypes
@@ -1356,6 +1355,7 @@ public class WASIBridgeToHost: WASI {
     private let args: [String]
     private let environment: [String: String]
     private var fdTable: FdTable
+    private var randomGenerator: RandomBufferGenerator
 
     public init(
         args: [String] = [],
@@ -1363,7 +1363,8 @@ public class WASIBridgeToHost: WASI {
         preopens: [String: String] = [:],
         stdin: FileDescriptor = .standardInput,
         stdout: FileDescriptor = .standardOutput,
-        stderr: FileDescriptor = .standardError
+        stderr: FileDescriptor = .standardError,
+        randomGenerator: RandomBufferGenerator = SystemRandomNumberGenerator()
     ) throws {
         self.args = args
         self.environment = environment
@@ -1386,6 +1387,7 @@ public class WASIBridgeToHost: WASI {
             }
         }
         self.fdTable = fdTable
+        self.randomGenerator = randomGenerator
     }
 
     public var wasiHostModules: [String: WASIHostModule] { _hostModules }
@@ -1860,7 +1862,7 @@ public class WASIBridgeToHost: WASI {
     func random_get(buffer: UnsafeGuestPointer<UInt8>, length: WASIAbi.Size) {
         guard length > 0 else { return }
         buffer.withHostPointer(count: Int(length)) {
-            swift_stdlib_random($0.baseAddress!, Int(length))
+            self.randomGenerator.fill(buffer: $0)
         }
     }
 }
diff --git a/Tests/WASITests/RandomBufferGeneratorTests.swift b/Tests/WASITests/RandomBufferGeneratorTests.swift
new file mode 100644
index 00000000..85a9751a
--- /dev/null
+++ b/Tests/WASITests/RandomBufferGeneratorTests.swift
@@ -0,0 +1,35 @@
+import XCTest
+
+@testable import WASI
+
+final class RandomBufferGeneratorTests: XCTestCase {
+    struct DeterministicGenerator: RandomNumberGenerator, RandomBufferGenerator {
+        var items: [UInt64]
+
+        mutating func next() -> UInt64 {
+            items.removeFirst()
+        }
+    }
+    func testDefaultFill() {
+        var generator = DeterministicGenerator(items: [
+            0x0123456789abcdef, 0xfedcba9876543210, 0xdeadbeefbaddcafe
+        ])
+        for (bufferSize, expectedBytes): (Int, [UInt8]) in [
+            (10, [0xef, 0xcd, 0xab, 0x89, 0x67, 0x45, 0x23, 0x01, 0x10, 0x32]),
+            (2, [0xfe, 0xca]),
+            (0, [])
+        ] {
+            var buffer: [UInt8] = Array(repeating: 0, count: bufferSize)
+            buffer.withUnsafeMutableBufferPointer {
+                generator.fill(buffer: $0)
+            }
+            let expected: [UInt8]
+#if _endian(little)
+            expected = expectedBytes
+#else
+            expected = Array(expectedBytes.reversed())
+#endif
+            XCTAssertEqual(buffer, expected)
+        }
+    }
+}