Re-seed a CtrDrbg in a fresh context; keep raw add_source! objects alive

src/ctr_drbg.jl

@@ -1,13 +1,15 @@
 mutable struct CtrDrbg  <: AbstractRNG
     data::Ptr{Cvoid}
-    entropy::Union{Cvoid, Entropy}
+    # The `Entropy` mbedtls reseeds from, or `nothing` until seeding succeeds.
+    entropy::Union{Nothing, Entropy}
     # MbedTLS_jll is built without `MBEDTLS_THREADING_C`, so the C context has no mutex, and
     # every `SSLContext` set up with an `SSLConfig` draws from the config's generator.
     lock::ReentrantLock
 
     function CtrDrbg()
         ctx = new()
         ctx.lock = ReentrantLock()
+        ctx.entropy = nothing
         ctx.data = Libc.malloc(1000)  # 344
         ccall((:mbedtls_ctr_drbg_init, libmbedcrypto), Cvoid, (Ptr{Cvoid},), ctx.data)
         finalizer(ctx->begin
@@ -26,28 +28,37 @@ end
 
 function seed!(rng::CtrDrbg, entropy::Entropy, pdata)
     entropy_func = @cfunction(f_entropy, Cint, (Ref{Entropy}, Ptr{UInt8}, Csize_t))
-    Base.@lock rng.lock @err_check ccall((:mbedtls_ctr_drbg_seed, libmbedcrypto), Cint,
-        (Ptr{Cvoid}, Ptr{Cvoid}, Any, Ptr{Cvoid}, Csize_t),
-        rng.data, entropy_func, entropy, pdata, sizeof(pdata))
-    # mbedtls keeps `entropy` and reseeds from it; this also marks `rng` as seeded.
-    rng.entropy = entropy
+    Base.@lock rng.lock begin
+        # mbedtls seeds a context once: seeding again needs a fresh context, and a failed
+        # seeding leaves it unusable until seeded again.
+        rng.entropy = nothing
+        ccall((:mbedtls_ctr_drbg_free, libmbedcrypto), Cvoid, (Ptr{Cvoid},), rng.data)
+        ccall((:mbedtls_ctr_drbg_init, libmbedcrypto), Cvoid, (Ptr{Cvoid},), rng.data)
+        @err_check ccall((:mbedtls_ctr_drbg_seed, libmbedcrypto), Cint,
+            (Ptr{Cvoid}, Ptr{Cvoid}, Any, Ptr{Cvoid}, Csize_t),
+            rng.data, entropy_func, entropy, pdata, sizeof(pdata))
+        # mbedtls keeps `entropy` and reseeds from it.
+        rng.entropy = entropy
+    end
     rng
 end
 
 seed!(rng::CtrDrbg, entropy) = seed!(rng, entropy, UInt8[])
 
 # Fill `nbytes` bytes at `ptr`, at most `MBEDTLS_CTR_DRBG_MAX_REQUEST` bytes per call.
 function random!(rng::CtrDrbg, ptr::Ptr{UInt8}, nbytes::Integer)
-    # mbedtls crashes when it generates from a context that was never seeded.
-    isdefined(rng, :entropy) ||
-        throw(ArgumentError("a `CtrDrbg` must be seeded with `MbedTLS.seed!` before use"))
-    Base.@lock rng.lock while nbytes > 0
-        n = min(nbytes, MBEDTLS_CTR_DRBG_MAX_REQUEST)
-        @err_check ccall((:mbedtls_ctr_drbg_random, libmbedcrypto), Cint,
-            (Ptr{Cvoid}, Ptr{UInt8}, Csize_t),
-            rng.data, ptr, n)
-        ptr += n
-        nbytes -= n
+    Base.@lock rng.lock begin
+        # mbedtls crashes when it generates from a context that was never seeded.
+        rng.entropy === nothing &&
+            throw(ArgumentError("a `CtrDrbg` must be seeded with `MbedTLS.seed!` before use"))
+        while nbytes > 0
+            n = min(nbytes, MBEDTLS_CTR_DRBG_MAX_REQUEST)
+            @err_check ccall((:mbedtls_ctr_drbg_random, libmbedcrypto), Cint,
+                (Ptr{Cvoid}, Ptr{UInt8}, Csize_t),
+                rng.data, ptr, n)
+            ptr += n
+            nbytes -= n
+        end
     end
 end
 

src/entropy.jl

@@ -20,27 +20,27 @@ mutable struct Entropy
 end
 
 function add_source!(ctx::Entropy, f_source::Ptr, f, threshold, strong)
-    ret = ccall((:mbedtls_entropy_add_source, libmbedcrypto), Cint,
-      (Ptr{Cvoid}, Ptr{Cvoid}, Any, Csize_t, Cint),
-      ctx.data, f_source, f, threshold, strong)
-    Int(ret)
+    Base.@lock ctx.lock begin
+        push!(ctx.sources, f)
+        # Pass the stored object: mbedtls keeps this pointer, and passing `f` directly can hand
+        # it a temporary box (for a closure, say) that nothing keeps alive.
+        ret = ccall((:mbedtls_entropy_add_source, libmbedcrypto), Cint,
+          (Ptr{Cvoid}, Ptr{Cvoid}, Any, Csize_t, Cint),
+          ctx.data, f_source, ctx.sources[end], threshold, strong)
+        Int(ret)
+    end
 end
 
 function jl_entropy(f, output, len, olen)
     output_jl = unsafe_wrap(Array, convert(Ptr{UInt8}, output), len)
-    sz = f(output_jl)
+    # mbedtls reads `olen` bytes from its `len`-byte buffer.
+    sz = clamp(Int(f(output_jl)), 0, Int(len))
     unsafe_store!(convert(Ptr{Csize_t}, olen), Csize_t(sz))
     return Cint(0)
 end
 
-function add_source!(ctx::Entropy, f::Function, threshold, strong)
-    Base.@lock ctx.lock begin
-        push!(ctx.sources, f)
-        # Pass the stored object: mbedtls keeps this pointer, and passing `f` directly can hand
-        # it a temporary box (for a closure, say) that nothing keeps alive.
-        add_source!(ctx, c_entropy[], ctx.sources[end], threshold, strong ? 1 : 0)
-    end
-end
+add_source!(ctx::Entropy, f::Function, threshold, strong) =
+    add_source!(ctx, c_entropy[], f, threshold, strong ? 1 : 0)
 
 const c_entropy = Ref{Ptr{Cvoid}}(C_NULL)
 function __entropyinit__()

test/runtests.jl

@@ -184,6 +184,17 @@ let rng = MbedTLS.CtrDrbg()
     @test 0 <= rand(rng, Float64) < 1
     @test rand(rng, 1:10) in 1:10
     @test all(x -> 0 <= x < 1, rand!(rng, zeros(100)))
+
+    # Seeding again after use works, and the generator later reseeds from the new Entropy.
+    for _ in 1:400
+        rand(rng, UInt8)
+    end
+    MbedTLS.seed!(rng, MbedTLS.Entropy())
+    GC.gc()
+    for _ in 1:10_001  # past the reseed interval
+        rand(rng, UInt8)
+    end
+    @test rand(rng, UInt64) isa UInt64
 end
 
 # Basic TLS client functionality