From fca40d776465528f351b8f2d1dba7012b6fd53db Mon Sep 17 00:00:00 2001
From: Didactic Drunk <1479616+didactic-drunk@users.noreply.github.com>
Date: Mon, 14 Jun 2021 17:02:37 -0700
Subject: [PATCH] Sodium::SecureBuffer is now a Crypto::Secret

---
 shard.yml                         |   4 +
 spec/sodium/secure_buffer_spec.cr |  43 +++++----
 src/sodium/kdf.cr                 |   8 +-
 src/sodium/secure_buffer.cr       | 152 ++++--------------------------
 4 files changed, 56 insertions(+), 151 deletions(-)

diff --git a/shard.yml b/shard.yml
index fe51ca9..a77c077 100644
--- a/shard.yml
+++ b/shard.yml
@@ -10,6 +10,10 @@ targets:
     main: examples/blake2b_hash.cr
   pwhash_selector:
     main: examples/pwhash_selector.cr
+dependencies:
+  crypto-secret:
+    github: didactic-drunk/crypto-secret.cr
+    branch: main
 libraries:
   libsodium: ">= 1.0.18"
 license: MIT
diff --git a/spec/sodium/secure_buffer_spec.cr b/spec/sodium/secure_buffer_spec.cr
index 368581f..58dfe8f 100644
--- a/spec/sodium/secure_buffer_spec.cr
+++ b/spec/sodium/secure_buffer_spec.cr
@@ -7,8 +7,10 @@ end
 describe Sodium::SecureBuffer do
   it "allocates empty" do
     buf = Sodium::SecureBuffer.new 5
-    buf.to_slice.each do |b|
-      b.should eq 0xdb_u8
+    buf.readonly do |slice|
+      slice.each do |b|
+        b.should eq 0xdb_u8
+      end
     end
 
     buf.noaccess
@@ -17,18 +19,22 @@ describe Sodium::SecureBuffer do
   end
 
   it "allocates random" do
-    buf = Sodium::SecureBuffer.random 5
-    buf.to_slice.bytesize.should eq 5
-    buf.wipe
+    buf1 = Sodium::SecureBuffer.random 5
+    buf2 = Sodium::SecureBuffer.random 5
+    (buf1 == buf2).should be_false
+    buf1.wipe
   end
 
   it "copies and erases" do
     bytes = Bytes.new(5) { 1_u8 }
 
     buf = Sodium::SecureBuffer.new bytes, erase: true
-    buf.to_slice.bytesize.should eq 5
-    buf.to_slice.each do |b|
-      b.should eq 1_u8
+    buf.readonly do |slice|
+      slice.bytesize.should eq 5
+
+      slice.each do |b|
+        b.should eq 1_u8
+      end
     end
 
     bytes.to_slice.each do |b|
@@ -37,16 +43,21 @@ describe Sodium::SecureBuffer do
   end
 
   it "dups without crashing" do
-    buf = Sodium::SecureBuffer.new 5
-    buf.readwrite
+    buf1 = Sodium::SecureBuffer.new 5
+    buf1.noaccess
 
-    buf2 = buf.dup
-    buf2.@state.should eq Sodium::SecureBuffer::State::Readwrite
+    buf2 = buf1.dup
+    buf2.@state.should eq Sodium::SecureBuffer::State::Noaccess
 
-    buf[0] = 1_u8
-    buf.to_slice.hexstring.should_not eq buf2.to_slice.hexstring
-    buf2[0] = 1_u8
-    buf.to_slice.hexstring.should eq buf2.to_slice.hexstring
+    buf1.readwrite do |slice|
+      slice[0] = 1_u8
+    end
+    buf1.hexstring.should_not eq buf2.hexstring
+
+    buf2.readwrite do |slice|
+      slice[0] = 1_u8
+    end
+    buf1.hexstring.should eq buf2.hexstring
   end
 
   it "transitions correctly" do
diff --git a/src/sodium/kdf.cr b/src/sodium/kdf.cr
index 1ed52ec..71fb459 100644
--- a/src/sodium/kdf.cr
+++ b/src/sodium/kdf.cr
@@ -79,9 +79,11 @@ module Sodium
       end
 
       subkey = SecureBuffer.new subkey_size
-      @sbuf.readonly do
-        if (ret = LibSodium.crypto_kdf_derive_from_key(subkey, subkey.bytesize, subkey_id, context, self.to_slice)) != 0
-          raise Sodium::Error.new("crypto_kdf_derive_from_key returned #{ret} (subkey size is probably out of range)")
+      subkey.readwrite do |sub_slice|
+        @sbuf.readonly do
+          if (ret = LibSodium.crypto_kdf_derive_from_key(sub_slice, sub_slice.bytesize, subkey_id, context, self.to_slice)) != 0
+            raise Sodium::Error.new("crypto_kdf_derive_from_key returned #{ret} (subkey size is probably out of range)")
+          end
         end
       end
 
diff --git a/src/sodium/secure_buffer.cr b/src/sodium/secure_buffer.cr
index 8b02384..883ceb6 100644
--- a/src/sodium/secure_buffer.cr
+++ b/src/sodium/secure_buffer.cr
@@ -1,5 +1,6 @@
 require "./lib_sodium"
 require "./wipe"
+require "crypto-secret/stateful"
 
 module Sodium
   # Allocate guarded memory using [sodium_malloc](https://libsodium.gitbook.io/doc/memory_management)
@@ -7,44 +8,17 @@ module Sodium
   # #initialize returns readonly or readwrite for thread safety
   # When state changes are required (such as using #noaccess) and the buffer is accessed from multiple threads wrap each #readonly/#readwrite block in a lock.
   class SecureBuffer
