add checksum for the python tchannel. It doesn't cover streaming message

yet

Author: Junchao Wu

requirements-test.txt

@@ -12,5 +12,8 @@ flake8==2.2.5
 # Optional dependency, but must be tested er'ry time
 tornado>=4.0,<5.0
 
+# checksum calculation
+pyfarmhash
+
 # Mocks
 doubles

tchannel/exceptions.py

@@ -29,6 +29,11 @@ class ReadException(TChannelException):
     pass
 
 
+class InvalidChecksumException(TChannelException):
+    """Represent invalid checksum type in the message"""
+    pass
+
+
 class TChannelApplicationException(TChannelException):
     """The remote application returned an exception.
 

tchannel/messages/call_request.py

@@ -42,7 +42,11 @@ def __init__(
         self.tracing = tracing or common.Tracing(0, 0, 0, 0)
         self.service = service or ''
         self.headers = dict(headers) if headers else {}
-        self.checksum = (common.ChecksumType.none, None)
+        if checksum is not None:
+            checksum = common.ChecksumType.standardize(checksum)
+        self.checksum = checksum or \
+            (common.ChecksumType.none, None)
+
         self.arg_1 = arg_1 or ''
         self.arg_2 = arg_2 or ''
         self.arg_3 = arg_3 or ''

tchannel/messages/call_response.py

@@ -38,7 +38,11 @@ def __init__(
         self.code = code
         self.tracing = tracing or common.Tracing(0, 0, 0, 0)
         self.headers = dict(headers) if headers else {}
-        self.checksum = (common.ChecksumType.none, None)
+        if checksum is not None:
+            checksum = common.ChecksumType.standardize(checksum)
+        self.checksum = checksum or \
+            (common.ChecksumType.none, None)
+
         self.arg_1 = arg_1 or ''
         self.arg_2 = arg_2 or ''
         self.arg_3 = arg_3 or ''

tchannel/messages/common.py

@@ -1,9 +1,13 @@
 from __future__ import absolute_import
 
+import zlib
+
 from collections import namedtuple
 from enum import IntEnum
 
 from .. import rw
+from ..exceptions import InvalidChecksumException
+from .types import Types
 
 PROTOCOL_VERSION = 0x02
 
@@ -24,6 +28,11 @@ class ChecksumType(IntEnum):
     crc32 = 0x01
     farm32 = 0x02
 
+    @staticmethod
+    def standardize(checksum):
+        return (ChecksumType(checksum[0]), checksum[1])
+
+
 checksum_rw = rw.switch(
     rw.number(1),   # csumtype:1
     {
@@ -32,3 +41,64 @@ class ChecksumType(IntEnum):
         ChecksumType.farm32: rw.number(4),  # csum:4
     }
 )
+
+
+CHECKSUM_MSG_TYPES = [Types.CALL_REQ,
+                      Types.CALL_REQ_CONTINUE,
+                      Types.CALL_RES,
+                      Types.CALL_RES_CONTINUE]
+
+
+def compute_checksum(checksum_type, args, csum=0):
+    if csum is None:
+        csum = 0
+
+    if checksum_type == ChecksumType.none:
+        return None
+    elif checksum_type == ChecksumType.crc32:
+        for arg in args:
+            csum = zlib.crc32(arg, csum) & 0xffffffff
+    # TODO figure out farm32 cross platform issue
+    elif checksum_type == ChecksumType.farm32:
+        raise NotImplementedError()
+    else:
+        raise InvalidChecksumException()
+
+    return csum
+
+
+def generate_checksum(message):
+    """Generate checksum for messages with
+        CALL_REQ, CALL_REQ_CONTINUE,
+        CALL_RES,CALL_RES_CONTINUE types
+
+    :param message: outgoing message
+    """
+    if message.message_type in CHECKSUM_MSG_TYPES:
+        csum = compute_checksum(
+            message.checksum[0],
+            [message.arg_1,
+             message.arg_2,
+             message.arg_3])
+
+        message.checksum = (message.checksum[0], csum)
+
+
+def verify_checksum(message):
+    """
+    :return return True if message checksum type is None
+    or checksum is correct
+    """
+    if message.message_type in CHECKSUM_MSG_TYPES:
+        csum = compute_checksum(
+            message.checksum[0],
+            [message.arg_1,
+             message.arg_2,
+             message.arg_3])
+
+        if csum == message.checksum[1]:
+            return True
+        else:
+            return False
+    else:
+        return True

tchannel/tornado/connection.py

@@ -9,6 +9,7 @@
 
 from tornado import gen
 from tornado import iostream
+from tornado.concurrent import Future
 
 from .. import frame
 from .. import messages
@@ -18,9 +19,8 @@
 from ..exceptions import ConnectionClosedException
 from ..messages import CallResponseMessage
 from ..messages.types import Types
-from ..messages.common import PROTOCOL_VERSION
+from ..messages.common import PROTOCOL_VERSION, generate_checksum
 
-from tornado.concurrent import Future
 
 log = logging.getLogger('tchannel')
 
@@ -114,6 +114,7 @@ def on_error(future):
 
     def frame_and_write(self, message, message_id=None):
         # TODO: track awaiting responses in here
+        generate_checksum(message)
         message_id = message_id or self.next_message_id()
 
         if message.message_type in (

tchannel/tornado/inbound_server.py

@@ -1,12 +1,14 @@
 #!/usr/bin/env python
 from __future__ import absolute_import
 
+import socket
 import sys
 import tornado.ioloop
 import tornado.tcpserver
-import socket
 
-from tchannel.tornado.connection import TornadoConnection
+from ..tornado.connection import TornadoConnection
+from ..exceptions import InvalidChecksumException
+from ..messages.common import verify_checksum
 
 
 class InboundServer(tornado.tcpserver.TCPServer):
@@ -48,4 +50,8 @@ def preprocess_request(self, context, conn):
         :param context: a context contains call request message
         :param conn: incoming tornado connection
         """
