Skip to content

Commit b3072ea

Browse files
qmuntalqmuntal
authored
Merge pull request #142 from microsoft/dev/mclayton/update-addCleanup
Dev/mclayton/update add cleanup
2 parents 7f762f5 + 4e25a67 commit b3072ea

File tree

12 files changed

+140
-90
lines changed

12 files changed

+140
-90
lines changed

‎cng/aes.go‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,10 @@ func NewAESCipher(key []byte) (cipher.Block, error) {
3030
return nil, err
3131
}
3232
c := &aesCipher{kh: kh, key: bytes.Clone(key)}
33-
runtime.SetFinalizer(c, (*aesCipher).finalize)
33+
addCleanupKey(c, kh)
3434
return c, nil
3535
}
3636

37-
func (c *aesCipher) finalize() {
38-
bcrypt.DestroyKey(c.kh)
39-
}
40-
4137
func (c *aesCipher) BlockSize() int { return aesBlockSize }
4238

4339
// validateAndClipInputs checks that dst and src meet the [cipher.Block]
@@ -165,15 +161,11 @@ func newCBC(encrypt bool, alg string, key, iv []byte) *cbcCipher {
165161
panic(err)
166162
}
167163
x := &cbcCipher{kh: kh, encrypt: encrypt, blockSize: blockSize}
168-
runtime.SetFinalizer(x, (*cbcCipher).finalize)
164+
addCleanupKey(x, kh)
169165
x.SetIV(iv)
170166
return x
171167
}
172168

173-
func (x *cbcCipher) finalize() {
174-
bcrypt.DestroyKey(x.kh)
175-
}
176-
177169
func (x *cbcCipher) BlockSize() int { return x.blockSize }
178170

179171
func (x *cbcCipher) CryptBlocks(dst, src []byte) {
@@ -247,17 +239,13 @@ type aesGCM struct {
247239
maskInitialized bool
248240
}
249241

250-
func (g *aesGCM) finalize() {
251-
bcrypt.DestroyKey(g.kh)
252-
}
253-
254242
func newGCM(key []byte, tls cipherGCMTLS) (*aesGCM, error) {
255243
kh, err := newCipherHandle(bcrypt.AES_ALGORITHM, bcrypt.CHAIN_MODE_GCM, key)
256244
if err != nil {
257245
return nil, err
258246
}
259247
g := &aesGCM{kh: kh, tls: tls}
260-
runtime.SetFinalizer(g, (*aesGCM).finalize)
248+
addCleanupKey(g, kh)
261249
return g, nil
262250
}
263251

‎cng/chacha20poly1305.go‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -40,16 +40,10 @@ func NewChaCha20Poly1305(key []byte) (cipher.AEAD, error) {
4040
return nil, err
4141
}
4242
c := &chacha20poly1305{kh: kh}
43-
runtime.SetFinalizer(c, (*chacha20poly1305).finalize)
43+
addCleanupKey(c, kh)
4444
return c, nil
4545
}
4646

47-
func (c *chacha20poly1305) finalize() {
48-
if c.kh != 0 {
49-
bcrypt.DestroyKey(c.kh)
50-
}
51-
}
52-
5347
func (c *chacha20poly1305) NonceSize() int {
5448
return chacha20Poly1305NonceSize
5549
}

‎cng/des.go‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ func NewDESCipher(key []byte) (cipher.Block, error) {
2929
return nil, err
3030
}
3131
c := &desCipher{kh: kh, alg: bcrypt.DES_ALGORITHM, key: bytes.Clone(key)}
32-
runtime.SetFinalizer(c, (*desCipher).finalize)
32+
addCleanupKey(c, kh)
3333
return c, nil
3434
}
3535

@@ -39,14 +39,10 @@ func NewTripleDESCipher(key []byte) (cipher.Block, error) {
3939
return nil, err
4040
}
4141
c := &desCipher{kh: kh, alg: bcrypt.DES3_ALGORITHM, key: bytes.Clone(key)}
42-
runtime.SetFinalizer(c, (*desCipher).finalize)
42+
addCleanupKey(c, kh)
4343
return c, nil
4444
}
4545

46-
func (c *desCipher) finalize() {
47-
bcrypt.DestroyKey(c.kh)
48-
}
49-
5046
func (c *desCipher) BlockSize() int { return desBlockSize }
5147

