From 0f8b427bf4e5e7db348f5efc0b319889beeb3aca Mon Sep 17 00:00:00 2001
From: Chuck Walbourn <walbourn@users.noreply.github.com>
Date: Wed, 6 Dec 2023 19:17:02 -0800
Subject: [PATCH] ADO OneFuzz pipeline updated with setup script (#420)

---
 build/DirectXTex-OneFuzz.yml |   8 ++-
 build/OneFuzzConfig.json     |   5 +-
 build/onefuzz-setup.ps1      | 104 +++++++++++++++++++++++++++++++++++
 3 files changed, 114 insertions(+), 3 deletions(-)
 create mode 100644 build/onefuzz-setup.ps1

diff --git a/build/DirectXTex-OneFuzz.yml b/build/DirectXTex-OneFuzz.yml
index 32dfce9..df0e1e2 100644
--- a/build/DirectXTex-OneFuzz.yml
+++ b/build/DirectXTex-OneFuzz.yml
@@ -118,6 +118,13 @@ jobs:
             Invoke-WebRequest $url -o $target
         }
 
+  - task: PowerShell@2
+    displayName: Copy OneFuzz setup script
+    inputs:
+      targetType: 'inline'
+      script: |
+        Copy-Item -Path .\build\onefuzz-setup.ps1 -Destination .drop/setup.ps1
+
   - task: MSBuild@1
     displayName: 'Copy ASan binaries'
     inputs:
@@ -141,4 +148,3 @@ jobs:
       onefuzzDropPAT: $(ONEFUZZ_PAT)
       SYSTEM_ACCESSTOKEN: $(System.AccessToken)
       onefuzzBugFilingPAT: $(BUGFILING_PAT)
-  
\ No newline at end of file
diff --git a/build/OneFuzzConfig.json b/build/OneFuzzConfig.json
index b63c367..30f7272 100644
--- a/build/OneFuzzConfig.json
+++ b/build/OneFuzzConfig.json
@@ -14,7 +14,7 @@
           "ProjectName": "Direct3D",
           "TargetName": "DirectXTex",
           "TargetOptions": [
-             "-rss_limit_mb=4096"
+             " -rss_limit_mb=4096"
           ]
         }
       ],
@@ -22,7 +22,8 @@
         "fuzzloaders.exe",
         "fuzzloaders.pdb",
         "clang_rt.asan_dynamic-x86_64.dll",
-        "msdia140.dll"
+        "msdia140.dll",
+        "setup.ps1"
       ],
       "AdoTemplate": {
         "Org": "microsoft",
diff --git a/build/onefuzz-setup.ps1 b/build/onefuzz-setup.ps1
new file mode 100644
index 0000000..55fc44f
--- /dev/null
+++ b/build/onefuzz-setup.ps1
@@ -0,0 +1,104 @@
+function Execute-Setup {
+    # Temporary work-around while OneFuzz does not run script from setup dir
+    Set-Location -Path $PSScriptRoot
+
+    Write-Log "Executing custom setup script in $(pwd)"
+
+    # Exclude any uploaded DLL from known DLLs
+    gci -filter '*.dll' | Exclude-Library
+
+    # Update environment values for ASAN_OPTIONS.
+    # These environment variables affect how the ASan runtime operates. 
+    # Use the 'allocator_may_return_null=1' environment variable if you have compiled
+    # your fuzzer using clang with the ASan flags
+    #
+    # Use 'windows_hook_rtl_allocators=true:allocator_may_return_null=1' if your fuzzer has
+    # been compiled using MSVC with the ASan flags
+    # 
+    # $AsanOptions = 'windows_hook_rtl_allocators=true:allocator_may_return_null=1'
+    # $AsanOptions = 'allocator_may_return_null=1'
+
+    $AsanOptions = 'windows_hook_rtl_allocators=true:allocator_may_return_null=1'
+
+    # Use the 'machine' scope to make this permanent because the machine will reboot  
+    [Environment]::SetEnvironmentVariable('ASAN_OPTIONS', $AsanOptions, 'Machine')
+    Write-Log "Set ASAN_OPTIONS to $AsanOptions"
+
+    # Done. Useful to know that the script did not prematurely error out
+    Write-Log 'Setup script finished successfully'
+}
+
+# Write log statements into the event log.
+function Write-Log {
+    Param(
+        [Parameter(Position=0,
+                   Mandatory,
+                   ValueFromPipeline,
+                   ValueFromRemainingArguments)]
+        [String[]] $Messages,
+        [Parameter()] [Int] $EventId=0)
+    Begin {
+        $EventSource = 'onefuzz-setup.ps1'
+        $EventLog = 'Application'
+        if (![System.Diagnostics.EventLog]::SourceExists($EventSource)) {
+            New-EventLog -LogName $EventLog -Source $EventSource
+        }
+    }
+    Process {
+        $Messages.ForEach({
+            Write-EventLog -LogName $EventLog -Source $EventSource -EventId $EventId -EntryType Information -Message $_
+        })
+    }
+    End {}
+}
+
+# This function is used to exclude DLL's that the fuzzer is dependent on. The dependent DLL's
+# have been built with ASan and copied into the setup directory along with the fuzzer. 
+function Exclude-Library {
+    Param(
+        [Parameter(Position=0,
+                   Mandatory,
+                   ValueFromPipeline,
+                   ValueFromRemainingArguments)]
+        [String[]] $Libraries)
+    Begin {
+        $Path = 'HKLM:System\CurrentControlSet\Control\Session Manager'
+        $Name = 'ExcludeFromKnownDLLs'
+    }
+    Process {
+        # Get existing excluded libraries
+        $ExistingProperty = Get-ItemProperty -Path $Path -Name $Name -ErrorAction SilentlyContinue
+        $ExistingExclusions = @()
+
+        # Normalize DLL name to lowercase for comparison
+        if ($ExistingProperty -ne $null) {
+            $ExistingExclusions = $ExistingProperty.$Name.ForEach("ToLower")
+        }
+
+        # Normalize DLL name to lowercase for comparison, remove duplicates
+        $Libs = $Libraries.ForEach("ToLower") | Select-Object -Unique
+        Write-Log "Excluding libraries $Libs"
+
+        # Discard empty strings and libraries already excluded
+        $Libs = $Libs.Where({$_.Length -gt 0 -and !($ExistingExclusions.Contains($_))})
+
+        # If anything remains either add or update registry key
+        if ($Libs.Length -gt 0) {
+            if ($ExistingProperty -eq $null) {
+                # Create registry key to exclude DLLs
+                New-ItemProperty -Path $Path -Name $Name -PropertyType MultiString -Value $Libs
+                Write-Log "Created known DLLs exclusions with $Libs"
+            } else {
+                # Update registry key to exclude DLLs
+                Set-ItemProperty -Path $Path -Name $Name -Value ($ExistingProperty.$Name + $Libs)
+                Write-Log "Updated known DLLs exclusions with $Libs"
+            }
+        } else {
+            # DLLs already excluded
+            Write-Log "Known DLL exclusions already exist for $Libraries"
+        }
+    }
+    End {}
+}
+
+Execute-Setup