Prevent DOS attack on UDP echo servers by distinguishing between an echo request

net/tools/testserver/echo_message.py

Index: net/tools/testserver/echo_message.py
===================================================================
--- net/tools/testserver/echo_message.py	(revision 0)
+++ net/tools/testserver/echo_message.py	(revision 0)
@@ -0,0 +1,386 @@
+#!/usr/bin/python2.4
+# Copyright (c) 2011 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Provides utility functions for TCP/UDP echo servers and clients.
+
+This program has classes and functions to encode, decode, calculate checksum
+and verify the "echo request" and "echo response" messages. "echo request"
+message is an echo message sent from the client to the server. "echo response"
+message is a response from the server to the "echo request" message from the
+client.
+
+The format of "echo request" message is
+<version><checksum><payload_size><payload>. <version> is the version number
+of the "echo request" protocol. <checksum> is the checksum of the <payload>.
+<payload_size> is the size of the <payload>. <payload> is the echo message.
+
+The format of "echo response" message is
+<version><checksum><payload_size><key><encoded_payload>.<version>,
+<checksum> and <payload_size> are same as what is in the "echo request" message.
+<encoded_payload> is encoded version of the <payload>. <key> is a randomly
+generated key that is used to encode/decode the <payload>.
+"""
+
+__author__ = 'rtenneti@google.com (Raman Tenneti)'
+
+
+from itertools import cycle
+from itertools import izip
+import random
+
+
+class EchoHeader(object):
+  """Class to keep header info of the EchoRequest and EchoResponse messages.
+
+  This class knows how to parse the checksum, payload_size from the
+  "echo request" and "echo response" messages. It holds the checksum,
+  payload_size of the "echo request" and "echo response" messages.
+  """
+
+  # This specifies the version.
+  VERSION_STRING = '01'
+
+  # This specifies the starting position of the checksum and length of the
+  # checksum. Maximum value for the checksum is less than (2 ** 31 - 1).
+  CHECKSUM_START = 2
+  CHECKSUM_LENGTH = 10
+  CHECKSUM_FORMAT = '%010d'
+  CHECKSUM_END = CHECKSUM_START + CHECKSUM_LENGTH
+
+  # This specifies the starting position of the <payload_size> and length of the
+  # <payload_size>. Maximum number of bytes that can be sent in the <payload> is
+  # 9,999,999.
+  PAYLOAD_SIZE_START = CHECKSUM_END
+  PAYLOAD_SIZE_LENGTH = 7
+  PAYLOAD_SIZE_FORMAT = '%07d'
+  PAYLOAD_SIZE_END = PAYLOAD_SIZE_START + PAYLOAD_SIZE_LENGTH
+
+  def __init__(self, checksum=0, payload_size=0):
+    """Initializes the checksum and payload_size of self (EchoHeader).
+
+    Args:
+      checksum: (int)
+        The checksum of the payload.
+      payload_size: (int)
+        The size of the payload.
+    """
+    self.checksum = checksum
+    self.payload_size = payload_size
+
+  def ParseAndInitialize(self, echo_message):
+    """Parses the echo_message and initializes self with the parsed data.
+
+    This method extracts checksum, and payload_size from the echo_message
+    (echo_message could be either echo_request or echo_response messages) and
+    initializes self (EchoHeader) with checksum and payload_size.
+
+    Args:
+      echo_message: (string)
+        The string representation of EchoRequest or EchoResponse objects.
+    Raises:
+      ValueError: Invalid data
+    """
+    if not echo_message or len(echo_message) < EchoHeader.PAYLOAD_SIZE_END:
+      raise ValueError('Invalid data:%s' % echo_message)
+    self.checksum = int(echo_message[
+        EchoHeader.CHECKSUM_START:EchoHeader.CHECKSUM_END])
+    self.payload_size = int(echo_message[
+        EchoHeader.PAYLOAD_SIZE_START:EchoHeader.PAYLOAD_SIZE_END])
+
+  def InitializeFromPayload(self, payload):
+    """Initializes the EchoHeader object with the payload.
+
+    It calculates checksum for the payload and initializes self (EchoHeader)
+    with the calculated checksum and size of the payload.
+
+    This method is used by the client code during testing.
+
+    Args:
+      payload: (string)
+        The payload is the echo string (like 'hello').
+    Raises:
+      ValueError: Invalid data
+    """
+    if not payload:
+      raise ValueError('Invalid data:%s' % payload)
+    self.payload_size = len(payload)
+    self.checksum = Checksum(payload, self.payload_size)
+
+  def __str__(self):
+    """String representation of the self (EchoHeader).
+
+    Returns:
+      A string representation of self (EchoHeader).
+    """
+    checksum_string = EchoHeader.CHECKSUM_FORMAT % self.checksum
+    payload_size_string = EchoHeader.PAYLOAD_SIZE_FORMAT % self.payload_size
+    return EchoHeader.VERSION_STRING + checksum_string + payload_size_string
+
+
+class EchoRequest(EchoHeader):
+  """Class holds data specific to the "echo request" message.
+
+  This class holds the payload extracted from the "echo request" message.
+  """
+
+  # This specifies the starting position of the <payload>.
+  PAYLOAD_START = EchoHeader.PAYLOAD_SIZE_END
+
+  def __init__(self):
+    """Initializes EchoRequest object."""
+    EchoHeader.__init__(self)
+    self.payload = ''
+
+  def ParseAndInitialize(self, echo_request_data):
+    """Parses and Initializes the EchoRequest object from the echo_request_data.
+
+    This method extracts the header information (checksum and payload_size) and
+    payload from echo_request_data.
+
+    Args:
+      echo_request_data: (string)
+        The string representation of EchoRequest object.
+    Raises:
+      ValueError: Invalid data
+    """
+    EchoHeader.ParseAndInitialize(self, echo_request_data)
+    if len(echo_request_data) <= EchoRequest.PAYLOAD_START:
+      raise ValueError('Invalid data:%s' % echo_request_data)
+    self.payload = echo_request_data[EchoRequest.PAYLOAD_START:]
+
+  def InitializeFromPayload(self, payload):
+    """Initializes the EchoRequest object with payload.
+
+    It calculates checksum for the payload and initializes self (EchoRequest)
+    object.
+
+    Args:
+      payload: (string)
+        The payload string for which "echo request" needs to be constructed.
+    """
+    EchoHeader.InitializeFromPayload(self, payload)
+    self.payload = payload
+
+  def __str__(self):
+    """String representation of the self (EchoRequest).
+
+    Returns:
+      A string representation of self (EchoRequest).
+    """
+    return EchoHeader.__str__(self) + self.payload
+
+
+class EchoResponse(EchoHeader):
+  """Class holds data specific to the "echo response" message.
+
+  This class knows how to parse the "echo response" message. This class holds
+  key, encoded_payload and decoded_payload of the "echo response" message.
+  """
+
+  # This specifies the starting position of the |key_| and length of the |key_|.
+  # Minimum and maximum values for the |key_| are 100,000 and 999,999.
+  KEY_START = EchoHeader.PAYLOAD_SIZE_END
+  KEY_LENGTH = 6
+  KEY_FORMAT = '%06d'
+  KEY_END = KEY_START + KEY_LENGTH
+  KEY_MIN_VALUE = 0
+  KEY_MAX_VALUE = 999999
+
+  # This specifies the starting position of the <encoded_payload> and length
+  # of the <encoded_payload>.
+  ENCODED_PAYLOAD_START = KEY_END
+
+  def __init__(self, key='', encoded_payload='', decoded_payload=''):
+    """Initializes the EchoResponse object."""
+    EchoHeader.__init__(self)
+    self.key = key
+    self.encoded_payload = encoded_payload
+    self.decoded_payload = decoded_payload
+
+  def ParseAndInitialize(self, echo_response_data=None):
+    """Parses and Initializes the EchoResponse object from echo_response_data.
+
+    This method calls EchoHeader to extract header information from the
+    echo_response_data and it then extracts key and encoded_payload from the
+    echo_response_data. It holds the decoded payload of the encoded_payload.
+
+    Args:
+      echo_response_data: (string)
+        The string representation of EchoResponse object.
+    Raises:
+      ValueError: Invalid echo_request_data
+    """
+    EchoHeader.ParseAndInitialize(self, echo_response_data)
+    if len(echo_response_data) <= EchoResponse.ENCODED_PAYLOAD_START:
+      raise ValueError('Invalid echo_response_data:%s' % echo_response_data)
+    self.key = echo_response_data[EchoResponse.KEY_START:EchoResponse.KEY_END]
+    self.encoded_payload = echo_response_data[
+        EchoResponse.ENCODED_PAYLOAD_START:]
+    self.decoded_payload = Crypt(self.encoded_payload, self.key)
+
+  def InitializeFromEchoRequest(self, echo_request):
+    """Initializes EchoResponse with the data from the echo_request object.
+
+    It gets the checksum, payload_size and payload from the echo_request object
+    and then encodes the payload with a random key. It also saves the payload
+    as decoded_payload.
+
+    Args:
+      echo_request: (EchoRequest)
+        The EchoRequest object which has "echo request" message.
+    """
+    self.checksum = echo_request.checksum
+    self.payload_size = echo_request.payload_size
+    self.key = (EchoResponse.KEY_FORMAT %
+                random.randrange(EchoResponse.KEY_MIN_VALUE,
+                                 EchoResponse.KEY_MAX_VALUE))
+    self.encoded_payload = Crypt(echo_request.payload, self.key)
+    self.decoded_payload = echo_request.payload
+
+  def __str__(self):
+    """String representation of the self (EchoResponse).
+
+    Returns:
+      A string representation of self (EchoResponse).
+    """
+    return EchoHeader.__str__(self) + self.key + self.encoded_payload
+
+
+def Crypt(payload, key):
+  """Encodes/decodes the payload with the key and returns encoded payload.
+
+  This method loops through the payload and XORs each byte with the key.
+
+  Args:
+    payload: (string)
+      The string to be encoded/decoded.
+    key: (string)
+      The key used to encode/decode the payload.
+
+  Returns:
+    An encoded/decoded string.
+  """
+  return ''.join(chr(ord(x) ^ ord(y)) for (x, y) in izip(payload, cycle(key)))
+
+
+def Checksum(payload, payload_size):
+  """Calculates the checksum of the payload.
+
+  Args:
+    payload: (string)
+      The payload string for which checksum needs to be calculated.
+    payload_size: (int)
+      The number of bytes in the payload.
+
+  Returns:
+    The checksum of the payload.
+  """
+  checksum = 0
+  length = min(payload_size, len(payload))
+  for i in range (0, length):
+    checksum += ord(payload[i])
+  return checksum
+
+
+def GetEchoRequestData(payload):
+  """Constructs an "echo request" message from the payload.
+
+  It builds an EchoRequest object from the payload and then returns a string
+  representation of the EchoRequest object.
+
+  This is used by the TCP/UDP echo clients to build the "echo request" message.
+
+  Args:
+    payload: (string)
+      The payload string for which "echo request" needs to be constructed.
+
+  Returns:
+    A string representation of the EchoRequest object.
+  Raises:
+    ValueError: Invalid payload
+  """
+  try:
+    echo_request = EchoRequest()
+    echo_request.InitializeFromPayload(payload)
+    return str(echo_request)
+  except (IndexError, ValueError):
+    raise ValueError('Invalid payload:%s' % payload)
+
+
+def GetEchoResponseData(echo_request_data):
+  """Verifies the echo_request_data and returns "echo response" message.
+
+  It builds the EchoRequest object from the echo_request_data and then verifies
+  the checksum of the EchoRequest is same as the calculated checksum of the
+  payload. If the checksums don't match then it returns None. It checksums
+  match, it builds the echo_response object from echo_request object and returns
+  string representation of the EchoResponse object.
+
+  This is used by the TCP/UDP echo servers.
+
+  Args:
+    echo_request_data: (string)
+      The string that echo servers send to the clients.
+
+  Returns:
+    A string representation of the EchoResponse object. It returns None if the
+    echo_request_data is not valid.
+  Raises:
+    ValueError: Invalid echo_request_data
+  """
+  try:
+    if not echo_request_data:
+      raise ValueError('Invalid payload:%s' % echo_request_data)
+
+    echo_request = EchoRequest()
+    echo_request.ParseAndInitialize(echo_request_data)
+
+    if Checksum(echo_request.payload,
+                echo_request.payload_size) != echo_request.checksum:
+      return None
+
+    echo_response = EchoResponse()
+    echo_response.InitializeFromEchoRequest(echo_request)
+
+    return str(echo_response)
+  except (IndexError, ValueError):
+    raise ValueError('Invalid payload:%s' % echo_request_data)
+
+
+def DecodeAndVerify(echo_request_data, echo_response_data):
+  """Decodes and verifies the echo_response_data.
+
+  It builds EchoRequest and EchoResponse objects from the echo_request_data and
+  echo_response_data. It returns True if the EchoResponse's payload and
+  checksum match EchoRequest's.
+
+  This is used by the TCP/UDP echo clients for testing purposes.
+
+  Args:
+    echo_request_data: (string)
+      The request clients sent to echo servers.
+    echo_response_data: (string)
+      The response clients received from the echo servers.
+
+  Returns:
+    True if echo_request_data and echo_response_data match.
+  Raises:
+    ValueError: Invalid echo_request_data or Invalid echo_response
+  """
+
+  try:
+    echo_request = EchoRequest()
+    echo_request.ParseAndInitialize(echo_request_data)
+  except (IndexError, ValueError):
+    raise ValueError('Invalid echo_request:%s' % echo_request_data)
+
+  try:
+    echo_response = EchoResponse()
+    echo_response.ParseAndInitialize(echo_response_data)
+  except (IndexError, ValueError):
+    raise ValueError('Invalid echo_response:%s' % echo_response_data)
+
+  return (echo_request.checksum == echo_response.checksum and
+          echo_request.payload == echo_response.decoded_payload)