From a4c5d05556ecd450acce5e20fcab7af5f282af2f Mon Sep 17 00:00:00 2001
From: Charles McGarvey <ccm@cpan.org>
Date: Tue, 16 Aug 2022 16:51:16 -0600
Subject: [PATCH] Fix memory protection of brand new databases

Fixes #4.
---
 Changes                |  1 +
 lib/File/KDBX.pm       | 20 ++++++++++++++------
 lib/File/KDBX/Entry.pm | 10 ++++++----
 t/entry.t              | 39 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 60 insertions(+), 10 deletions(-)

diff --git a/Changes b/Changes
index dcce162..a20e1c0 100644
--- a/Changes
+++ b/Changes
@@ -2,6 +2,7 @@ Revision history for File-KDBX.
 
 {{$NEXT}}
   * Fixed transform_rounds method to work with Argon KDF. Thanks HIGHTOWE.
+  * Fixed bug preventing memory protection on new databases. Thanks HIGHTOWE.
 
 0.905     2022-08-06 12:12:42-0600
   * Declared Time::Local 1.19 as a required dependency.
diff --git a/lib/File/KDBX.pm b/lib/File/KDBX.pm
index db40366..82796f3 100644
--- a/lib/File/KDBX.pm
+++ b/lib/File/KDBX.pm
@@ -1176,16 +1176,24 @@ sub _remove_safe { delete $SAFE{$_[0]} }
 sub lock {
     my $self = shift;
 
-    $self->_safe and return $self;
-
+    # Find things to lock:
     my @strings;
-
     $self->entries(history => 1)->each(sub {
-        push @strings, grep { $_->{protect} } values %{$_->strings}, values %{$_->binaries};
+        my $strings = $_->strings;
+        for my $string_key (keys %$strings) {
+            my $string = $strings->{$string_key};
+            push @strings, $string if $string->{protect} // $self->memory_protection($string_key);
+        }
+        push @strings, grep { $_->{protect} } values %{$_->binaries};
     });
+    return $self if !@strings;  # nothing to do
 
-    $self->_safe(File::KDBX::Safe->new(\@strings));
-
+    if (my $safe = $self->_safe) {
+        $safe->add(\@strings);
+    }
+    else {
+        $self->_safe(File::KDBX::Safe->new(\@strings));
+    }
     return $self;
 }
 
diff --git a/lib/File/KDBX/Entry.pm b/lib/File/KDBX/Entry.pm
index 0119b67..6a927c1 100644
--- a/lib/File/KDBX/Entry.pm
+++ b/lib/File/KDBX/Entry.pm
@@ -253,14 +253,16 @@ sub string {
 
     return $self->{strings}{$key} = $args{value} if is_plain_hashref($args{value});
 
+    # Auto-vivify the standard strings.
+    if (!exists $self->{strings}{$key} && $STANDARD_STRINGS{$key}) {
+        $args{value} //= '';
+        $args{protect} //= true if $self->_protect($key);
+    }
+
     while (my ($field, $value) = each %args) {
         $self->{strings}{$key}{$field} = $value;
     }
 
-    # Auto-vivify the standard strings.
-    if ($STANDARD_STRINGS{$key}) {
-        return $self->{strings}{$key} //= {value => '', $self->_protect($key) ? (protect => true) : ()};
-    }
     return $self->{strings}{$key};
 }
 
diff --git a/t/entry.t b/t/entry.t
index 1581608..a5700c9 100644
--- a/t/entry.t
+++ b/t/entry.t
@@ -84,6 +84,17 @@ subtest 'Accessors' => sub {
     $entry->creation_time('2022-02-02 12:34:56');
     cmp_ok $entry->creation_time->epoch, '==', 1643805296, 'Creation time coerced into a Time::Piece (epoch)';
     is $entry->creation_time->datetime, '2022-02-02T12:34:56', 'Creation time coerced into a Time::Piece';
+
+    $entry->username('foo');
+    cmp_deeply $entry->strings->{UserName}, {
+        value   => 'foo',
+    }, 'Username setter works';
+
+    $entry->password('bar');
+    cmp_deeply $entry->strings->{Password}, {
+        value   => 'bar',
+        protect => bool(1),
+    }, 'Password setter works';
 };
 
 subtest 'Custom icons' => sub {
@@ -169,4 +180,32 @@ subtest 'Auto-type' => sub {
     is $keys, 'blah', 'Select the correct association';
 };
 
+subtest 'Memory protection' => sub {
+    my $kdbx = File::KDBX->new;
+
+    is exception { $kdbx->lock }, undef, 'Can lock empty database';
+    $kdbx->unlock;  # should be no-op since nothing was locked
+
+    my $entry = $kdbx->root->add_entry(
+        title    => 'My Bank',
+        username => 'mreynolds',
+        password => 's3cr3t',
+    );
+    $entry->string(Custom => 'foo', protect => 1);
+    $entry->binary(Binary => 'bar', protect => 1);
+    $entry->binary(UnprotectedBinary => 'baz');
+
+    is exception { $kdbx->lock }, undef, 'Can lock new database';
+    is $entry->username, 'mreynolds', 'UserName does not get locked';
+    is $entry->password, undef, 'Password is lockable';
+    is $entry->string_value('Custom'), undef, 'Custom is lockable';
+    is $entry->binary_value('Binary'), undef, 'Binary is lockable';
+    is $entry->binary_value('UnprotectedBinary'), 'baz', 'Unprotected binary does not get locked';
+
+    $kdbx->unlock;
+    is $entry->password, 's3cr3t', 'Password is unlockable';
+    is $entry->string_value('Custom'), 'foo', 'Custom is unlockable';
+    is $entry->binary_value('Binary'), 'bar', 'Binary is unlockable';
+};
+
 done_testing;
-- 
2.45.2