From 4b762f9cd0a68ed547096b678487c0830a2fd16f Mon Sep 17 00:00:00 2001
From: Zeke Gabrielse <ezekg@yahoo.com>
Date: Fri, 12 Jul 2024 12:00:40 -0500
Subject: [PATCH] add export data validation

---
 lib/keygen/exporter.rb                 |  2 +-
 lib/keygen/importer.rb                 |  6 +++---
 lib/keygen/importer/v1/deserializer.rb |  5 +++++
 lib/keygen/importer/v1/reader.rb       |  8 +++++++-
 spec/lib/keygen/importer_spec.rb       | 16 ++++++++++++++++
 5 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/lib/keygen/exporter.rb b/lib/keygen/exporter.rb
index 546b0c4eec..a1e7aa60f2 100644
--- a/lib/keygen/exporter.rb
+++ b/lib/keygen/exporter.rb
@@ -28,7 +28,7 @@ def exporter_class_for(version:)
       when 1
         V1::Exporter
       else
-        raise UnsupportedVersionError.new, "Unsupported export version: #{version}"
+        raise UnsupportedVersionError.new, "unsupported export version: #{version}"
       end
     end
   end
diff --git a/lib/keygen/importer.rb b/lib/keygen/importer.rb
index 7b76a907c8..3452a4c81f 100644
--- a/lib/keygen/importer.rb
+++ b/lib/keygen/importer.rb
@@ -8,6 +8,8 @@ module Keygen
   module Importer
     class UnsupportedVersionError < StandardError; end
     class InvalidSecretKeyError < StandardError; end
+    class InvalidDataError < StandardError; end
+    class InvalidChunkError < InvalidDataError; end
 
     extend self
 
@@ -19,8 +21,6 @@ def import(from:, secret_key: nil)
       importer       = importer_class.new(secret_key:)
 
       importer.import(from:)
-    rescue OpenSSL::Cipher::CipherError
-      raise InvalidSecretKeyError.new, 'Secret key is invalid'
     end
 
     private
@@ -30,7 +30,7 @@ def importer_class_for(version:)
       when 1
         V1::Importer
       else
-        raise UnsupportedVersionError.new, "Unsupported import version: #{version}"
+        raise UnsupportedVersionError.new, "unsupported import version: #{version}"
       end
     end
   end
diff --git a/lib/keygen/importer/v1/deserializer.rb b/lib/keygen/importer/v1/deserializer.rb
index 36ab6e794a..e70981022f 100644
--- a/lib/keygen/importer/v1/deserializer.rb
+++ b/lib/keygen/importer/v1/deserializer.rb
@@ -16,6 +16,11 @@ def deserialize(data)
           unencrypted  = decrypt(decompressed)
 
           unpack(unencrypted)
+        rescue MessagePack::MalformedFormatError,
+               OpenSSL::Cipher::CipherError
+          raise InvalidSecretKeyError.new, 'secret key is invalid'
+        rescue Zlib::Error
+          raise InvalidDataError.new, 'data is invalid'
         end
 
         private
diff --git a/lib/keygen/importer/v1/reader.rb b/lib/keygen/importer/v1/reader.rb
index 9c2d53fb7e..a08ff71b21 100644
--- a/lib/keygen/importer/v1/reader.rb
+++ b/lib/keygen/importer/v1/reader.rb
@@ -11,7 +11,13 @@ def read_chunk
           chunk_size = read_chunk_size
           return if chunk_size.zero?
 
-          read(chunk_size)
+          chunk = read(chunk_size)
+
+          unless chunk.size == chunk_size
+            raise InvalidChunkError.new, 'chunk is invalid'
+          end
+
+          chunk
         end
       end
     end
diff --git a/spec/lib/keygen/importer_spec.rb b/spec/lib/keygen/importer_spec.rb
index 88f834b6ac..bfec9dad7c 100644
--- a/spec/lib/keygen/importer_spec.rb
+++ b/spec/lib/keygen/importer_spec.rb
@@ -28,6 +28,22 @@
       .to raise_error Keygen::Importer::UnsupportedVersionError
   end
 
+  it 'should raise for invalid data' do
+    data = 'zzz'
+    size = [data.bytesize].pack('Q>')
+
+    expect { Keygen::Importer.import(from: StringIO.new(1.chr + size + data)) }
+      .to raise_error Keygen::Importer::InvalidDataError
+  end
+
+  it 'should raise for invalid chunk' do
+    data = 'zz'
+    size = ['zzz'.bytesize].pack('Q>')
+
+    expect { Keygen::Importer.import(from: StringIO.new(1.chr + size + data)) }
+      .to raise_error Keygen::Importer::InvalidChunkError
+  end
+
   context 'with encryption' do
     let(:secret_key) { SecureRandom.hex }