From 216d8a3be98e17a664afd6a40a98667c3ccc1f0c Mon Sep 17 00:00:00 2001
From: Gary Kotton <gkotton@vmware.com>
Date: Tue, 24 Oct 2017 23:41:56 -0700
Subject: [PATCH] Provide a callback to inject headers to the NSX

Enable application to bind to a callback that will inject
headers to the requests sent to the NSX. This can for example
be: 'X-NSX-EUSER' which will provide the NSX context of the user
on behalf of whom the opertaion is done.

A new method set_inject_headers_callback(in is added. This method
receives a callback that should return a dictionary of headers to
be added.

Change-Id: I90fa7ea8c7828bf97aec66321e4169588108760e
---
 vmware_nsxlib/tests/unit/v3/test_client.py | 27 ++++++++++++++++++++++
 vmware_nsxlib/v3/client.py                 |  4 ++++
 vmware_nsxlib/v3/utils.py                  |  6 +++++
 3 files changed, 37 insertions(+)

diff --git a/vmware_nsxlib/tests/unit/v3/test_client.py b/vmware_nsxlib/tests/unit/v3/test_client.py
index 85204445..9160ce7a 100644
--- a/vmware_nsxlib/tests/unit/v3/test_client.py
+++ b/vmware_nsxlib/tests/unit/v3/test_client.py
@@ -23,6 +23,7 @@ from vmware_nsxlib.tests.unit.v3 import mocks
 from vmware_nsxlib.tests.unit.v3 import nsxlib_testcase
 from vmware_nsxlib.v3 import client
 from vmware_nsxlib.v3 import exceptions as nsxlib_exc
+from vmware_nsxlib.v3 import utils
 
 
 LOG = log.getLogger(__name__)
@@ -275,6 +276,32 @@ class NsxV3RESTClientTestCase(nsxlib_testcase.NsxClientTestCase):
                 _verb_response_code, verb,
                 requests.codes.NOT_FOUND, 202)
 
+    def test_inject_headers_callback(self):
+
+        self.injected = None
+
+        def inject_header():
+            self.injected = True
+            return {}
+
+        utils.set_inject_headers_callback(inject_header)
+        api = self.new_mocked_client(
+            client.RESTClient,
+            url_prefix='/v1/api')
+
+        api.list()
+        injected_headers = {}
+        assert_call(
+            'get', api,
+            'https://1.2.3.4/v1/api',
+            headers=_headers(**injected_headers))
+
+        api = self.new_mocked_client(
+            client.RESTClient,
+            url_prefix='/v1/api')
+        utils.set_inject_headers_callback(None)
+        self.assertIsNotNone(self.injected)
+
 
 class NsxV3JSONClientTestCase(nsxlib_testcase.NsxClientTestCase):
 
diff --git a/vmware_nsxlib/v3/client.py b/vmware_nsxlib/v3/client.py
index 783721db..0855e2e3 100644
--- a/vmware_nsxlib/v3/client.py
+++ b/vmware_nsxlib/v3/client.py
@@ -178,6 +178,10 @@ class RESTClient(object):
                    silent=False):
         request_headers = headers.copy() if headers else {}
         request_headers.update(self._default_headers)
+        if utils.INJECT_HEADERS_CALLBACK:
+            inject_headers = utils.INJECT_HEADERS_CALLBACK()
+            request_headers.update(inject_headers)
+
         request_url = self._build_url(url)
 
         do_request = getattr(self._conn, method.lower())
diff --git a/vmware_nsxlib/v3/utils.py b/vmware_nsxlib/v3/utils.py
index b9123ef5..039b6411 100644
--- a/vmware_nsxlib/v3/utils.py
+++ b/vmware_nsxlib/v3/utils.py
@@ -32,6 +32,12 @@ MAX_RESOURCE_TYPE_LEN = 20
 MAX_TAG_LEN = 40
 DEFAULT_MAX_ATTEMPTS = 10
 DEFAULT_CACHE_AGE_SEC = 600
+INJECT_HEADERS_CALLBACK = None
+
+
+def set_inject_headers_callback(callback):
+    global INJECT_HEADERS_CALLBACK
+    INJECT_HEADERS_CALLBACK = callback
 
 
 def _validate_resource_type_length(resource_type):