Minor refactor

newrelic · Oct 27, 2023 · 2444790 · 2444790
1 parent 9fbe401
commit 2444790
Showing 1 changed file with 23 additions and 9 deletions.
diff --git a/...Relic/Agent/Extensions/Providers/Wrapper/AspNetCore6Plus/BrowserInjectingStreamWrapper.cs b/...Relic/Agent/Extensions/Providers/Wrapper/AspNetCore6Plus/BrowserInjectingStreamWrapper.cs
@@ -20,9 +20,6 @@ public class BrowserInjectingStreamWrapper : Stream
         private HttpContext _context;
         private bool _isContentLengthSet;
 
-
-        private const string InjectingRUM = "InjectingRUM";
-
         public BrowserInjectingStreamWrapper(IAgent agent, Stream baseStream, HttpContext context)
         {
             _agent = agent;
@@ -66,16 +63,16 @@ public override void SetLength(long value)
 
         public override void Write(byte[] buffer, int offset, int count)
         {
-            if (!_context.Items.ContainsKey(InjectingRUM))  // pass through without modification if we're already in the middle of injecting
+            if (!CurrentlyInjecting())  // pass through without modification if we're already in the middle of injecting
             {
                 if (IsHtmlResponse())
                 {
                     // Set a flag on the context to indicate we're in the middle of injecting - prevents multiple recursions when response compression is in use
-                    _context.Items.Add(InjectingRUM, null);
+                    StartInjecting();
                     var curBuf = buffer.AsMemory(offset, count).ToArray();
                     _agent.TryInjectBrowserScriptAsync(_context.Response.ContentType, _context.Request.Path.Value, curBuf, _baseStream)
                         .GetAwaiter().GetResult();
-                    _context.Items.Remove(InjectingRUM);
+                    FinishInjecting();
 
                     return;
                 }
@@ -88,14 +85,14 @@ public override void Write(byte[] buffer, int offset, int count)
 
         public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
         {
-            if (!_context.Items.ContainsKey(InjectingRUM))  // pass through without modification if we're already in the middle of injecting
+            if (!CurrentlyInjecting())  // pass through without modification if we're already in the middle of injecting
             {
                 if (IsHtmlResponse())
                 {
                     // Set a flag on the context to indicate we're in the middle of injecting - prevents multiple recursions when response compression is in use
-                    _context.Items.Add(InjectingRUM, null);
+                    StartInjecting();
                     await _agent.TryInjectBrowserScriptAsync(_context.Response.ContentType, _context.Request.Path.Value, buffer.ToArray(), _baseStream);
-                    _context.Items.Remove(InjectingRUM);
+                    FinishInjecting();
 
                     return;
                 }
@@ -105,6 +102,23 @@ public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, Cancella
                 await _baseStream.WriteAsync(buffer, cancellationToken);
         }
 
+        private const string InjectingRUM = "InjectingRUM";
+
+        private void FinishInjecting()
+        {
+            _context.Items.Remove(InjectingRUM);
+        }
+
+        private void StartInjecting()
+        {
+            _context.Items.Add(InjectingRUM, null);
+        }
+
+        private bool CurrentlyInjecting()
+        {
+            return _context.Items.ContainsKey(InjectingRUM);
+        }
+
         public override async ValueTask DisposeAsync()
         {
             _context = null;