From 722281a8db0c8e41859d07df8ea2ee8d7809ceeb Mon Sep 17 00:00:00 2001
From: Will Horne <will.horne@cloudrunner.com>
Date: Tue, 15 Oct 2024 21:40:44 -0700
Subject: [PATCH] Capture ScriptBlock source code from the topmost frame if
 available as a fallback

Ensure StackTraceString matches InvocationInfo when ScriptName is blank
---
 CHANGELOG.md                                  |  1 +
 .../Sentry/private/StackTraceProcessor.ps1    | 36 +++++++++++++++----
 modules/Sentry/public/Out-Sentry.ps1          |  9 +++++
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0b7e1ad..2647afb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### Fixes
 
 - StackTrace parsing on Windows Powershell 5.1 ([#50](https://github.com/getsentry/sentry-powershell/pull/50))
+- Context lines captured for PowerShell executed directly as a ScriptBlock  ([#60](https://github.com/getsentry/sentry-powershell/pull/60))
 
 ### Dependencies
 
diff --git a/modules/Sentry/private/StackTraceProcessor.ps1 b/modules/Sentry/private/StackTraceProcessor.ps1
index 6ad1bca..b1a6bff 100644
--- a/modules/Sentry/private/StackTraceProcessor.ps1
+++ b/modules/Sentry/private/StackTraceProcessor.ps1
@@ -4,6 +4,7 @@ class StackTraceProcessor : SentryEventProcessor
     [System.Management.Automation.InvocationInfo]$InvocationInfo
     [System.Management.Automation.CallStackFrame[]]$StackTraceFrames
     [string[]]$StackTraceString
+    [string[]]$ScriptBlockSource
     hidden [string[]] $modulePaths
     hidden [hashtable] $pwshModules = @{}
 
@@ -178,7 +179,7 @@ class StackTraceProcessor : SentryEventProcessor
             # Update module info
             $this.SetModule($sentryFrame)
             $sentryFrame.InApp = [string]::IsNullOrEmpty($sentryFrame.Module)
-            $this.SetContextLines($sentryFrame)
+            $this.SetContextLines($sentryFrame, $this.ScriptBlockSource)
         }
 
         $sentryFrames.Reverse()
@@ -190,7 +191,14 @@ class StackTraceProcessor : SentryEventProcessor
     hidden [Sentry.SentryStackFrame] CreateFrame([System.Management.Automation.InvocationInfo] $info)
     {
         $sentryFrame = [Sentry.SentryStackFrame]::new()
-        $sentryFrame.AbsolutePath = $info.ScriptName
+        if ("" -eq $info.ScriptName)
+        {
+            $sentryFrame.AbsolutePath = "<No file>"
+        }
+        else
+        {
+            $sentryFrame.AbsolutePath = $info.ScriptName
+        }
         $sentryFrame.LineNumber = $info.ScriptLineNumber
         $sentryFrame.ColumnNumber = $info.OffsetInLine
         $sentryFrame.ContextLine = $info.Line.TrimEnd()
@@ -275,18 +283,34 @@ class StackTraceProcessor : SentryEventProcessor
         }
     }
 
-    hidden SetContextLines([Sentry.SentryStackFrame] $sentryFrame)
+    hidden SetContextLines([Sentry.SentryStackFrame] $sentryFrame, [string[]] $scriptBlockSource)
     {
-        if ([string]::IsNullOrEmpty($sentryFrame.AbsolutePath) -or $sentryFrame.LineNumber -lt 1)
+        if ($sentryFrame.LineNumber -lt 1)
         {
             return
         }
 
-        if ((Test-Path $sentryFrame.AbsolutePath -IsValid) -and (Test-Path $sentryFrame.AbsolutePath -PathType Leaf))
+        $lines = $null
+
+        if (-not [string]::IsNullOrEmpty($sentryFrame.AbsolutePath) -and (Test-Path $sentryFrame.AbsolutePath -IsValid) -and (Test-Path $sentryFrame.AbsolutePath -PathType Leaf))
         {
             try
             {
                 $lines = Get-Content $sentryFrame.AbsolutePath -TotalCount ($sentryFrame.LineNumber + 5)
+            }
+            catch
+            {
+                Write-Warning "Failed to read context lines for $($sentryFrame.AbsolutePath): $_"
+            }
+        } elseif ($null -ne $scriptBlockSource) {
+            Write-Debug "Using fallback ScriptBlockSource for context lines."
+            $lines = $scriptBlockSource | Select-Object -First ($sentryFrame.LineNumber + 5)
+        }
+
+        if ($null -ne $lines)
+        {
+            try
+            {
                 if ($null -eq $sentryFrame.ContextLine)
                 {
                     $sentryFrame.ContextLine = $lines[$sentryFrame.LineNumber - 1]
@@ -303,7 +327,7 @@ class StackTraceProcessor : SentryEventProcessor
             }
             catch
             {
-                Write-Warning "Failed to read context lines for $($sentryFrame.AbsolutePath): $_"
+                Write-Warning "Failed to process context lines for $($sentryFrame.AbsolutePath): $_"
             }
         }
     }
diff --git a/modules/Sentry/public/Out-Sentry.ps1 b/modules/Sentry/public/Out-Sentry.ps1
index c605a1c..e868017 100644
--- a/modules/Sentry/public/Out-Sentry.ps1
+++ b/modules/Sentry/public/Out-Sentry.ps1
@@ -94,6 +94,15 @@ function Out-Sentry
             return
         }
 
+        # Use the PSCallStack to capture the source code of the main script that is being executed.
+        # This is used as a fallback in case the code is executed directly as a ScriptBlock from a hosted .NET environment with no .ps1 file.
+        # If the top frame is Script (i.e. the code is executed as a script file), we don't need to capture the source code.
+        $TopFrame = Get-PSCallStack | Select-Object -Last 1
+        if ("Script" -ne $TopFrame.InvocationInfo.MyCommand.CommandType)
+        {
+            $processor.ScriptBlockSource = $TopFrame.InvocationInfo.MyCommand.ScriptBlock.ToString() -split "`r`n"
+        }
+
         if ($options.AttachStackTrace -and $null -eq $processor.StackTraceFrames -and $null -eq $processor.StackTraceString)
         {
             $processor.StackTraceFrames = Get-PSCallStack | Select-Object -Skip 1