Fix memory leak when excpetions occur (#814)

shadowsocks/crypto/aead.py

@@ -206,10 +206,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)
@@ -265,6 +267,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:]
@@ -288,6 +291,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:]

shadowsocks/crypto/openssl.py

@@ -165,6 +165,7 @@ class OpenSSLAeadCrypto(OpenSSLCryptoBase, AeadCryptoBase):
             None
         )
         if not r:
+            self.clean()
             raise Exception('Set ivlen failed')
 
         self.cipher_ctx_init()
@@ -201,6 +202,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):
@@ -216,6 +218,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]
 
@@ -231,6 +234,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]
@@ -255,6 +259,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:])

shadowsocks/crypto/sodium.py

@@ -217,6 +217,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, crypto_path=None):