-    class Error < Sodium::Error
-      class KeyWiped < Error
-      end
+    include Crypto::Secret::Stateful
 
-      class InvalidStateTransition < Error
-      end
+    #    @state = State::Readwrite
 
-      # Check RLIMIT_MEMLOCK if you receive this
-      class OutOfMemory < Error
-      end
-    end
-
-    enum State
-      Cloning
-      Wiped
-      Noaccess
-      Readonly
-      Readwrite
-    end
-
-    @state = State::Readwrite
-
-    getter bytesize
-
-    delegate :+, :[], :[]=, :hexstring, to: to_slice
+    getter bytesize : Int32
 
     def initialize(@bytesize : Int32)
       @ptr = LibSodium.sodium_malloc @bytesize
       raise Error::OutOfMemory.new if @ptr.null?
     end
 
-    # Returns a **readonly** random SecureBuffer.
-    def self.random(size)
-      buf = new(size)
-      Random::Secure.random_bytes buf.to_slice
-      buf.readonly
-    end
-
     # Copies bytes to a **readonly** SecureBuffer.
     # Optionally erases bytes after copying if erase is set
     # Returns a **readonly** SecureBuffer.
@@ -61,30 +35,16 @@ module Sodium
       initialize sbuf.bytesize
 
       # Maybe not thread safe
-      sbuf.readonly do
-        sbuf.to_slice.copy_to self.to_slice
+      sbuf.readonly do |s1|
+        self.to_slice do |s2|
+          s1.copy_to s2.to_slice
+        end
       end
 
       @state = State::Cloning
       set_state sbuf.@state
     end
 
-    # WARNING: Not thread safe
-    def wipe
-      return if @state == State::Wiped
-      readwrite
-      Sodium.memzero self.to_slice
-      @state = State::Wiped
-      noaccess!
-    end
-
-    # WARNING: Not thread safe
-    def wipe
-      yield
-    ensure
-      wipe
-    end
-
     # :nodoc:
     def finalize
       LibSodium.sodium_free @ptr
@@ -103,6 +63,11 @@ module Sodium
       Slice(UInt8).new @ptr, @bytesize
     end
 
+    def to_slice(& : Bytes -> Nil)
+      yield Bytes.new @ptr, @bytesize
+    end
+
+    # :nodoc:
     def to_unsafe
       @ptr
     end
@@ -112,98 +77,21 @@ module Sodium
       self.class.new self
     end
 
-    # Temporarily make buffer readonly within the block returning to the prior state on exit.
-    # WARNING: Not thread safe unless this object is readonly or readwrite
-    def readonly
-      with_state State::Readonly do
-        yield
-      end
-    end
-
-    # Temporarily make buffer readwrite within the block returning to the prior state on exit.
-    # WARNING: Not thread safe unless this object is **readwrite**
-    def readwrite
-      with_state State::Readwrite do
-        yield
-      end
-    end
-
-    # Makes a region allocated using sodium_malloc() or sodium_allocarray() inaccessible. It cannot be read or written, but the data are preserved.
-    # WARNING: Not thread safe
-    def noaccess
-      raise Error::KeyWiped.new if @state == State::Wiped
-      noaccess!
-      @state = State::Noaccess
-      self
-    end
-
-    # Also used by #wipe
-    private def noaccess!
-      if LibSodium.sodium_mprotect_noaccess(@ptr) != 0
-        raise "sodium_mprotect_noaccess"
-      end
-      self
-    end
-
-    # Marks a region allocated using sodium_malloc() or sodium_allocarray() as read-only.
-    # WARNING: Not thread safe
-    def readonly
-      raise Error::KeyWiped.new if @state == State::Wiped
-      if LibSodium.sodium_mprotect_readonly(@ptr) != 0
-        raise "sodium_mprotect_readonly"
-      end
-      @state = State::Readonly
-      self
-    end
-
-    # Marks a region allocated using sodium_malloc() or sodium_allocarray() as readable and writable, after having been protected using sodium_mprotect_readonly() or sodium_mprotect_noaccess().
-    # WARNING: Not thread safe
-    def readwrite
-      raise Error::KeyWiped.new if @state == State::Wiped
+    protected def readwrite_impl : Nil
       if LibSodium.sodium_mprotect_readwrite(@ptr) != 0
         raise "sodium_mprotect_readwrite"
       end
-      @state = State::Readwrite
-      self
     end
 
-    # Timing safe memory compare.
-    def ==(other : self)
-      Sodium.memcmp self.to_slice, other.to_slice
-    end
-
-    # Timing safe memory compare.
-    def ==(other : Bytes)
-      Sodium.memcmp self.to_slice, other
-    end
-
-    # WARNING: Not thread safe
-    private def set_state(new_state : State)
-      return if @state == new_state
-
-      case new_state
-      when State::Readwrite; readwrite
-      when State::Readonly ; readonly
-      when State::Noaccess ; noaccess
-      when State::Wiped    ; raise Error::InvalidStateTransition.new
-      else
-        raise "unknown state #{new_state}"
+    protected def readonly_impl : Nil
+      if LibSodium.sodium_mprotect_readonly(@ptr) != 0
+        raise "sodium_mprotect_readonly"
       end
     end
 
-    # WARNING: Only thread safe when current state >= requested state
-    private def with_state(new_state : State)
-      old_state = @state
-      # Only change when new_state needs more access than @state.
-      if old_state >= new_state
-        yield
-      else
-        begin
-          set_state new_state
-          yield
-        ensure
-          set_state old_state
-        end
+    protected def noaccess_impl : Nil
+      if LibSodium.sodium_mprotect_noaccess(@ptr) != 0
+        raise "sodium_mprotect_noaccess"
       end
     end
   end