1 files changed, 60 insertions, 2 deletions

diff --git a/test/storage/PinTest.mc b/test/storage/PinTest.mc
index 0212936..55adaa2 100644
--- a/test/storage/PinTest.mc
+++ b/test/storage/PinTest.mc
@@ -1,5 +1,5 @@
 // Author:  cirakg & claude ai
-// Summary: Tests for PINs: create, unlock, wrong PIN, the lockout schedule on a fake clock, erase, change, reset, auto-lock.
+// Summary: Tests for PINs: create, unlock, wrong PIN, 4-digit rule, legacy PINs, the lockout schedule on a fake clock, erase, change, reset, auto-lock.
 
 import Toybox.Lang;
 import Toybox.Test;
@@ -39,6 +39,64 @@ function pinRejectsMalformedPins(logger as Test.Logger) as Boolean {
     return KeyManager.isValidPin(Bytes.fromUtf8("0000")) && KeyManager.isValidPin(Bytes.fromUtf8("0123456789"));
 }
 
+// A new PIN has exactly KeyManager.PIN_DIGITS digits; other lengths are
+// refused before any derivation.
+(:test)
+function pinSetRequiresFourDigits(logger as Test.Logger) as Boolean {
+    var bad = ["123", "12345", "123456", "0123456789", "12a4"];
+    for (var i = 0; i < bad.size(); i++) {
+        if (KeyManager.isValidNewPin(Bytes.fromUtf8(bad[i]))) {
+            logger.error("accepted new PIN '" + bad[i] + "'");
+            return false;
+        }
+        var code = 0;
+        try {
+            new Vault(new MemoryBackend(), new FakeTime()).beginSetPin(StorageFixtures.pin(bad[i]), StorageFixtures.ITERATIONS);
+        } catch (err instanceof VaultError) {
+            code = err.code;
+        }
+        if (code != VaultError.INVALID) {
+            logger.error("set PIN '" + bad[i] + "' gave code " + code);
+            return false;
+        }
+    }
+    return KeyManager.isValidNewPin(Bytes.fromUtf8("0000")) && KeyManager.isValidNewPin(Bytes.fromUtf8("9876"));
+}
+
+// Setting or changing the PIN stores the 4-digit marker; forgetting the
+// PIN removes it.
+(:test)
+function pinMarkerIsSetBySetAndChange(logger as Test.Logger) as Boolean {
+    var store = new MemoryBackend();
+    var vault = new Vault(store, new FakeTime());
+    StorageFixtures.setPin(vault, StorageFixtures.PIN);
+    var afterSet = StorageFixtures.pinMarked(store) && !vault.isLegacyPin();
+    store.deleteValue(KeyManager.PIN_LENGTH_KEY);
+    var legacyWithout = vault.isLegacyPin();
+    StorageFixtures.setPin(vault, StorageFixtures.OTHER_PIN); // change while unlocked
+    var afterChange = StorageFixtures.pinMarked(store) && !vault.isLegacyPin();
+    vault.resetPin();
+    return afterSet && legacyWithout && afterChange && store.getValue(KeyManager.PIN_LENGTH_KEY) == null
+        && !vault.isLegacyPin();
+}
+
+// A 6-digit PIN from before the 4-digit rule still unlocks (and a wrong
+// one fails without an exception); a new 4-digit PIN then replaces it.
+(:test)
+function pinLegacySixDigitWrapStillUnlocks(logger as Test.Logger) as Boolean {
+    var store = new MemoryBackend();
+    StorageFixtures.setLegacyPin(store, StorageFixtures.LEGACY_PIN);
+    var vault = new Vault(store, new FakeTime());
+    var legacy = vault.hasPin() && vault.isLegacyPin();
+    var wrong = !StorageFixtures.unlock(vault, "2468101");
+    var right = StorageFixtures.unlock(vault, StorageFixtures.LEGACY_PIN) && vault.session.isUnlocked();
+    StorageFixtures.setPin(vault, StorageFixtures.PIN);
+    vault.lock();
+    var oldFails = !StorageFixtures.unlock(vault, StorageFixtures.LEGACY_PIN);
+    var newWorks = StorageFixtures.unlock(vault, StorageFixtures.PIN);
+    return legacy && wrong && right && !vault.isLegacyPin() && oldFails && newWorks;
+}
+
 (:test)
 function lockoutDelaySchedule(logger as Test.Logger) as Boolean {
     var expected = [0, 0, 0, 0, 0, 30, 60, 120, 240, 480, 960, 1920, 3600, 3600, 3600];
@@ -170,7 +228,7 @@ function pinChangeKeepsEntries(logger as Test.Logger) as Boolean {
     vault.lock();
     var lockedCode = 0;
     try {
-        vault.beginSetPin(StorageFixtures.pin("999999"), StorageFixtures.ITERATIONS);
+        vault.beginSetPin(StorageFixtures.pin("9999"), StorageFixtures.ITERATIONS);
     } catch (err instanceof VaultError) {
         lockedCode = err.code;
     }