Make it possible to save and restore the state of a ChaChaStream

Add a getstate() function that saves the state of a ChaChaStream as a
ChaChaState NamedTuple, and add constructors to restore the ChaChaStream
from that saved state.

Author: kernelmethod

src/ChaChaCiphers.jl

@@ -8,5 +8,6 @@ include("generation.jl")
 
 export ChaCha
 export ChaChaStream, ChaCha20Stream, ChaCha12Stream
+export getstate
 
 end

src/core.jl

@@ -1,5 +1,6 @@
 using ChaChaCiphers.ChaCha: CHACHA_BLOCK_SIZE
 using Random: AbstractRNG
+using StaticArrays
 
 ### UnsafeView
 
@@ -26,11 +27,34 @@ function Base.getindex(a::UnsafeView{T}, i::UnitRange) where T
     UnsafeView(pointer(a) + view_start, len)
 end
 
+### ChaChaState
+
+"""
+    ChaChaState
+
+A `NamedTuple` storing the current state of a ChaCha keystream. `ChaChaState`
+can be used to save and restore the state of a keystream.
+"""
+const ChaChaState = @NamedTuple begin
+    key :: SVector{8,UInt32}
+    nonce :: UInt64
+    counter :: UInt64
+    position :: Int
+    doublerounds :: Int
+end
+
+
 ### AbstractChaChaStream
 
 # Number of bytes to store in an AbstractChaChaStream
 const STREAM_BUFFER_SIZE = CHACHA_BLOCK_SIZE
 
+if STREAM_BUFFER_SIZE % CHACHA_BLOCK_SIZE != 0
+    error("STREAM_BUFFER_SIZE must be a multiple of the CHACHA_BLOCK_SIZE")
+end
+
+const STREAM_BUFFER_BLOCKS = STREAM_BUFFER_SIZE ÷ CHACHA_BLOCK_SIZE
+
 """
     AbstractChaChaStream
 
@@ -44,4 +68,21 @@ abstract type AbstractChaChaStream <: AbstractRNG end
     T(SVector{8,UInt32}(rand(RandomDevice(), UInt32, 8)))
 (::Type{T})(key; kws...) where {T <: AbstractChaChaStream} =
     T(key, UInt64(0); kws...)
+(::Type{T})(state::ChaChaState) where {T <: AbstractChaChaStream} =
+    T(state.key, state.nonce, state.counter, state.position; doublerounds=state.doublerounds)
 
+"""
+    getstate(stream::AbstractChaChaStream)
+
+Return a `NamedTuple` containing enough state of the input
+`AbstractChaChaStream` to be able to reproduce it.
+"""
+function getstate(stream::AbstractChaChaStream)
+    (;
+        :key => key(stream),
+        :nonce => nonce(stream),
+        :counter => counter(stream) - STREAM_BUFFER_BLOCKS,
+        :position => position(stream),
+        :doublerounds => doublerounds(stream),
+    )
+end

src/keystream.jl

@@ -17,20 +17,28 @@ mutable struct ChaChaStream <: AbstractChaChaStream
     position :: Int
     doublerounds :: Int
 
-    function ChaChaStream(key, nonce, counter = UInt64(0); doublerounds = 10)
-        if doublerounds ∉ (6, 10)
-            error("`doublerounds` must be equal to 6 or 10")
+    function ChaChaStream(
+        key,
+        nonce,
+        counter = UInt64(0),
+        position = 1;
+        doublerounds = 10
+    )
+        if doublerounds < 0 || !iseven(doublerounds)
+            error("`doublerounds` must be an even positive number")
         end
 
         key = SVector{8,UInt32}(key)
         buffer = MVector{STREAM_BUFFER_SIZE,UInt8}(undef)
         stream = new(key, nonce, counter, buffer, 1, doublerounds)
         _refresh_buffer!(stream)
+        stream.position = position
+
+        stream
     end
 end
 
 # Constructors
-
 ChaCha20Stream(args...) = ChaChaStream(args...; doublerounds=10)
 ChaCha12Stream(args...) = ChaChaStream(args...; doublerounds=6)
 
@@ -39,6 +47,12 @@ Base.show(io::IO, rng::ChaChaStream) =
 
 buffer_size(stream::ChaChaStream) = STREAM_BUFFER_SIZE - stream.position
 
+key(stream::ChaChaStream) = stream.key
+nonce(stream::ChaChaStream) = stream.nonce
+counter(stream::ChaChaStream) = stream.counter
+position(stream::ChaChaStream) = stream.position
+doublerounds(stream::ChaChaStream) = stream.doublerounds
+
 @generated function _refresh_buffer!(stream::ChaChaStream)
     local blocks_in_buffer, rem = divrem(STREAM_BUFFER_SIZE, CHACHA_BLOCK_SIZE)
     local words_per_block = div(CHACHA_BLOCK_SIZE, sizeof(UInt32))

test/test_keystream.jl

@@ -46,5 +46,26 @@ using Test
         @test isapprox(mean(samples), 0., atol=1e-2)
         @test isapprox(std(samples), 1., atol=1e-2)
     end
+
+    @testset "Save and restore a keystream" begin
+        # Create a stream, run some operations on it,
+        # and save its state. Ensure that we can
+        # reproduce the stream from its saved state.
+        stream = ChaCha12Stream()
+        rand(stream, UInt8, 3_000)
+        stream_repro = getstate(stream) |> ChaChaStream
+
+        rand_orig = randn(stream, 5_000)
+        rand_repro = randn(stream_repro, 5_000)
+        @test rand_orig == rand_repro
+
+        stream = ChaCha20Stream()
+        randn(stream, Float64, 3_000)
+        stream_repro = getstate(stream) |> ChaChaStream
+
+        rand_orig = rand(stream, 1:10, 1024)
+        rand_repro = rand(stream_repro, 1:10, 1024)
+        @test rand_orig == rand_repro
+    end
 end