-        self.req_handler.handle_request(context, conn)
+        if verify_checksum(context.message):
+            self.req_handler.handle_request(context, conn)
+        else:
+            # TODO return Error message
+            raise InvalidChecksumException()

tests/test_checksum.py

@@ -0,0 +1,36 @@
+from __future__ import absolute_import
+import pytest
+
+from tchannel import frame
+from tchannel.frame_reader import FrameReader
+from tchannel import messages
+from tchannel.messages import CallRequestMessage, ChecksumType
+from tchannel.io import BytesIO
+from tchannel.messages.common import generate_checksum, verify_checksum
+
+
+@pytest.mark.parametrize('checksum_type, seed', [
+    (ChecksumType.none, 0),
+    (ChecksumType.crc32, 0x0812fa3f),
+])
+def test_checksum(checksum_type, seed):
+    message = CallRequestMessage()
+    message.checksum = (checksum_type, seed)
+    generate_checksum(message)
+    message_id = 32
+    payload = messages.RW[message.message_type].write(
+        message, BytesIO()
+    ).getvalue()
+
+    f = frame.Frame(
+        header=frame.FrameHeader(
+            message_type=message.message_type,
+            message_id=message_id,
+        ),
+        payload=payload
+    )
+
+    inframe = frame.frame_rw.write(f, BytesIO()).getvalue()
+    reader = FrameReader(BytesIO(inframe))
+    msg = reader.read().next().message
+    assert verify_checksum(msg)

tests/test_messages.py

@@ -109,7 +109,7 @@ def test_valid_ping_request():
         'tracing': messages.Tracing(0, 0, 0, 0),
         'service': 'kodenom',
         'headers': {},
-        'checksum': (messages.ChecksumType.none, 0),
+        'checksum': (messages.ChecksumType.none, None),
         'arg_1': None,
         'arg_2': None,
         'arg_3': None,