From b212f6bacd16a9ede403cc7b1ea9d5b7c4d3032c Mon Sep 17 00:00:00 2001
From: Didactic Drunk <1479616+didactic-drunk@users.noreply.github.com>
Date: Fri, 13 Sep 2019 21:33:21 -0700
Subject: [PATCH] Fix Sodium::SecureBuffer state transitions.

---
 spec/sodium/secure_buffer_spec.cr |  6 +++---
 src/sodium/secret_box.cr          | 24 +++++++++++++-----------
 src/sodium/secure_buffer.cr       | 13 +++++++++----
 3 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/spec/sodium/secure_buffer_spec.cr b/spec/sodium/secure_buffer_spec.cr
index dba4d09..6544047 100644
--- a/spec/sodium/secure_buffer_spec.cr
+++ b/spec/sodium/secure_buffer_spec.cr
@@ -51,17 +51,17 @@ describe Sodium::SecureBuffer do
 
     buf.noaccess
     buf.@state.should eq Sodium::SecureBuffer::State::Noaccess
-    buf.readonly { }
+    buf.readonly { buf.@state.should eq Sodium::SecureBuffer::State::Readonly }
     buf.@state.should eq Sodium::SecureBuffer::State::Noaccess
 
     buf.readonly
     buf.@state.should eq Sodium::SecureBuffer::State::Readonly
-    buf.readwrite { }
+    buf.readwrite { buf.@state.should eq Sodium::SecureBuffer::State::Readwrite }
     buf.@state.should eq Sodium::SecureBuffer::State::Readonly
 
     buf.readwrite
     buf.@state.should eq Sodium::SecureBuffer::State::Readwrite
-    buf.readonly { }
+    buf.readonly { buf.@state.should eq Sodium::SecureBuffer::State::Readwrite }
     buf.@state.should eq Sodium::SecureBuffer::State::Readwrite
 
     buf.wipe
diff --git a/src/sodium/secret_box.cr b/src/sodium/secret_box.cr
index 628b3c7..98deb0e 100644
--- a/src/sodium/secret_box.cr
+++ b/src/sodium/secret_box.cr
@@ -21,19 +21,19 @@ module Sodium
     MAC_SIZE   = LibSodium.crypto_secretbox_macbytes.to_i
 
     # Returns key
-    delegate to_slice, to: @buf
+    delegate to_slice, to: @key
 
     # Generate a new random key held in a SecureBuffer.
     def initialize
-      @buf = SecureBuffer.random KEY_SIZE
+      @key = SecureBuffer.random KEY_SIZE
     end
 
     # Use an existing SecureBuffer.
-    def initialize(@buf : SecureBuffer)
-      if @buf.bytesize != KEY_SIZE
-        raise ArgumentError.new("Secret key must be #{KEY_SIZE} bytes, got #{@buf.bytesize}")
+    def initialize(@key : SecureBuffer)
+      if @key.bytesize != KEY_SIZE
+        raise ArgumentError.new("Secret key must be #{KEY_SIZE} bytes, got #{@key.bytesize}")
       end
-      @buf.readonly
+      @key.readonly
     end
 
     # Copy bytes to a new SecureBuffer
@@ -43,7 +43,7 @@ module Sodium
       if bytes.bytesize != KEY_SIZE
         raise ArgumentError.new("Secret key must be #{KEY_SIZE} bytes, got #{bytes.bytesize}")
       end
-      @buf = SecureBuffer.new bytes, erase: erase
+      @key = SecureBuffer.new bytes, erase: erase
     end
 
     # Encrypts data and returns {ciphertext, nonce}
@@ -59,9 +59,10 @@ module Sodium
         raise ArgumentError.new("dst.bytesize must be src.bytesize + MAC_SIZE, got #{dst.bytesize}")
       end
       nonce.used!
-      if LibSodium.crypto_secretbox_easy(dst, src, src.bytesize, nonce.to_slice, self.to_slice) != 0
-        raise Sodium::Error.new("crypto_secretbox_easy")
+      r = @key.readonly do
+        LibSodium.crypto_secretbox_easy(dst, src, src.bytesize, nonce.to_slice, @key)
       end
+      raise Sodium::Error.new("crypto_secretbox_easy") if r != 0
       {dst, nonce}
     end
 
@@ -80,9 +81,10 @@ module Sodium
       if dst.bytesize != (src.bytesize - MAC_SIZE)
         raise ArgumentError.new("dst.bytesize must be src.bytesize - MAC_SIZE, got #{dst.bytesize}")
       end
-      if LibSodium.crypto_secretbox_open_easy(dst, src, src.bytesize, nonce.to_slice, self.to_slice) != 0
-        raise Sodium::Error::DecryptionFailed.new("crypto_secretbox_easy")
+      r = @key.readonly do
+        LibSodium.crypto_secretbox_open_easy(dst, src, src.bytesize, nonce.to_slice, @key)
       end
+      raise Sodium::Error::DecryptionFailed.new("crypto_secretbox_easy") if r != 0
       dst
     end
 
diff --git a/src/sodium/secure_buffer.cr b/src/sodium/secure_buffer.cr
index 2a6b90d..05713a3 100644
--- a/src/sodium/secure_buffer.cr
+++ b/src/sodium/secure_buffer.cr
@@ -68,7 +68,12 @@ module Sodium
     # Returns key
     # May permanently set key to readonly depending on class usage.
     def to_slice
-      readonly if @state == State::Noaccess
+      case @state
+      when State::Noaccess, State::Wiped
+        readonly
+      else
+        # Ok
+      end
       Slice(UInt8).new @ptr, @bytesize
     end
 
@@ -154,16 +159,16 @@ module Sodium
     end
 
     private def with_state(new_state : State)
+      old_state = @state
       # Only change when new_state needs more access than @state.
-      if new_state >= @state
+      if old_state >= new_state
         yield
       else
         begin
-          old_state = @state
           set_state new_state
           yield
         ensure
-          set_state old_state if old_state
+          set_state old_state
         end
       end
     end