From 87633e07c722a802d3fd2555657b79076f02c844 Mon Sep 17 00:00:00 2001
From: Branden Clark <branden.clark@datadoghq.com>
Date: Mon, 2 Dec 2024 23:39:38 -0500
Subject: [PATCH] move msgo root logic to python

---
 omnibus/config/software/datadog-agent.rb |  6 +-----
 tasks/go.py                              |  8 ++++++++
 tasks/omnibus.py                         | 15 ++++++++++++++-
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/omnibus/config/software/datadog-agent.rb b/omnibus/config/software/datadog-agent.rb
index bb25815f886e25..1bacfa89cf6806 100644
--- a/omnibus/config/software/datadog-agent.rb
+++ b/omnibus/config/software/datadog-agent.rb
@@ -63,11 +63,7 @@
   # Use msgo toolchain when fips mode is enabled
   if fips_mode?
     if windows_target?
-      go_version = ENV['GO_VERSION'] || File.read("#{project_dir}/.go-version").strip
-      if go_version.empty?
-        raise "GO_VERSION is not set and .go-version file is not present"
-      end
-      msgoroot = "C:\\msgo\\#{go_version}\\go"
+      msgoroot = ENV['MSGO_ROOT']
       # ensure msgoroot\bin\go.exe file exists
       if !File.exist?("#{msgoroot}\\bin\\go.exe")
         raise "msgo go.exe not found at #{msgoroot}\\bin\\go.exe"
diff --git a/tasks/go.py b/tasks/go.py
index 0aefe9624115ca..325ecdaabeb1f1 100644
--- a/tasks/go.py
+++ b/tasks/go.py
@@ -715,3 +715,11 @@ def mod_diffs(_, targets):
                 print(f"  - Version {version} in:")
                 for path in paths:
                     print(f"    * {path}")
+
+
+def repo_go_version():
+    """
+    Returns the go version specified in the .go-version file
+    """
+    with open('.go-version') as f:
+        return f.read().strip()
diff --git a/tasks/omnibus.py b/tasks/omnibus.py
index 649992fdcfb44b..6e9799500de03e 100644
--- a/tasks/omnibus.py
+++ b/tasks/omnibus.py
@@ -6,7 +6,7 @@
 from invoke.exceptions import Exit, UnexpectedExit
 
 from tasks.flavor import AgentFlavor
-from tasks.go import deps
+from tasks.go import deps, repo_go_version
 from tasks.libs.common.omnibus import (
     install_dir_for_project,
     omnibus_compute_cache_key,
@@ -136,6 +136,19 @@ def get_omnibus_env(
 
     if fips_mode:
         env['FIPS_MODE'] = 'true'
+        if sys.platform == 'win32' and not os.environ.get('MSGO_ROOT'):
+            # Point omnibus at the msgo root
+            # TODO: idk how to do this in omnibus datadog-agent.rb
+            #       because `File.read` is executed when the script is loaded,
+            #       not when the `command`s are run and the source tree is not
+            #       available at that time.
+            #       Comments from the Linux FIPS PR discussed wanting to centralize
+            #       the msgo root logic, so this can be updated then.
+            go_version = repo_go_version()
+            env['MSGO_ROOT'] = f'C:\\msgo\\{go_version}\\go'
+            gobinpath = f"{env['MSGO_ROOT']}\\bin\\go.exe"
+            if not os.path.exists(gobinpath):
+                raise Exit(f"msgo go.exe not found at {gobinpath}")
 
     # We need to override the workers variable in omnibus build when running on Kubernetes runners,
     # otherwise, ohai detect the number of CPU on the host and run the make jobs with all the CPU.