Fix memory leak when excpetions occur

This commit is contained in:
Zou Yong 2017-03-29 01:42:24 +08:00
parent 1222fb19a6
commit 55837f7d0d
3 changed files with 12 additions and 0 deletions

View file

@ -202,10 +202,12 @@ class AeadCryptoBase(object):
# network byte order
ctext = [self.aead_encrypt(pack("!H", plen & AEAD_CHUNK_SIZE_MASK))]
if len(ctext[0]) != AEAD_CHUNK_SIZE_LEN + self._tlen:
self.clean()
raise Exception("size length invalid")
ctext.append(self.aead_encrypt(data))
if len(ctext[1]) != plen + self._tlen:
self.clean()
raise Exception("data length invalid")
return b''.join(ctext)
@ -261,6 +263,7 @@ class AeadCryptoBase(object):
plen = self.aead_decrypt(data[:hlen])
plen, = unpack("!H", plen)
if plen & AEAD_CHUNK_SIZE_MASK != plen or plen <= 0:
self.clean()
raise Exception('Invalid message length')
return plen, data[hlen:]
@ -284,6 +287,7 @@ class AeadCryptoBase(object):
plaintext = self.aead_decrypt(data[:plen + self._tlen])
if len(plaintext) != plen:
self.clean()
raise Exception("plaintext length invalid")
return plaintext, data[plen + self._tlen:]

View file

@ -163,6 +163,7 @@ class OpenSSLAeadCrypto(OpenSSLCryptoBase, AeadCryptoBase):
None
)
if not r:
self.clean()
raise Exception('Set ivlen failed')
self.cipher_ctx_init()
@ -199,6 +200,7 @@ class OpenSSLAeadCrypto(OpenSSLCryptoBase, AeadCryptoBase):
c_int(tag_len), c_char_p(tag)
)
if not r:
self.clean()
raise Exception('Set tag failed')
def get_tag(self):
@ -214,6 +216,7 @@ class OpenSSLAeadCrypto(OpenSSLCryptoBase, AeadCryptoBase):
c_int(tag_len), byref(tag_buf)
)
if not r:
self.clean()
raise Exception('Get tag failed')
return tag_buf.raw[:tag_len]
@ -229,6 +232,7 @@ class OpenSSLAeadCrypto(OpenSSLCryptoBase, AeadCryptoBase):
byref(buf), byref(cipher_out_len)
)
if not r:
self.clean()
# print(self._nonce.raw, r, cipher_out_len)
raise Exception('Finalize cipher failed')
return buf.raw[:cipher_out_len.value]
@ -253,6 +257,7 @@ class OpenSSLAeadCrypto(OpenSSLCryptoBase, AeadCryptoBase):
"""
clen = len(data)
if clen < self._tlen:
self.clean()
raise Exception('Data too short')
self.set_tag(data[clen - self._tlen:])

View file

@ -196,6 +196,9 @@ class SodiumCrypto(object):
# strip off the padding
return buf.raw[padding:padding + l]
def clean(self):
pass
class SodiumAeadCrypto(AeadCryptoBase):
def __init__(self, cipher_name, key, iv, op):