Add comment and wrap AddCleanup

Author: Michelle Clayton

cng/aes.go

@@ -30,7 +30,7 @@ func NewAESCipher(key []byte) (cipher.Block, error) {
 		return nil, err
 	}
 	c := &aesCipher{kh: kh, key: bytes.Clone(key)}
-	runtime.AddCleanup(c, destroyKey, kh)
+	addCleanupKey(c, kh)
 	return c, nil
 }
 
@@ -161,7 +161,7 @@ func newCBC(encrypt bool, alg string, key, iv []byte) *cbcCipher {
 		panic(err)
 	}
 	x := &cbcCipher{kh: kh, encrypt: encrypt, blockSize: blockSize}
-	runtime.AddCleanup(x, destroyKey, kh)
+	addCleanupKey(x, kh)
 	x.SetIV(iv)
 	return x
 }
@@ -245,7 +245,7 @@ func newGCM(key []byte, tls cipherGCMTLS) (*aesGCM, error) {
 		return nil, err
 	}
 	g := &aesGCM{kh: kh, tls: tls}
-	runtime.AddCleanup(g, destroyKey, kh)
+	addCleanupKey(g, kh)
 	return g, nil
 }
 

cng/chacha20poly1305.go

@@ -40,7 +40,7 @@ func NewChaCha20Poly1305(key []byte) (cipher.AEAD, error) {
 		return nil, err
 	}
 	c := &chacha20poly1305{kh: kh}
-	runtime.AddCleanup(c, destroyKey, kh)
+	addCleanupKey(c, kh)
 	return c, nil
 }
 

cng/des.go

@@ -29,7 +29,7 @@ func NewDESCipher(key []byte) (cipher.Block, error) {
 		return nil, err
 	}
 	c := &desCipher{kh: kh, alg: bcrypt.DES_ALGORITHM, key: bytes.Clone(key)}
-	runtime.AddCleanup(c, destroyKey, kh)
+	addCleanupKey(c, kh)
 	return c, nil
 }
 
@@ -39,7 +39,7 @@ func NewTripleDESCipher(key []byte) (cipher.Block, error) {
 		return nil, err
 	}
 	c := &desCipher{kh: kh, alg: bcrypt.DES3_ALGORITHM, key: bytes.Clone(key)}
-	runtime.AddCleanup(c, destroyKey, kh)
+	addCleanupKey(c, kh)
 	return c, nil
 }
 

cng/dsa.go

@@ -147,7 +147,7 @@ func NewPrivateKeyDSA(params DSAParameters, X, Y BigInt) (*PrivateKeyDSA, error)
 		return nil, err
 	}
 	k := &PrivateKeyDSA{params, X, Y, hkey}
-	runtime.AddCleanup(k, destroyKey, hkey)
+	addCleanupKey(k, hkey)
 	return k, nil
 }
 
@@ -166,7 +166,7 @@ func NewPublicKeyDSA(params DSAParameters, Y BigInt) (*PublicKeyDSA, error) {
 		return nil, err
 	}
 	k := &PublicKeyDSA{params, Y, hkey}
-	runtime.AddCleanup(k, destroyKey, hkey)
+	addCleanupKey(k, hkey)
 	return k, nil
 }
 

cng/ecdh.go

@@ -128,7 +128,7 @@ func GenerateKeyECDH(curve string) (*PrivateKeyECDH, []byte, error) {
 	bytes = bytes[hdr.KeySize*2:]
 
 	k := &PrivateKeyECDH{hkey, isNIST(curve)}
-	runtime.AddCleanup(k, destroyKey, hkey)
+	addCleanupKey(k, hkey)
 	return k, bytes, nil
 }
 
@@ -163,7 +163,7 @@ func NewPublicKeyECDH(curve string, bytes []byte) (*PublicKeyECDH, error) {
 		return nil, err
 	}
 	k := &PublicKeyECDH{hkey, append([]byte(nil), bytes...), nil}
-	runtime.AddCleanup(k, destroyKey, hkey)
+	addCleanupKey(k, hkey)
 	return k, nil
 }
 
@@ -192,7 +192,7 @@ func NewPrivateKeyECDH(curve string, key []byte) (*PrivateKeyECDH, error) {
 		return nil, err
 	}
 	k := &PrivateKeyECDH{hkey, nist}
-	runtime.AddCleanup(k, destroyKey, hkey)
+	addCleanupKey(k, hkey)
 	return k, nil
 }
 
@@ -210,6 +210,8 @@ func (k *PrivateKeyECDH) PublicKey() (*PublicKeyECDH, error) {
 		// Only include X.
 		bytes = data[:hdr.KeySize]
 	}
