From f0fad5289ca192a09772b5d0fef782681cf40920 Mon Sep 17 00:00:00 2001
From: Julien Duponchelle <julien@gns3.net>
Date: Fri, 19 Aug 2016 11:05:54 +0200
Subject: [PATCH] Correclty cleanup packet capture objects when closing server

Fix #592
---
 gns3server/controller/compute.py     | 15 ++++++++++++++-
 gns3server/controller/link.py        | 23 ++++++++++++-----------
 gns3server/controller/udp_link.py    |  2 +-
 gns3server/utils/asyncio/__init__.py |  1 +
 tests/utils/test_asyncio.py          |  2 ++
 5 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/gns3server/controller/compute.py b/gns3server/controller/compute.py
index f175ab35..145de540 100644
--- a/gns3server/controller/compute.py
+++ b/gns3server/controller/compute.py
@@ -295,11 +295,24 @@ class Compute:
         :returns: A file stream
         """
 
+        # Due to Python 3.4 limitation we can't use with and asyncio
+        # https://www.python.org/dev/peps/pep-0492/
+        # that why we wrap the answer
+        class StreamResponse:
+            def __init__(self, response):
+                self._response = response
+
+            def __enter__(self):
+                return self._response.content
+
+            def __exit__(self):
+                self._response.close()
+
         url = self._getUrl("/projects/{}/stream/{}".format(project.id, path))
         response = yield from self._session().request("GET", url, auth=self._auth)
         if response.status == 404:
             raise aiohttp.web.HTTPNotFound(text="{} not found on compute".format(path))
-        return response.content
+        return StreamResponse(response)
 
     @asyncio.coroutine
     def http_query(self, method, path, data=None, **kwargs):
diff --git a/gns3server/controller/link.py b/gns3server/controller/link.py
index 43ef7b79..845b0b34 100644
--- a/gns3server/controller/link.py
+++ b/gns3server/controller/link.py
@@ -116,17 +116,18 @@ class Link:
         Dump a pcap file on disk
         """
 
-        stream = yield from self.read_pcap_from_source()
-        with open(self.capture_file_path, "wb+") as f:
-            while self._capturing:
-                # We read 1 bytes by 1 otherwise the remaining data is not read if the traffic stops
-                data = yield from stream.read(1)
-                if data:
-                    f.write(data)
-                    # Flush to disk otherwise the live is not really live
-                    f.flush()
-                else:
-                    break
+        stream_content = yield from self.read_pcap_from_source()
+        with stream_content as stream:
+            with open(self.capture_file_path, "wb+") as f:
+                while self._capturing:
+                    # We read 1 bytes by 1 otherwise the remaining data is not read if the traffic stops
+                    data = yield from stream.read(1)
+                    if data:
+                        f.write(data)
+                        # Flush to disk otherwise the live is not really live
+                        f.flush()
+                    else:
+                        break
 
     @asyncio.coroutine
     def stop_capture(self):
diff --git a/gns3server/controller/udp_link.py b/gns3server/controller/udp_link.py
index 095707c1..96f2a6ce 100644
--- a/gns3server/controller/udp_link.py
+++ b/gns3server/controller/udp_link.py
@@ -120,7 +120,7 @@ class UDPLink(Link):
         :returns: Node where the capture should run
         """
 
-        # use the local node first to save bandwidth
+        # use the local node first to save bandwidth
         for node in self._nodes:
             if node["node"].compute.id == "local" and node["node"].node_type not in [""]:  # FIXME
                 return node
diff --git a/gns3server/utils/asyncio/__init__.py b/gns3server/utils/asyncio/__init__.py
index 1ddfcc76..e0577af8 100644
--- a/gns3server/utils/asyncio/__init__.py
+++ b/gns3server/utils/asyncio/__init__.py
@@ -127,3 +127,4 @@ def wait_for_named_pipe_creation(pipe_path, timeout=60):
         else:
             return
     raise asyncio.TimeoutError()
+
diff --git a/tests/utils/test_asyncio.py b/tests/utils/test_asyncio.py
index b4fc7ae6..3c5d617f 100644
--- a/tests/utils/test_asyncio.py
+++ b/tests/utils/test_asyncio.py
@@ -67,3 +67,5 @@ def test_wait_for_process_termination(loop):
     exec = wait_for_process_termination(process, timeout=0.5)
     with pytest.raises(asyncio.TimeoutError):
         loop.run_until_complete(asyncio.async(exec))
+
+