Add encrypt, decrypt, encrypt!, and decrypt! functions for ChaChaStream

Add methods to use a keystream to encrypt or decrypt a string. I've also
maded some slight improvements and simplifications to other keystream
functions, e.g., making the _refresh_buffer! and _fetch_one! functions
no longer be generated.

Author: kernelmethod

src/ChaChaCiphers.jl

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

src/generation.jl

@@ -9,6 +9,9 @@ using Random: SamplerType
 
 Random.rng_native_52(::ChaChaStream) = UInt64
 
+Random.rand(rng::ChaChaStream, ::Type{T}) where {T <: BitInteger} =
+    _fetch_one!(rng, T)
+
 Random.rand(rng::ChaChaStream, T::Random.SamplerType{<:BitInteger}) =
     _fetch_one!(rng, T[])
 

src/keystream.jl

@@ -45,34 +45,30 @@ ChaCha12Stream(args...) = ChaChaStream(args...; doublerounds=6)
 Base.show(io::IO, rng::ChaChaStream) =
     write(io, "ChaChaStream(key=$(rng.key), nonce=$(rng.nonce), counter=$(rng.counter))")
 
-buffer_size(stream::ChaChaStream) = STREAM_BUFFER_SIZE - stream.position
+buffer_size(stream::ChaChaStream) = length(stream.buffer) - stream.position + 1
+
+# Accessors; must be defined for AbstractChaChaStream.
 
 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))
-
-    if rem != 0
-        error("STREAM_BUFFER_SIZE must be a multiple of CHACHA_BLOCK_SIZE")
-    end
-
-    quote
-        _fill_blocks!(stream.buffer, stream, $(blocks_in_buffer))
-        stream.position = 1
-        stream
-    end
+function _refresh_buffer!(stream::ChaChaStream)
+    _fill_blocks!(
+        stream.buffer,
+        stream,
+        STREAM_BUFFER_BLOCKS
+    )
+    stream.position = 1
+    stream
 end
 
 # Fill a buffer with a block of the keystream
 function _fill_blocks!(
     buffer::AbstractVector{T}, stream::ChaChaStream, nblocks::Int
 ) where {T <: BitInteger}
-    # Buffer length when viewed as an array of u32
     bufsize_u32 = div(length(buffer) * sizeof(T), sizeof(UInt32))
 
     GC.@preserve buffer begin
