From e6cc97fb674a3d86e56bada35545d297c1dfb0c1 Mon Sep 17 00:00:00 2001
From: Joshua Harlow <harlowja@yahoo-inc.com>
Date: Wed, 11 Jun 2014 11:22:39 -0700
Subject: [PATCH] Add sanity tests for profiler and hmac usage

Create a few sanity test cases that validate that
the hmac and profiler integration work as expected
when invalid, fake and valid usage occurs.

Change-Id: Ifb4c3763b2fb2caaa3c10f722538add785ffc086
---
 osprofiler/web.py |  4 ++--
 tests/test_web.py | 49 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/osprofiler/web.py b/osprofiler/web.py
index dc056bc..c2a9861 100644
--- a/osprofiler/web.py
+++ b/osprofiler/web.py
@@ -26,8 +26,8 @@ def add_trace_id_header(headers):
     p = profiler.get_profiler()
     if p:
         idents = {"base_id": p.get_base_id(), "parent_id": p.get_id()}
-        raw_content = json.dumps(idents)
-        headers["X-Trace-Info"] = utils.binary_encode(raw_content)
+        raw_content = utils.binary_encode(json.dumps(idents))
+        headers["X-Trace-Info"] = raw_content
         if p.hmac_key:
             headers["X-Trace-HMAC"] = generate_hmac(raw_content, p.hmac_key)
 
diff --git a/tests/test_web.py b/tests/test_web.py
index c7362c3..2fc522b 100644
--- a/tests/test_web.py
+++ b/tests/test_web.py
@@ -56,6 +56,55 @@ class WebMiddlewareTestCase(test.TestCase):
         web.add_trace_id_header(headers)
         self.assertEqual(old_headers, headers)
 
+    def test_hmac_generation(self):
+        profiler.init(base_id="b", parent_id="a", hmac_key="secret_password")
+        headers = {
+            'Content-Type': 'text/javascript',
+        }
+        web.add_trace_id_header(headers)
+        self.assertIn('X-Trace-HMAC', headers)
+        self.assertTrue(len(headers['X-Trace-HMAC']) > 0)
+
+    def test_hmac_no_generation(self):
+        profiler.init(base_id="b", parent_id="a")
+        headers = {
+            'Content-Type': 'text/javascript',
+        }
+        web.add_trace_id_header(headers)
+        self.assertNotIn('X-Trace-HMAC', headers)
+        self.assertIn('X-Trace-Info', headers)
+        self.assertEqual(2, len(headers))
+
+    def test_hmac_validation(self):
+        profiler.init(base_id="b", parent_id="a", hmac_key="secret_password")
+        headers = {
+            'Content-Type': 'text/javascript',
+        }
+        web.add_trace_id_header(headers)
+        content = headers.get("X-Trace-Info")
+        web.validate_hmac(content, headers['X-Trace-HMAC'], "secret_password")
+
+    def test_invalid_hmac(self):
+        profiler.init(base_id="b", parent_id="a", hmac_key="secret_password")
+        headers = {
+            'Content-Type': 'text/javascript',
+        }
+        web.add_trace_id_header(headers)
+        content = headers.get("X-Trace-Info")
+        content += "_changed"
+        self.assertRaises(IOError, web.validate_hmac, content,
+                          headers['X-Trace-HMAC'], "secret_password")
+
+    def test_hmac_faked(self):
+        headers = {
+            'Content-Type': 'text/javascript',
+            'X-Trace-HMAC': 'fake',
+            'X-Trace-Info': '{}',
+        }
+        content = headers.get("X-Trace-Info")
+        self.assertRaises(IOError, web.validate_hmac, content,
+                          headers['X-Trace-HMAC'], 'secret_password')
+
     def test_wsgi_middleware_no_trace(self):
         request = mock.MagicMock()
         request.get_response.return_value = "yeah!"