5248
func (c *desCipher) Encrypt(dst, src []byte) {

‎cng/dsa.go‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -96,10 +96,6 @@ type PrivateKeyDSA struct {
9696
hkey bcrypt.KEY_HANDLE
9797
}
9898

99-
func (k *PrivateKeyDSA) finalize() {
100-
bcrypt.DestroyKey(k.hkey)
101-
}
102-
10399
// PublicKeyDSA represents a DSA public key.
104100
type PublicKeyDSA struct {
105101
DSAParameters
@@ -108,10 +104,6 @@ type PublicKeyDSA struct {
108104
hkey bcrypt.KEY_HANDLE
109105
}
110106

111-
func (k *PublicKeyDSA) finalize() {
112-
bcrypt.DestroyKey(k.hkey)
113-
}
114-
115107
// GenerateKeyDSA generates a new private DSA key using the given parameters.
116108
func GenerateKeyDSA(params DSAParameters) (x, y BigInt, err error) {
117109
h, err := loadDSA()
@@ -155,7 +147,7 @@ func NewPrivateKeyDSA(params DSAParameters, X, Y BigInt) (*PrivateKeyDSA, error)
155147
return nil, err
156148
}
157149
k := &PrivateKeyDSA{params, X, Y, hkey}
158-
runtime.SetFinalizer(k, (*PrivateKeyDSA).finalize)
150+
addCleanupKey(k, hkey)
159151
return k, nil
160152
}
161153

@@ -174,7 +166,7 @@ func NewPublicKeyDSA(params DSAParameters, Y BigInt) (*PublicKeyDSA, error) {
174166
return nil, err
175167
}
176168
k := &PublicKeyDSA{params, Y, hkey}
177-
runtime.SetFinalizer(k, (*PublicKeyDSA).finalize)
169+
addCleanupKey(k, hkey)
178170
return k, nil
179171
}
180172

‎cng/ecdh.go‎

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -58,21 +58,11 @@ type PublicKeyECDH struct {
5858
priv *PrivateKeyECDH
5959
}
6060

61-
func (k *PublicKeyECDH) finalize() {
62-
if k.priv == nil {
63-
bcrypt.DestroyKey(k.hkey)
64-
}
65-
}
66-
6761
type PrivateKeyECDH struct {
6862
hkey bcrypt.KEY_HANDLE
6963
isNIST bool
7064
}
7165

72-
func (k *PrivateKeyECDH) finalize() {
73-
bcrypt.DestroyKey(k.hkey)
74-
}
75-
7666
func ECDH(priv *PrivateKeyECDH, pub *PublicKeyECDH) ([]byte, error) {
7767
// First establish the shared secret.
7868
var secret bcrypt.SECRET_HANDLE
@@ -138,7 +128,7 @@ func GenerateKeyECDH(curve string) (*PrivateKeyECDH, []byte, error) {
138128
bytes = bytes[hdr.KeySize*2:]
139129

140130
k := &PrivateKeyECDH{hkey, isNIST(curve)}
141-
runtime.SetFinalizer(k, (*PrivateKeyECDH).finalize)
131+
addCleanupKey(k, hkey)
142132
return k, bytes, nil
143133
}
144134

@@ -173,7 +163,7 @@ func NewPublicKeyECDH(curve string, bytes []byte) (*PublicKeyECDH, error) {
173163
return nil, err
174164
}
175165
k := &PublicKeyECDH{hkey, append([]byte(nil), bytes...), nil}
176-
runtime.SetFinalizer(k, (*PublicKeyECDH).finalize)
166+
addCleanupKey(k, hkey)
177167
return k, nil
178168
}
179169

@@ -202,7 +192,7 @@ func NewPrivateKeyECDH(curve string, key []byte) (*PrivateKeyECDH, error) {
202192
return nil, err
203193
}
204194
k := &PrivateKeyECDH{hkey, nist}
205-
runtime.SetFinalizer(k, (*PrivateKeyECDH).finalize)
195+
addCleanupKey(k, hkey)
206196
return k, nil
207197
}
208198

