Repository: Azure/WALinuxAgent
Pull request #2078 (2020-11-10 21:16:42)

Add wire server address as an environment variable
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description
Partners such as VM extension owners need a simple and reliable way to obtain the WireServer IP address so they can communicate with WireProtocol.

use the “AZURE_GUEST_AGENT_WIRE_PROTOCOL_ADDRESS” environment variable

TFS Task: https://msazure.visualstudio.com/One/_workitems/edit/8378126

Issue # <!-- if any -->
<!--
Please add an informative description that covers that changes made by the pull request. 
This checklist is used to make sure that common issues in a pull request are addressed.
This will expedite the process of getting your pull request merged and avoid extra work on your part to fix issues discovered during the review process.
-->

---

### PR information
- [ ] The title of the PR is clear and informative.
- [ ] There are a small number of commits, each of which has an informative message. This means that previously merged commits do not appear in the history of the PR. For information on cleaning up the commits in your pull request, [see this page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).
- [ ] Except for special cases involving multiple contributors, the PR is started from a fork of the main repository, not a branch.
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made and Travis.CI is passing.

### Quality of Code and Contribution Guidelines
- [ ] I have read the [contribution guidelines](https://github.com/Azure/WALinuxAgent/blob/master/.github/CONTRIBUTING.md).


Patch:
diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py
index 3891ba5e..bfac3bfa 100644
--- a/azurelinuxagent/ga/exthandlers.py
+++ b/azurelinuxagent/ga/exthandlers.py
@@ -98,14 +98,16 @@ class ValidHandlerStatus(object):  # pylint: disable=R0903
 _EXTENSION_TERMINAL_STATUSES = [ValidHandlerStatus.error, ValidHandlerStatus.success]
 
 
+# too-few-public-methods<R0903> Disabled: This class is used as an Enum
 class ExtCommandEnvVariable(object):  # pylint: disable=R0903
     Prefix = "AZURE_GUEST_AGENT"
-    DisableReturnCode = "%s_DISABLE_CMD_EXIT_CODE" % Prefix
-    UninstallReturnCode = "%s_UNINSTALL_CMD_EXIT_CODE" % Prefix
-    ExtensionPath = "%s_EXTENSION_PATH" % Prefix
-    ExtensionVersion = "%s_EXTENSION_VERSION" % Prefix
+    DisableReturnCode = "{0}_DISABLE_CMD_EXIT_CODE".format(Prefix)
+    UninstallReturnCode = "{0}_UNINSTALL_CMD_EXIT_CODE".format(Prefix)
+    ExtensionPath = "{0}_EXTENSION_PATH".format(Prefix)
+    ExtensionVersion = "{0}_EXTENSION_VERSION".format(Prefix)
     ExtensionSeqNumber = "ConfigSequenceNumber"  # At par with Windows Guest Agent
-    UpdatingFromVersion = "%s_UPDATING_FROM_VERSION" % Prefix
+    UpdatingFromVersion = "{0}_UPDATING_FROM_VERSION".format(Prefix)
+    WireProtocolAddress = "{0}_WIRE_PROTOCOL_ADDRESS".format(Prefix)
 
 
 def get_traceback(e):  # pylint: disable=R1710,C0103
@@ -1331,9 +1333,12 @@ class ExtHandlerInstance(object):  # pylint: disable=R0904
                     env = {}
                 env.update(os.environ)
                 # Always add Extension Path and version to the current launch_command (Ask from publishers)
-                env.update({ExtCommandEnvVariable.ExtensionPath: base_dir,
-                            ExtCommandEnvVariable.ExtensionVersion: str(self.ext_handler.properties.version),
-                            ExtCommandEnvVariable.ExtensionSeqNumber: str(self.get_seq_no())})
+                env.update({
+                    ExtCommandEnvVariable.ExtensionPath: base_dir,
+                    ExtCommandEnvVariable.ExtensionVersion: str(self.ext_handler.properties.version),
+                    ExtCommandEnvVariable.ExtensionSeqNumber: str(self.get_seq_no()),
+                    ExtCommandEnvVariable.WireProtocolAddress: self.protocol.get_endpoint()
+                })
 
                 try:
                     # Some extensions erroneously begin cmd with a slash; don't interpret those


Test patch:
diff --git a/tests/ga/test_exthandlers.py b/tests/ga/test_exthandlers.py
index fd79d9ba..58305539 100644
--- a/tests/ga/test_exthandlers.py
+++ b/tests/ga/test_exthandlers.py
@@ -19,15 +19,18 @@ import json
 import os
 import subprocess
 import time
+import uuid
 
 from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
 from azurelinuxagent.common.event import AGENT_EVENT_FILE_EXTENSION, WALAEventOperation
 from azurelinuxagent.common.exception import ExtensionError, ExtensionErrorCodes
 from azurelinuxagent.common.protocol.restapi import ExtensionStatus, Extension, ExtHandler, ExtHandlerProperties
 from azurelinuxagent.common.protocol.util import ProtocolUtil
+from azurelinuxagent.common.protocol.wire import WireProtocol
 from azurelinuxagent.common.utils.extensionprocessutil import TELEMETRY_MESSAGE_MAX_LEN, format_stdout_stderr, \
     read_output
-from azurelinuxagent.ga.exthandlers import parse_ext_status, ExtHandlerInstance, ExtCommandEnvVariable, ExtensionStatusError
+from azurelinuxagent.ga.exthandlers import parse_ext_status, ExtHandlerInstance, ExtCommandEnvVariable, \
+    ExtensionStatusError
 from tests.protocol import mockwiredata
 from tests.protocol.mocks import mock_wire_protocol
 from tests.tools import AgentTestCase, patch, mock_sleep, clear_singleton_instances
@@ -299,7 +302,7 @@ class LaunchCommandTestCase(AgentTestCase):
         ext_handler_properties.version = "1.2.3"
         self.ext_handler = ExtHandler(name='foo')
         self.ext_handler.properties = ext_handler_properties
-        self.ext_handler_instance = ExtHandlerInstance(ext_handler=self.ext_handler, protocol=None)
+        self.ext_handler_instance = ExtHandlerInstance(ext_handler=self.ext_handler, protocol=WireProtocol("1.2.3.4"))
 
         self.mock_get_base_dir = patch("azurelinuxagent.ga.exthandlers.ExtHandlerInstance.get_base_dir", lambda *_: self.tmp_dir)
         self.mock_get_base_dir.start()
@@ -669,9 +672,13 @@ sys.stderr.write("STDERR")
 
     def test_it_should_contain_all_helper_environment_variables(self):
 
-        helper_env_vars = {ExtCommandEnvVariable.ExtensionSeqNumber: self.ext_handler_instance.get_seq_no(),
+        wire_ip = str(uuid.uuid4())
+        ext_handler_instance = ExtHandlerInstance(ext_handler=self.ext_handler, protocol=WireProtocol(wire_ip))
+
+        helper_env_vars = {ExtCommandEnvVariable.ExtensionSeqNumber: ext_handler_instance.get_seq_no(),
                            ExtCommandEnvVariable.ExtensionPath: self.tmp_dir,
-                           ExtCommandEnvVariable.ExtensionVersion: self.ext_handler_instance.ext_handler.properties.version}
+                           ExtCommandEnvVariable.ExtensionVersion: ext_handler_instance.ext_handler.properties.version,
+                           ExtCommandEnvVariable.WireProtocolAddress: wire_ip}
 
         command = """
             printenv | grep -E '(%s)'
@@ -680,7 +687,7 @@ sys.stderr.write("STDERR")
         test_file = self.create_script('printHelperEnvironments.sh', command)
 
         with patch("subprocess.Popen", wraps=subprocess.Popen) as patch_popen:
-            output = self.ext_handler_instance.launch_command(test_file)
+            output = ext_handler_instance.launch_command(test_file)
 
             args, kwagrs = patch_popen.call_args  # pylint: disable=unused-variable
             without_os_env = dict((k, v) for (k, v) in kwagrs['env'].items() if k not in os.environ)