From 55837f7d0d1c5df1b6b1aef441584566d30faf34 Mon Sep 17 00:00:00 2001
From: Zou Yong <zouy02@gmail.com>
Date: Wed, 29 Mar 2017 01:42:24 +0800
Subject: [PATCH] Fix memory leak when excpetions occur

---
 shadowsocks/crypto/aead.py    | 4 ++++
 shadowsocks/crypto/openssl.py | 5 +++++
 shadowsocks/crypto/sodium.py  | 3 +++
 3 files changed, 12 insertions(+)

diff --git a/shadowsocks/crypto/aead.py b/shadowsocks/crypto/aead.py
index 23c0b36..e3c2bec 100644
--- a/shadowsocks/crypto/aead.py
+++ b/shadowsocks/crypto/aead.py
@@ -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:]
diff --git a/shadowsocks/crypto/openssl.py b/shadowsocks/crypto/openssl.py
index 6d4d154..9a48189 100644
--- a/shadowsocks/crypto/openssl.py
+++ b/shadowsocks/crypto/openssl.py
@@ -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:])
diff --git a/shadowsocks/crypto/sodium.py b/shadowsocks/crypto/sodium.py
index 69196a3..ea9c600 100644
--- a/shadowsocks/crypto/sodium.py
+++ b/shadowsocks/crypto/sodium.py
@@ -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):