@@ -220,8 +210,9 @@ func (k *PrivateKeyECDH) PublicKey() (*PublicKeyECDH, error) {
220210
// Only include X.
221211
bytes = data[:hdr.KeySize]
222212
}
213+
// No cleanup needed: pub.priv prevents k from being garbage collected,
214+
// so k's cleanup will destroy the shared hkey when both are unreachable.
223215
pub := &PublicKeyECDH{k.hkey, bytes, k}
224-
runtime.SetFinalizer(pub, (*PublicKeyECDH).finalize)
225216
return pub, nil
226217
}
227218

‎cng/ecdh_test.go‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,3 +158,107 @@ func hexDecode(t *testing.T, s string) []byte {
158158
}
159159
return b
160160
}
161+
162+
func BenchmarkGenerateKeyECDH(b *testing.B) {
163+
for _, curve := range []string{"P-256", "P-384", "P-521", "X25519"} {
164+
b.Run(curve, func(b *testing.B) {
165+
b.ReportAllocs()
166+
b.ResetTimer()
167+
for i := 0; i < b.N; i++ {
168+
_, _, err := cng.GenerateKeyECDH(curve)
169+
if err != nil {
170+
b.Fatal(err)
171+
}
172+
}
173+
})
174+
}
175+
}
176+
177+
func BenchmarkECDH(b *testing.B) {
178+
for _, curve := range []string{"P-256", "P-384", "P-521", "X25519"} {
179+
b.Run(curve, func(b *testing.B) {
180+
aliceKey, _, err := cng.GenerateKeyECDH(curve)
181+
if err != nil {
182+
b.Fatal(err)
183+
}
184+
bobKey, _, err := cng.GenerateKeyECDH(curve)
185+
if err != nil {
186+
b.Fatal(err)
187+
}
188+
bobPub, err := bobKey.PublicKey()
189+
if err != nil {
190+
b.Fatal(err)
191+
}
192+
b.ReportAllocs()
193+
b.ResetTimer()
194+
for i := 0; i < b.N; i++ {
195+
_, err := cng.ECDH(aliceKey, bobPub)
196+
if err != nil {
197+
b.Fatal(err)
198+
}
199+
}
200+
})
201+
}
202+
}
203+
204+
func BenchmarkNewPrivateKeyECDH(b *testing.B) {
205+
for _, curve := range []string{"P-256", "P-384", "P-521", "X25519"} {
206+
b.Run(curve, func(b *testing.B) {
207+
_, privBytes, err := cng.GenerateKeyECDH(curve)
208+
if err != nil {
209+
b.Fatal(err)
210+
}
211+
b.ReportAllocs()
212+
b.ResetTimer()
213+
for i := 0; i < b.N; i++ {
214+
_, err := cng.NewPrivateKeyECDH(curve, privBytes)
215+
if err != nil {
216+
b.Fatal(err)
217+
}
218+
}
219+
})
220+
}
221+
}
222+
223+
func BenchmarkNewPublicKeyECDH(b *testing.B) {
224+
for _, curve := range []string{"P-256", "P-384", "P-521", "X25519"} {
225+
b.Run(curve, func(b *testing.B) {
226+
key, _, err := cng.GenerateKeyECDH(curve)
227+
if err != nil {
228+
b.Fatal(err)
229+
}
230+
pub, err := key.PublicKey()
231+
if err != nil {
232+
b.Fatal(err)
233+
}
234+
pubBytes := pub.Bytes()
235+
b.ReportAllocs()
236+
b.ResetTimer()
237+
for i := 0; i < b.N; i++ {
238+
_, err := cng.NewPublicKeyECDH(curve, pubBytes)
239+
if err != nil {
240+
b.Fatal(err)
241+
}
242+
}
243+
})
244+
}
245+
}
246+
247+
func BenchmarkPublicKeyECDH(b *testing.B) {
248+
for _, curve := range []string{"P-256", "P-384", "P-521", "X25519"} {
249+
b.Run(curve, func(b *testing.B) {
250+
key, _, err := cng.GenerateKeyECDH(curve)
251+
if err != nil {
252+
b.Fatal(err)
253+
}
254+
b.ReportAllocs()
255+
b.ResetTimer()
256+
for i := 0; i < b.N; i++ {
257+
_, err := key.PublicKey()
258+
if err != nil {
259+
b.Fatal(err)
260+
}
261+
}
262+
})
263+
}
264+
}