@@ -98,25 +94,21 @@ function _fill_blocks!(
 end
 
 # Fetch bytes from the ChaChaStream buffer and increment the position index
-@generated function _fetch_one!(stream::ChaChaStream, ::Type{T}) where {T <: BitInteger}
-    local type_size = sizeof(T)
-
-    quote
-        if buffer_size(stream) < $(type_size)
-            _refresh_buffer!(stream)
-        end
+function _fetch_one!(stream::ChaChaStream, ::Type{T}) where {T <: BitInteger}
+    if buffer_size(stream) < sizeof(T)
+        _refresh_buffer!(stream)
+    end
 
-        buffer = stream.buffer
+    buffer = stream.buffer
 
-        val = GC.@preserve buffer begin
-            p = pointer(buffer, stream.position)
-            p = Base.unsafe_convert(Ptr{T}, p)
-            unsafe_load(p)
-        end
-
-        stream.position += $(type_size)
-        val
+    val = GC.@preserve buffer begin
+        p = pointer(buffer, stream.position)
+        p = Base.unsafe_convert(Ptr{T}, p)
+        unsafe_load(p)
     end
+
+    stream.position += sizeof(T)
+    val
 end
 
 function _fill_buffer!(dest::AbstractVector{UInt8}, stream::ChaChaStream)
@@ -138,12 +130,9 @@ function _fill_buffer!(dest::AbstractVector{UInt8}, stream::ChaChaStream)
     # Instead of repeatedly operating on the stream buffer, we insert
     # as many blocks of the keystream into the destination as possible
     (n_blocks, rem) = divrem(length(dest) - bfsize, CHACHA_BLOCK_SIZE)
-    n_blocks = Int(n_blocks)
-    #slice = view(dest, bfsize + 1:bfsize + n_blocks * CHACHA_BLOCK_SIZE)
     GC.@preserve dest begin
         sp = pointer(dest, bfsize + 1)
         slice = unsafe_wrap(Vector{UInt8}, sp, n_blocks * CHACHA_BLOCK_SIZE, own=false)
-        #state = unsafe_wrap(Vector{UInt32}, sp, bufsize_u32, own=false)
         _fill_blocks!(slice, stream, n_blocks)
     end
 
@@ -154,3 +143,28 @@ function _fill_buffer!(dest::AbstractVector{UInt8}, stream::ChaChaStream)
 
     dest
 end
+
+#=
+Stream cipher methods
+=#
+
+"""
+    decrypt(stream::ChaChaStream, x)
+    decrypt!(stream::ChaChaStream, x)
+
+Decrypt an encrypted vector `x` using the keystream from a [`ChaChaStream`](@ref).
+"""
+decrypt!(stream::ChaChaStream, x) = encrypt!(stream, x)
+decrypt(stream::ChaChaStream, x) = decrypt!(stream, copy(x))
+decrypt(stream::ChaChaStream, x::AbstractString) = decrypt!(stream, collect(codeunits(x)))
+
+"""
+    encrypt(stream::ChaChaStream, x)
+    encrypt!(stream::ChaChaStream, x)
+
+Encrypt a vector or string `x` using the keystream from a [`ChaChaStream`](@ref).
+"""
+encrypt!(stream::ChaChaStream, x::DenseVector{T}) where {T <: BitInteger} =
+    map!(u -> u ⊻ _fetch_one!(stream, T), x, x)
+encrypt(stream::ChaChaStream, x) = encrypt!(stream, copy(x))
+encrypt(stream::ChaChaStream, x::AbstractString) = encrypt!(stream, collect(codeunits(x)))

test/test_keystream.jl

@@ -67,5 +67,114 @@ using Test
         rand_repro = rand(stream_repro, 1:10, 1024)
         @test rand_orig == rand_repro
     end
+
+    @testset "Encrypt data with a keystream" begin
+        # Ref: IETF RFC 8439, Sec. A.2
+        # https://datatracker.ietf.org/doc/html/rfc8439#appendix-A.2
+
+        # Test Vector #1:
+        # ==============
+        stream = ChaCha20Stream(zeros(UInt32, 8), 0)
+        state = getstate(stream)
+        text = zeros(UInt8, 64)
+
+        ciphertext = encrypt(stream, text)
+        plaintext = decrypt(ChaChaStream(state), ciphertext)
+
+        test_vector = Vector{UInt8}([
+            0x76, 0xb8, 0xe0, 0xad, 0xa0, 0xf1, 0x3d, 0x90, 0x40, 0x5d, 0x6a, 0xe5, 0x53, 0x86, 0xbd, 0x28,
+            0xbd, 0xd2, 0x19, 0xb8, 0xa0, 0x8d, 0xed, 0x1a, 0xa8, 0x36, 0xef, 0xcc, 0x8b, 0x77, 0x0d, 0xc7,
+            0xda, 0x41, 0x59, 0x7c, 0x51, 0x57, 0x48, 0x8d, 0x77, 0x24, 0xe0, 0x3f, 0xb8, 0xd8, 0x4a, 0x37,
+            0x6a, 0x43, 0xb8, 0xf4, 0x15, 0x18, 0xa1, 0x1c, 0xc3, 0x87, 0xb6, 0x69, 0xb2, 0xee, 0x65, 0x86
+        ])
+        @test ciphertext == test_vector
+        @test plaintext == text
+
+
+        # Test Vector #2:
+        # ==============
+        key = Vector{UInt8}([
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
+        ])
+        key = reinterpret(UInt32, key)
+        counter = UInt64(1)
+        nonce = UInt64(0x0200000000000000)
+        stream = ChaCha20Stream(key, nonce, counter)
+        state = getstate(stream)
+        text =
+            "Any submission to the IETF intended by the Contributor for " *
+            "publication as all or part of an IETF Internet-Draft or RFC " *
+            "and any statement made within the context of an IETF activity " *
+            "is considered an \"IETF Contribution\". Such statements include " *
+            "oral statements in IETF sessions, as well as written and " *
+            "electronic communications made at any time or place, which are " *
+            "addressed to"
+
+        ciphertext = encrypt(stream, text)
+        plaintext = decrypt(ChaChaStream(state), ciphertext)
+        test_vector = Vector{UInt8}([
+            0xa3, 0xfb, 0xf0, 0x7d, 0xf3, 0xfa, 0x2f, 0xde, 0x4f, 0x37, 0x6c, 0xa2, 0x3e, 0x82, 0x73, 0x70,
+            0x41, 0x60, 0x5d, 0x9f, 0x4f, 0x4f, 0x57, 0xbd, 0x8c, 0xff, 0x2c, 0x1d, 0x4b, 0x79, 0x55, 0xec,
+            0x2a, 0x97, 0x94, 0x8b, 0xd3, 0x72, 0x29, 0x15, 0xc8, 0xf3, 0xd3, 0x37, 0xf7, 0xd3, 0x70, 0x05,
+            0x0e, 0x9e, 0x96, 0xd6, 0x47, 0xb7, 0xc3, 0x9f, 0x56, 0xe0, 0x31, 0xca, 0x5e, 0xb6, 0x25, 0x0d,
+            0x40, 0x42, 0xe0, 0x27, 0x85, 0xec, 0xec, 0xfa, 0x4b, 0x4b, 0xb5, 0xe8, 0xea, 0xd0, 0x44, 0x0e,
+            0x20, 0xb6, 0xe8, 0xdb, 0x09, 0xd8, 0x81, 0xa7, 0xc6, 0x13, 0x2f, 0x42, 0x0e, 0x52, 0x79, 0x50,
+            0x42, 0xbd, 0xfa, 0x77, 0x73, 0xd8, 0xa9, 0x05, 0x14, 0x47, 0xb3, 0x29, 0x1c, 0xe1, 0x41, 0x1c,
+            0x68, 0x04, 0x65, 0x55, 0x2a, 0xa6, 0xc4, 0x05, 0xb7, 0x76, 0x4d, 0x5e, 0x87, 0xbe, 0xa8, 0x5a,
+            0xd0, 0x0f, 0x84, 0x49, 0xed, 0x8f, 0x72, 0xd0, 0xd6, 0x62, 0xab, 0x05, 0x26, 0x91, 0xca, 0x66,
+            0x42, 0x4b, 0xc8, 0x6d, 0x2d, 0xf8, 0x0e, 0xa4, 0x1f, 0x43, 0xab, 0xf9, 0x37, 0xd3, 0x25, 0x9d,
+            0xc4, 0xb2, 0xd0, 0xdf, 0xb4, 0x8a, 0x6c, 0x91, 0x39, 0xdd, 0xd7, 0xf7, 0x69, 0x66, 0xe9, 0x28,
+            0xe6, 0x35, 0x55, 0x3b, 0xa7, 0x6c, 0x5c, 0x87, 0x9d, 0x7b, 0x35, 0xd4, 0x9e, 0xb2, 0xe6, 0x2b,
+            0x08, 0x71, 0xcd, 0xac, 0x63, 0x89, 0x39, 0xe2, 0x5e, 0x8a, 0x1e, 0x0e, 0xf9, 0xd5, 0x28, 0x0f,
+            0xa8, 0xca, 0x32, 0x8b, 0x35, 0x1c, 0x3c, 0x76, 0x59, 0x89, 0xcb, 0xcf, 0x3d, 0xaa, 0x8b, 0x6c,
+            0xcc, 0x3a, 0xaf, 0x9f, 0x39, 0x79, 0xc9, 0x2b, 0x37, 0x20, 0xfc, 0x88, 0xdc, 0x95, 0xed, 0x84,
+            0xa1, 0xbe, 0x05, 0x9c, 0x64, 0x99, 0xb9, 0xfd, 0xa2, 0x36, 0xe7, 0xe8, 0x18, 0xb0, 0x4b, 0x0b,
+            0xc3, 0x9c, 0x1e, 0x87, 0x6b, 0x19, 0x3b, 0xfe, 0x55, 0x69, 0x75, 0x3f, 0x88, 0x12, 0x8c, 0xc0,
+            0x8a, 0xaa, 0x9b, 0x63, 0xd1, 0xa1, 0x6f, 0x80, 0xef, 0x25, 0x54, 0xd7, 0x18, 0x9c, 0x41, 0x1f,
+            0x58, 0x69, 0xca, 0x52, 0xc5, 0xb8, 0x3f, 0xa3, 0x6f, 0xf2, 0x16, 0xb9, 0xc1, 0xd3, 0x00, 0x62,
+            0xbe, 0xbc, 0xfd, 0x2d, 0xc5, 0xbc, 0xe0, 0x91, 0x19, 0x34, 0xfd, 0xa7, 0x9a, 0x86, 0xf6, 0xe6,
+            0x98, 0xce, 0xd7, 0x59, 0xc3, 0xff, 0x9b, 0x64, 0x77, 0x33, 0x8f, 0x3d, 0xa4, 0xf9, 0xcd, 0x85,
+            0x14, 0xea, 0x99, 0x82, 0xcc, 0xaf, 0xb3, 0x41, 0xb2, 0x38, 0x4d, 0xd9, 0x02, 0xf3, 0xd1, 0xab,
+            0x7a, 0xc6, 0x1d, 0xd2, 0x9c, 0x6f, 0x21, 0xba, 0x5b, 0x86, 0x2f, 0x37, 0x30, 0xe3, 0x7c, 0xfd,
+            0xc4, 0xfd, 0x80, 0x6c, 0x22, 0xf2, 0x21
+        ])
+        @test ciphertext == test_vector
+        @test String(plaintext) == text
+
+        # Test Vector #3:
+        # ==============
+        key = Vector{UInt8}([
+            0x1c, 0x92, 0x40, 0xa5, 0xeb, 0x55, 0xd3, 0x8a, 0xf3, 0x33, 0x88, 0x86, 0x04, 0xf6, 0xb5, 0xf0,
+            0x47, 0x39, 0x17, 0xc1, 0x40, 0x2b, 0x80, 0x09, 0x9d, 0xca, 0x5c, 0xbc, 0x20, 0x70, 0x75, 0xc0,
+        ])
+        key = reinterpret(UInt32, key)
+        nonce = UInt64(0x0200000000000000)
+        counter = UInt64(42)
+        stream = ChaCha20Stream(key, nonce, counter)
+        state = getstate(stream)
+
+        text = """
+            'Twas brillig, and the slithy toves
+            Did gyre and gimble in the wabe:
+            All mimsy were the borogoves,
+            And the mome raths outgrabe."""
+
+        ciphertext = encrypt(stream, text)
+        plaintext = decrypt(ChaChaStream(state), ciphertext)
+
+        test_vector = Vector{UInt8}([
+            0x62, 0xe6, 0x34, 0x7f, 0x95, 0xed, 0x87, 0xa4, 0x5f, 0xfa, 0xe7, 0x42, 0x6f, 0x27, 0xa1, 0xdf,
+            0x5f, 0xb6, 0x91, 0x10, 0x04, 0x4c, 0x0d, 0x73, 0x11, 0x8e, 0xff, 0xa9, 0x5b, 0x01, 0xe5, 0xcf,
+            0x16, 0x6d, 0x3d, 0xf2, 0xd7, 0x21, 0xca, 0xf9, 0xb2, 0x1e, 0x5f, 0xb1, 0x4c, 0x61, 0x68, 0x71,
+            0xfd, 0x84, 0xc5, 0x4f, 0x9d, 0x65, 0xb2, 0x83, 0x19, 0x6c, 0x7f, 0xe4, 0xf6, 0x05, 0x53, 0xeb,
+            0xf3, 0x9c, 0x64, 0x02, 0xc4, 0x22, 0x34, 0xe3, 0x2a, 0x35, 0x6b, 0x3e, 0x76, 0x43, 0x12, 0xa6,
+            0x1a, 0x55, 0x32, 0x05, 0x57, 0x16, 0xea, 0xd6, 0x96, 0x25, 0x68, 0xf8, 0x7d, 0x3f, 0x3f, 0x77,
+            0x04, 0xc6, 0xa8, 0xd1, 0xbc, 0xd1, 0xbf, 0x4d, 0x50, 0xd6, 0x15, 0x4b, 0x6d, 0xa7, 0x31, 0xb1,
+            0x87, 0xb5, 0x8d, 0xfd, 0x72, 0x8a, 0xfa, 0x36, 0x75, 0x7a, 0x79, 0x7a, 0xc1, 0x88, 0xd1,
+        ])
+        @test ciphertext == test_vector
+        @test String(plaintext) == text
+    end
 end