2 files changed, 79 insertions, 4 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;
     }
diff --git a/test/storage/StorageSupport.mc b/test/storage/StorageSupport.mc
index 51a0c7c..5f2ead3 100644
--- a/test/storage/StorageSupport.mc
+++ b/test/storage/StorageSupport.mc
@@ -143,8 +143,9 @@ class FakeTime extends TimeSource {
 (:test)
 module StorageFixtures {
 
-    const PIN = "246810";
-    const OTHER_PIN = "135799";
+    const PIN = "2468"; // KeyManager.PIN_DIGITS digits, like every new PIN
+    const OTHER_PIN = "1357";
+    const LEGACY_PIN = "246810"; // a 6-digit PIN from before the 4-digit rule
     const ITERATIONS = 2000; // Kdf.MIN_ITERATIONS: the cheapest allowed
 
     // A fake TOTP entry (the secret is the RFC 4226 test key, not a real one).
@@ -205,6 +206,22 @@ module StorageFixtures {
         vault.finishSetPin(job);
     }
 
+    // Store a wrap the way versions before the 4-digit rule did: any 4 to
+    // 10 digits and no PIN-length marker. The session stays locked.
+    function setLegacyPin(store as MemoryBackend, text as String) as Void {
+        var keyManager = new KeyManager(store);
+        var job = new PinJob(pin(text), new [KeyManager.SALT_BYTES]b, ITERATIONS);
+        Jobs.runToCompletion(job, Kdf.ITERATIONS_PER_STEP);
+        keyManager.storeWrap(job, keyManager.newMasterKey());
+        store.deleteValue(KeyManager.PIN_LENGTH_KEY);
+    }
+
+    // True when the store holds the 4-digit PIN marker.
+    function pinMarked(store as MemoryBackend) as Boolean {
+        var marker = store.getValue(KeyManager.PIN_LENGTH_KEY);
+        return marker instanceof Number && marker == KeyManager.PIN_DIGITS;
+    }
+
     // One unlock attempt; true when the PIN was right.
     function unlock(vault as Vault, text as String) as Boolean {
         var job = vault.beginUnlock(pin(text));