‎cng/ecdsa.go‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -90,14 +90,10 @@ func NewPublicKeyECDSA(curve string, X, Y BigInt) (*PublicKeyECDSA, error) {
9090
return nil, err
9191
}
9292
k := &PublicKeyECDSA{hkey}
93-
runtime.SetFinalizer(k, (*PublicKeyECDSA).finalize)
93+
addCleanupKey(k, hkey)
9494
return k, nil
9595
}
9696

97-
func (k *PublicKeyECDSA) finalize() {
98-
bcrypt.DestroyKey(k.hkey)
99-
}
100-
10197
type PrivateKeyECDSA struct {
10298
hkey bcrypt.KEY_HANDLE
10399
}
@@ -112,14 +108,10 @@ func NewPrivateKeyECDSA(curve string, X, Y, D BigInt) (*PrivateKeyECDSA, error)
112108
return nil, err
113109
}
114110
k := &PrivateKeyECDSA{hkey}
115-
runtime.SetFinalizer(k, (*PrivateKeyECDSA).finalize)
111+
addCleanupKey(k, hkey)
116112
return k, nil
117113
}
118114

119-
func (k *PrivateKeyECDSA) finalize() {
120-
bcrypt.DestroyKey(k.hkey)
121-
}
122-
123115
// SignECDSA signs a hash (which should be the result of hashing a larger message),
124116
// using the private key, priv.
125117
//

‎cng/hash.go‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,13 @@ import (
2020
// maxHashSize is the size of SHA512 and SHA3_512, the largest hashes we support.
2121
const maxHashSize = 64
2222

23+
// addCleanupHash attaches a cleanup function to ptr that will destroy ctx.
24+
func addCleanupHash[T any](ptr *T, ctx bcrypt.HASH_HANDLE) {
25+
runtime.AddCleanup(ptr, func(ctx bcrypt.HASH_HANDLE) {
26+
bcrypt.DestroyHash(ctx)
27+
}, ctx)
28+
}
29+
2330
// SupportsHash returns true if a hash.Hash implementation is supported for h.
2431
func SupportsHash(h crypto.Hash) bool {
2532
switch h {
@@ -188,10 +195,6 @@ func newHash(id string) *Hash {
188195
return &Hash{alg: mustLoadHash(id, bcrypt.ALG_NONE_FLAG)}
189196
}
190197

191-
func (h *Hash) finalize() {
192-
bcrypt.DestroyHash(h.ctx)
193-
}
194-
195198
func (h *Hash) init() {
196199
defer runtime.KeepAlive(h)
197200
if h.ctx != 0 {
@@ -201,15 +204,15 @@ func (h *Hash) init() {
201204
if err != nil {
202205
panic(err)
203206
}
204-
runtime.SetFinalizer(h, (*Hash).finalize)
207+
addCleanupHash(h, h.ctx)
205208
}
206209

207210
func (h *Hash) Clone() (HashCloner, error) {
208211
defer runtime.KeepAlive(h)
209212
h2 := &Hash{alg: h.alg, key: bytes.Clone(h.key)}
210213
if h.ctx != 0 {
211214
hashClone(h.ctx, &h2.ctx)
212-
runtime.SetFinalizer(h2, (*Hash).finalize)
215+
addCleanupHash(h2, h2.ctx)
213216
}
214217
return h2, nil
215218
}

‎cng/keys.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,19 @@ package cng
88

99
import (
1010
"errors"
11+
"runtime"
1112
"unsafe"
1213

1314
"github.com/microsoft/go-crypto-winnative/internal/bcrypt"
1415
)
1516

17+
// addCleanupKey attaches a cleanup function to ptr that will destroy kh.
18+
func addCleanupKey[T any](ptr *T, kh bcrypt.KEY_HANDLE) {
19+
runtime.AddCleanup(ptr, func(kh bcrypt.KEY_HANDLE) {
20+
bcrypt.DestroyKey(kh)
21+
}, kh)
22+
}
23+
1624
const (
1725
sizeOfECCBlobHeader = uint32(unsafe.Sizeof(bcrypt.ECCKEY_BLOB{}))
1826
sizeOfRSABlobHeader = uint32(unsafe.Sizeof(bcrypt.RSAKEY_BLOB{}))

0 commit comments

Comments
 (0)