tpm2: re-use tpm connection when creating keydata with passphrase

Signed-off-by: Zeyad Gouda <[email protected]>
ZeyadYasser · Dec 5, 2024 · 56aea41 · 56aea41
1 parent 1224447
commit 56aea41
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 13 deletions.
diff --git a/tpm2/seal.go b/tpm2/seal.go
@@ -79,19 +79,24 @@ func makeKeyDataNoAuth(skd *SealedKeyData, role string, encryptedPayload []byte,
 	})
 }
 
-func makeKeyDataWithPassphraseConstructor(kdfOptions secboot.KDFOptions, passphrase string) keyDataConstructor {
+func makeKeyDataWithPassphraseConstructor(tpm *Connection, kdfOptions secboot.KDFOptions, passphrase string) keyDataConstructor {
 	return func(skd *SealedKeyData, role string, encryptedPayload []byte, kdfAlg crypto.Hash) (*secboot.KeyData, error) {
-		return secbootNewKeyDataWithPassphrase(&secboot.KeyWithPassphraseParams{
-			KeyParams: secboot.KeyParams{
-				Handle:           skd,
-				Role:             role,
-				EncryptedPayload: encryptedPayload,
-				PlatformName:     platformName,
-				KDFAlg:           kdfAlg,
-			},
-			KDFOptions:  kdfOptions,
-			AuthKeySize: skd.data.Public().NameAlg.Size(),
-		}, passphrase)
+		var keydata *secboot.KeyData
+		var err error
+		withTPMConnection(tpm, func() {
+			keydata, err = secbootNewKeyDataWithPassphrase(&secboot.KeyWithPassphraseParams{
+				KeyParams: secboot.KeyParams{
+					Handle:           skd,
+					Role:             role,
+					EncryptedPayload: encryptedPayload,
+					PlatformName:     platformName,
+					KDFAlg:           kdfAlg,
+				},
+				KDFOptions:  kdfOptions,
+				AuthKeySize: skd.data.Public().NameAlg.Size(),
+			}, passphrase)
+		})
+		return keydata, err
 	}
 }
 
@@ -304,5 +309,5 @@ func NewTPMPassphraseProtectedKey(tpm *Connection, params *PassphraseProtectKeyP
 		AuthMode:               secboot.AuthModePassphrase,
 		Role:                   params.Role,
 		PcrProfile:             params.PCRProfile,
-	}, sealer, makeKeyDataWithPassphraseConstructor(params.KDFOptions, passphrase), tpm.HmacSession())
+	}, sealer, makeKeyDataWithPassphraseConstructor(tpm, params.KDFOptions, passphrase), tpm.HmacSession())
 }
diff --git a/tpm2/tpm.go b/tpm2/tpm.go
@@ -185,3 +185,12 @@ func ConnectToDefaultTPM() (*Connection, error) {
 // ConnectToDefaultTPM. This can be overridden with a custom connection
 // function.
 var ConnectToTPM func() (*Connection, error) = ConnectToDefaultTPM
+
+func withTPMConnection(tpm *Connection, fn func()) {
+	old := ConnectToTPM
+	ConnectToTPM = func() (*Connection, error) {
+		return tpm, nil
+	}
+	defer func() { ConnectToTPM = old }()
+	fn()
+}