From 552904aedbca1c86fe7a008d566af640c9948d7f Mon Sep 17 00:00:00 2001
From: WooParadog <haochuan.guo@ele.me>
Date: Mon, 30 Oct 2017 12:02:26 +0800
Subject: [PATCH] Add on_tdecode_exception hook to catch decode/encode
 exceptions

---
 gunicorn_thrift/config.py                 | 17 +++++++++++++++++
 gunicorn_thrift/thriftpy_gevent_worker.py |  4 ++++
 gunicorn_thrift/thriftpy_sync_worker.py   |  4 ++++
 3 files changed, 25 insertions(+)

diff --git a/gunicorn_thrift/config.py b/gunicorn_thrift/config.py
index 889a04c..325d0d4 100644
--- a/gunicorn_thrift/config.py
+++ b/gunicorn_thrift/config.py
@@ -92,6 +92,23 @@ def on_connected(worker, addr):
         """
 
 
+class TDecodeExceptionRaised(Setting):
+    name = "on_tdecode_exception"
+    section = "Server Hooks"
+    validator = validate_callable(1)
+    type = six.callable
+
+    def on_tdecode_exception(err):
+        pass
+
+    default = staticmethod(on_tdecode_exception)
+    desc = """\
+        Called if tdecode exception is raised
+
+        The callable needs to accept one variable for the exception raised.
+        """
+
+
 class ClientConnectClosed(Setting):
     name = "post_connect_closed"
     section = "Server Hooks"
diff --git a/gunicorn_thrift/thriftpy_gevent_worker.py b/gunicorn_thrift/thriftpy_gevent_worker.py
index 474ea70..6f02214 100644
--- a/gunicorn_thrift/thriftpy_gevent_worker.py
+++ b/gunicorn_thrift/thriftpy_gevent_worker.py
@@ -28,6 +28,7 @@
 from thriftpy.transport import TTransportException
 from thriftpy.protocol.exc import TProtocolException
 from thriftpy.protocol.cybin import ProtocolError
+from thriftpy.thrift import TDecodeException
 
 from gunicorn.errors import AppImportError
 from gunicorn.workers.ggevent import GeventWorker
@@ -175,6 +176,9 @@ def handle(self, listener, client, addr):
             self.log.warning(
                 "Protocol error, is client(%s) correct? %s", addr, err
                 )
+        except TDecodeException as err:
+            self.log.exception('%r: %r', addr, err)
+            self.cfg.on_tdecode_exception(err)
         except socket.timeout:
             self.log.warning('Client timeout: %r', addr)
         except socket.error as e:
diff --git a/gunicorn_thrift/thriftpy_sync_worker.py b/gunicorn_thrift/thriftpy_sync_worker.py
index 70ee683..9baf351 100644
--- a/gunicorn_thrift/thriftpy_sync_worker.py
+++ b/gunicorn_thrift/thriftpy_sync_worker.py
@@ -12,6 +12,7 @@
 
 from thriftpy.transport import TSocket
 from thriftpy.transport import TTransportException
+from thriftpy.thrift import TDecodeException
 
 from gunicorn.errors import AppImportError
 from gunicorn.workers.sync import SyncWorker
@@ -68,6 +69,9 @@ def handle(self, listener, client, addr):
                     self.notify()
             except TTransportException:
                 pass
+        except TDecodeException as err:
+            self.log.exception('%r: %r', addr, err)
+            self.cfg.on_tdecode_exception(err)
         except socket.timeout:
             self.log.warning('Client timeout: %r', addr)
         except socket.error as e: