Keep slot-1 session in SignSequenceComplete

src/tpm2_wrap.c

@@ -6056,6 +6056,13 @@ int wolfTPM2_SignSequenceComplete(WOLFTPM2_DEV* dev,
     SignSequenceComplete_In signSeqCompleteIn;
     SignSequenceComplete_Out signSeqCompleteOut;
     WOLFTPM2_HANDLE seqHandleObj;
+#if MAX_SESSION_NUM > 1
+    TPM2_AUTH_SESSION session1;
+    int restore = 0;
+#endif
+#if MAX_SESSION_NUM > 2
+    int moved = 0;
+#endif
 
     if (dev == NULL || key == NULL || sig == NULL || sigSz == NULL) {
         return BAD_FUNC_ARG;
@@ -6077,6 +6084,26 @@ int wolfTPM2_SignSequenceComplete(WOLFTPM2_DEV* dev,
     XMEMSET(&seqHandleObj, 0, sizeof(seqHandleObj));
     seqHandleObj.hndl = sequenceHandle;
     wolfTPM2_SetAuthHandle(dev, 0, &seqHandleObj);
+#if MAX_SESSION_NUM > 1
+    /* SetAuthHandle resets slot 1 to password auth; restore the caller's
+     * slot-1 session (e.g. parameter encryption) after the command. While
+     * slot 2 is free, that session also protects this command from there.
+     * A policy-auth key runs on the slot-1 session itself; keep it live. */
+    if (!key->handle.policyAuth) {
+        XMEMCPY(&session1, &dev->session[1], sizeof(session1));
+        restore = 1;
+    #if MAX_SESSION_NUM > 2
+        if (TPM2_IS_HMAC_SESSION(session1.sessionHandle) &&
+                dev->session[2].sessionHandle == 0) {
+            XMEMCPY(&dev->session[2], &session1, sizeof(session1));
+            /* slot 2 authorizes no handle; a bound Name would fold the bind
+             * authValue into the param-enc key, which the TPM does not */
+            dev->session[2].name.size = 0;
+            moved = 1;
+        }
+    #endif
+    }
+#endif
     wolfTPM2_SetAuthHandle(dev, 1, &key->handle);
 
     XMEMSET(&signSeqCompleteIn, 0, sizeof(signSeqCompleteIn));
@@ -6089,6 +6116,24 @@ int wolfTPM2_SignSequenceComplete(WOLFTPM2_DEV* dev,
 
     XMEMSET(&signSeqCompleteOut, 0, sizeof(signSeqCompleteOut));
     rc = TPM2_SignSequenceComplete(&signSeqCompleteIn, &signSeqCompleteOut);
+#if MAX_SESSION_NUM > 1
+    if (restore) {
+    #if MAX_SESSION_NUM > 2
+        if (moved) {
+            /* keep the nonces the command advanced; a retired session
+             * stays cleared */
+            if (dev->session[2].sessionHandle == session1.sessionHandle) {
+                XMEMCPY(&dev->session[2].name, &session1.name,
+                    sizeof(session1.name));
+            }
+            XMEMCPY(&session1, &dev->session[2], sizeof(session1));
+            TPM2_ForceZero(&dev->session[2], sizeof(dev->session[2]));
+        }
+    #endif
+        XMEMCPY(&dev->session[1], &session1, sizeof(session1));
+        TPM2_ForceZero(&session1, sizeof(session1));
+    }
+#endif
     if (rc == TPM_RC_SUCCESS) {
         /* Extract signature based on algorithm */
         if (signSeqCompleteOut.signature.sigAlg == TPM_ALG_ECDSA ||

tests/unit_tests.c

@@ -10543,6 +10543,10 @@ static void test_wolfTPM2_MLDSA_SignSequence_NonEmptyAuth(WOLFTPM2_DEV* dev,
     int rc;
     WOLFTPM2_KEY key;
     TPMT_PUBLIC pub;
+#ifndef WOLFTPM2_NO_WOLFCRYPT
+    WOLFTPM2_SESSION sess;
+    TPM2B_NONCE nonceTPM;
+#endif
     static const byte gAuth[] = { 'p','q','c','_','a','u','t','h' };
     byte sig[5000];
     int sigSz = (int)sizeof(sig);
@@ -10562,12 +10566,139 @@ static void test_wolfTPM2_MLDSA_SignSequence_NonEmptyAuth(WOLFTPM2_DEV* dev,
     AssertIntEQ(rc, 0);
 
     test_wolfTPM2_MLDSA_SignSequence(dev, &key, gMsg, msgSz, sig, &sigSz);
+    if (sigSz == 0) {
+        /* sign sequence unsupported; the helper reported the skip */
+        wolfTPM2_UnloadHandle(dev, &key.handle);
+        return;
+    }
+
+#ifndef WOLFTPM2_NO_WOLFCRYPT
+    /* Signing must leave a caller's slot-1 param-encryption session in
+     * place; replacing it with password auth sends the next command's
+     * secrets (e.g. a Decapsulate shared secret) in the clear. */
+    rc = wolfTPM2_StartSession(dev, &sess, NULL, NULL, TPM_SE_HMAC,
+        TPM_ALG_CFB);
+    AssertIntEQ(rc, 0);
+    rc = wolfTPM2_SetAuthSession(dev, 1, &sess, (TPMA_SESSION_decrypt |
+        TPMA_SESSION_encrypt | TPMA_SESSION_continueSession));
+    AssertIntEQ(rc, 0);
+    XMEMCPY(&nonceTPM, &dev->session[1].nonceTPM, sizeof(nonceTPM));
+    sigSz = (int)sizeof(sig);
+    test_wolfTPM2_MLDSA_SignSequence(dev, &key, gMsg, msgSz, sig, &sigSz);
+    AssertIntEQ(dev->session[1].sessionHandle, sess.handle.hndl);
+#if MAX_SESSION_NUM > 2
+    /* With a free slot 2 the sign ran the session (nonce rolled) */
+    AssertIntNE(XMEMCMP(&nonceTPM, &dev->session[1].nonceTPM,
+        sizeof(nonceTPM)), 0);
+#endif
+    /* Verifying through the session checks the encrypted message and the
+     * session's nonces both survived. */
+    test_wolfTPM2_MLDSA_VerifySequence(dev, &key, gMsg, msgSz, sig, sigSz);
+    wolfTPM2_SetAuthSession(dev, 1, NULL, 0);
+    wolfTPM2_UnloadHandle(dev, &sess.handle);
+
+    /* Same with a session bound to the key and slot 1 holding the key's
+     * Name (as the NV wrappers leave it): while the session runs from
+     * slot 2 the bind authValue must stay out of the param-enc key. */
+    rc = wolfTPM2_StartSession(dev, &sess, NULL, &key.handle, TPM_SE_HMAC,
+        TPM_ALG_CFB);
+    AssertIntEQ(rc, 0);
+    rc = wolfTPM2_SetAuthSession(dev, 1, &sess, (TPMA_SESSION_decrypt |
+        TPMA_SESSION_encrypt | TPMA_SESSION_continueSession));
+    AssertIntEQ(rc, 0);
+    rc = wolfTPM2_SetAuthHandleName(dev, 1, &key.handle);
+    AssertIntEQ(rc, 0);
+    sigSz = (int)sizeof(sig);
+    test_wolfTPM2_MLDSA_SignSequence(dev, &key, gMsg, msgSz, sig, &sigSz);
+    AssertIntEQ(dev->session[1].sessionHandle, sess.handle.hndl);
+    AssertIntEQ(dev->session[1].name.size, key.handle.name.size);
+    /* Verify with password auth: a wrong decrypt above signed other bytes */
+    wolfTPM2_SetAuthSession(dev, 1, NULL, 0);
+    test_wolfTPM2_MLDSA_VerifySequence(dev, &key, gMsg, msgSz, sig, sigSz);
+    wolfTPM2_UnloadHandle(dev, &sess.handle);
+#endif
 
     wolfTPM2_UnloadHandle(dev, &key.handle);
     printf("Test TPM Wrapper: %-40s Passed\n",
         "ML-DSA Sign Seq w/ key auth:");
 }
 
+#ifndef WOLFTPM2_NO_WOLFCRYPT
+/* A policy-authorized key signs on the slot-1 policy session itself, so
+ * SignSequenceComplete must leave that session in slot 1 with the nonces
+ * the command rolled; a second sign on it fails if they went stale. */
+static void test_wolfTPM2_MLDSA_SignSequence_PolicyAuth(WOLFTPM2_DEV* dev,
+    const TPMT_PUBLIC* mldsaPub)
+{
+    int rc;
+    WOLFTPM2_KEY key;
+    WOLFTPM2_SESSION sess;
+    TPMT_PUBLIC pub;
+    TPM2B_NONCE nonceTPM;
+    word32 polSz;
+    static const byte gAuth[] = { 'p','o','l','_','a','u','t','h' };
+    byte sig[5000];
+    int sigSz = (int)sizeof(sig);
+    static const byte gMsg[] = "Policy-auth ML-DSA test message";
+    int msgSz = (int)sizeof(gMsg) - 1;
+
+    XMEMSET(&key, 0, sizeof(key));
+    XMEMCPY(&pub, mldsaPub, sizeof(pub));
+    /* policy-only: authPolicy = PolicyAuthValue, no userWithAuth */
+    pub.objectAttributes &= ~TPMA_OBJECT_userWithAuth;
+    XMEMSET(pub.authPolicy.buffer, 0, sizeof(pub.authPolicy.buffer));
+    polSz = (word32)TPM2_GetHashDigestSize(pub.nameAlg);
+    AssertIntEQ(wolfTPM2_PolicyHash(pub.nameAlg, pub.authPolicy.buffer,
+        &polSz, TPM_CC_PolicyAuthValue, NULL, 0), 0);
+    pub.authPolicy.size = (UINT16)polSz;
+
+    rc = wolfTPM2_CreatePrimaryKey(dev, &key, TPM_RH_OWNER, &pub,
+        gAuth, (int)sizeof(gAuth));
+    if (TEST_V185_UNSUPPORTED(rc, TPM_ALG_MLDSA)) {
+        printf("Test TPM Wrapper: %-40s Skipped (not supported)\n",
+            "ML-DSA Sign Seq w/ policy auth:");
+        return;
+    }
+    AssertIntEQ(rc, 0);
+    key.handle.policyAuth = 1;
+
+    /* bound: the PolicyAuthValue HMAC key layout needs a session key */
+    rc = wolfTPM2_StartSession(dev, &sess, NULL, &key.handle, TPM_SE_POLICY,
+        TPM_ALG_NULL);
+    AssertIntEQ(rc, 0);
+    AssertIntEQ(wolfTPM2_PolicyAuthValue(dev, &sess, gAuth,
+        (int)sizeof(gAuth)), 0);
+    rc = wolfTPM2_SetAuthSession(dev, 1, &sess, TPMA_SESSION_continueSession);
+    AssertIntEQ(rc, 0);
+    XMEMCPY(&nonceTPM, &dev->session[1].nonceTPM, sizeof(nonceTPM));
+
+    test_wolfTPM2_MLDSA_SignSequence(dev, &key, gMsg, msgSz, sig, &sigSz);
+    if (sigSz == 0) {
+        /* sign sequence unsupported; the helper reported the skip */
+        goto done;
+    }
+    AssertIntEQ(dev->session[1].sessionHandle, sess.handle.hndl);
+    AssertIntNE(XMEMCMP(&nonceTPM, &dev->session[1].nonceTPM,
+        sizeof(nonceTPM)), 0);
+
+    /* the TPM reset the policy digest on use; satisfy it again and sign */
+    AssertIntEQ(wolfTPM2_PolicyAuthValue(dev, &sess, NULL, 0), 0);
+    sigSz = (int)sizeof(sig);
+    test_wolfTPM2_MLDSA_SignSequence(dev, &key, gMsg, msgSz, sig, &sigSz);
+    AssertIntEQ(dev->session[1].sessionHandle, sess.handle.hndl);
+
+    wolfTPM2_SetAuthSession(dev, 1, NULL, 0);
+    test_wolfTPM2_MLDSA_VerifySequence(dev, &key, gMsg, msgSz, sig, sigSz);
+    printf("Test TPM Wrapper: %-40s Passed\n",
+        "ML-DSA Sign Seq w/ policy auth:");
+
+done:
+    wolfTPM2_SetAuthSession(dev, 1, NULL, 0);
+    wolfTPM2_UnloadHandle(dev, &sess.handle);
+    wolfTPM2_UnloadHandle(dev, &key.handle);
+}
+#endif
+
 /* Regression for the VerifySequenceComplete data-chain fix.
  *
  * The wrapper used to silently drop the data/dataSz arguments; the fix
@@ -11064,6 +11195,9 @@ static void test_wolfTPM2_PQC(void)
     }
     test_TPM2_SignSequenceStart_NoSession(&dev, &mldsaKey);
     test_wolfTPM2_MLDSA_SignSequence_NonEmptyAuth(&dev, &mldsaPub);
+#ifndef WOLFTPM2_NO_WOLFCRYPT
+    test_wolfTPM2_MLDSA_SignSequence_PolicyAuth(&dev, &mldsaPub);
+#endif
 
     wolfTPM2_UnloadHandle(&dev, &mldsaKey.handle);
 mldsa_done: