Add CNS function to clean up port matching requester IP address.

BUG=141537
TEST=unit tests

Review URL: https://chromiumcodereview.appspot.com/10824224

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@150705 0039d316-1c4b-4281-b951-d872f2087c98

2 files changed, 48 insertions, 2 deletions

diff --git a/media/tools/constrained_network_server/cns.py b/media/tools/constrained_network_server/cns.py
index 4fe31ee..9b0e96b 100755
--- a/media/tools/constrained_network_server/cns.py
+++ b/media/tools/constrained_network_server/cns.py
@@ -129,7 +129,7 @@ class PortAllocator(object):
       cherrypy.log('Error: %s\nOutput: %s' % (e.msg, e.error))
       return False
 
-  def Cleanup(self, all_ports):
+  def Cleanup(self, all_ports, request_ip=None):
     """Cleans up expired ports, or if all_ports=True, all allocated ports.
 
     By default, ports which haven't been used for self._expiry_time_secs are
@@ -137,13 +137,15 @@ class PortAllocator(object):
 
     Args:
       all_ports: Should all ports be torn down regardless of expiration?
+      request_ip: Tear ports matching the IP address regarless of expiration.
     """
     with self._port_lock:
       now = time.time()
       # Use .items() instead of .iteritems() so we can delete keys w/o error.
       for port, status in self._ports.items():
         expired = now - status['last_update'] > self._expiry_time_secs
-        if all_ports or expired:
+        matching_ip = request_ip and status['key'][0].startswith(request_ip)
+        if all_ports or expired or matching_ip:
           cherrypy.log('Cleaning up port %d' % port)
           self._DeletePort(port)
           del self._ports[port]
@@ -174,6 +176,19 @@ class ConstrainedNetworkServer(object):
     self._port_allocator = port_allocator
 
   @cherrypy.expose
+  def Cleanup(self):
+    """Cleans up all the ports allocated using the request IP address.
+
+    When requesting a constrained port, the cherrypy.request.remote.ip is used
+    as a key for that port (in addition to other request parameters).  Such
+    ports created for the same IP address are removed.
+    """
+    cherrypy.log('Cleaning up ports allocated by %s.' %
+                 cherrypy.request.remote.ip)
+    self._port_allocator.Cleanup(all_ports=False,
+                                 request_ip=cherrypy.request.remote.ip)
+
+  @cherrypy.expose
   def ServeConstrained(self, f=None, bandwidth=None, latency=None, loss=None,
                        new_port=False, no_cache=False, **kwargs):
     """Serves the requested file with the requested constraints.
diff --git a/media/tools/constrained_network_server/cns_test.py b/media/tools/constrained_network_server/cns_test.py
index 0a6e995..a3f1abf 100755
--- a/media/tools/constrained_network_server/cns_test.py
+++ b/media/tools/constrained_network_server/cns_test.py
@@ -115,6 +115,37 @@ class PortAllocatorTest(unittest.TestCase):
     self.assertEquals(set(self._pa._ports.keys()), set([
         cns._DEFAULT_CNS_PORT_RANGE[0], cns._DEFAULT_CNS_PORT_RANGE[0] + 1]))
 
+  def testPortAllocatorCleanMatchingIP(self):
+    # Setup PortAllocator w/o port expiration.
+    self._pa = cns.PortAllocator(cns._DEFAULT_CNS_PORT_RANGE, 0)
+
+    # Ensure Get() succeeds and returns the correct port.
+    self.assertEquals(self._pa.Get('ip1', t=1), cns._DEFAULT_CNS_PORT_RANGE[0])
+    self.assertEquals(self._pa.Get('ip1', t=2),
+                      cns._DEFAULT_CNS_PORT_RANGE[0] + 1)
+    self.assertEquals(self._pa.Get('ip1', t=3),
+                      cns._DEFAULT_CNS_PORT_RANGE[0] + 2)
+    self.assertEquals(self._pa.Get('ip2', t=1),
+                      cns._DEFAULT_CNS_PORT_RANGE[0] + 3)
+
+    self._pa.Cleanup(all_ports=False, request_ip='ip1')
+
+    self.assertEquals(self._pa._ports.keys(),
+                      [cns._DEFAULT_CNS_PORT_RANGE[0] + 3])
+    self.assertEquals(self._pa.Get('ip2'), cns._DEFAULT_CNS_PORT_RANGE[0])
+    self.assertEquals(self._pa.Get('ip1'), cns._DEFAULT_CNS_PORT_RANGE[0] + 1)
+
+    self._pa.Cleanup(all_ports=False, request_ip='ip2')
+    self.assertEquals(self._pa._ports.keys(),
+                      [cns._DEFAULT_CNS_PORT_RANGE[0] + 1])
+
+    self._pa.Cleanup(all_ports=False, request_ip='abc')
+    self.assertEquals(self._pa._ports.keys(),
+                      [cns._DEFAULT_CNS_PORT_RANGE[0] + 1])
+
+    self._pa.Cleanup(all_ports=False, request_ip='ip1')
+    self.assertEquals(self._pa._ports.keys(), [])
+
 
 class ConstrainedNetworkServerTest(unittest.TestCase):
   """End to end tests for ConstrainedNetworkServer system."""