+	// No cleanup needed: pub.priv prevents k from being garbage collected,
+	// so k's cleanup will destroy the shared hkey when both are unreachable.
 	pub := &PublicKeyECDH{k.hkey, bytes, k}
 	return pub, nil
 }

cng/ecdsa.go

@@ -90,7 +90,7 @@ func NewPublicKeyECDSA(curve string, X, Y BigInt) (*PublicKeyECDSA, error) {
 		return nil, err
 	}
 	k := &PublicKeyECDSA{hkey}
-	runtime.AddCleanup(k, destroyKey, hkey)
+	addCleanupKey(k, hkey)
 	return k, nil
 }
 
@@ -108,7 +108,7 @@ func NewPrivateKeyECDSA(curve string, X, Y, D BigInt) (*PrivateKeyECDSA, error)
 		return nil, err
 	}
 	k := &PrivateKeyECDSA{hkey}
-	runtime.AddCleanup(k, destroyKey, hkey)
+	addCleanupKey(k, hkey)
 	return k, nil
 }
 

cng/hash.go

@@ -20,9 +20,11 @@ import (
 // maxHashSize is the size of SHA512 and SHA3_512, the largest hashes we support.
 const maxHashSize = 64
 
-// destroyHash is a cleanup function for releasing bcrypt hash handles.
-func destroyHash(ctx bcrypt.HASH_HANDLE) {
-	bcrypt.DestroyHash(ctx)
+// addCleanupHash attaches a cleanup function to ptr that will destroy ctx.
+func addCleanupHash[T any](ptr *T, ctx bcrypt.HASH_HANDLE) {
+	runtime.AddCleanup(ptr, func(ctx bcrypt.HASH_HANDLE) {
+		bcrypt.DestroyHash(ctx)
+	}, ctx)
 }
 
 // SupportsHash returns true if a hash.Hash implementation is supported for h.
@@ -202,15 +204,15 @@ func (h *Hash) init() {
 	if err != nil {
 		panic(err)
 	}
-	runtime.AddCleanup(h, destroyHash, h.ctx)
+	addCleanupHash(h, h.ctx)
 }
 
 func (h *Hash) Clone() (HashCloner, error) {
 	defer runtime.KeepAlive(h)
 	h2 := &Hash{alg: h.alg, key: bytes.Clone(h.key)}
 	if h.ctx != 0 {
 		hashClone(h.ctx, &h2.ctx)
-		runtime.AddCleanup(h2, destroyHash, h2.ctx)
+		addCleanupHash(h2, h2.ctx)
 	}
 	return h2, nil
 }

cng/keys.go

@@ -8,14 +8,17 @@ package cng
 
 import (
 	"errors"
+	"runtime"
 	"unsafe"
 
 	"github.com/microsoft/go-crypto-winnative/internal/bcrypt"
 )
 
-// destroyKey is a cleanup function for releasing bcrypt key handles.
-func destroyKey(kh bcrypt.KEY_HANDLE) {
-	bcrypt.DestroyKey(kh)
+// addCleanupKey attaches a cleanup function to ptr that will destroy kh.
+func addCleanupKey[T any](ptr *T, kh bcrypt.KEY_HANDLE) {
+	runtime.AddCleanup(ptr, func(kh bcrypt.KEY_HANDLE) {
+		bcrypt.DestroyKey(kh)
+	}, kh)
 }
 
 const (

cng/rc4.go

@@ -25,7 +25,7 @@ func NewRC4Cipher(key []byte) (*RC4Cipher, error) {
 		return nil, err
 	}
 	c := &RC4Cipher{kh: kh}
-	runtime.AddCleanup(c, destroyKey, kh)
+	addCleanupKey(c, kh)
 	return c, nil
 }
 

cng/rsa.go

@@ -95,7 +95,7 @@ func NewPublicKeyRSA(N, E BigInt) (*PublicKeyRSA, error) {
 		return nil, err
 	}
 	k := &PublicKeyRSA{hkey, uint32(N.bitLen())}
-	runtime.AddCleanup(k, destroyKey, hkey)
+	addCleanupKey(k, hkey)
 	return k, nil
 }
 
@@ -117,7 +117,7 @@ func NewPrivateKeyRSA(N, E, D, P, Q, Dp, Dq, Qinv BigInt) (*PrivateKeyRSA, error
 		return nil, err
 	}
 	k := &PrivateKeyRSA{hkey, uint32(N.bitLen())}
-	runtime.AddCleanup(k, destroyKey, hkey)
+	addCleanupKey(k, hkey)
 	return k, nil
 }