1 files changed, 41 insertions, 42 deletions

diff --git a/src/Agent.Windows/Network/TaskNetworkPolicy.cs b/src/Agent.Windows/Network/TaskNetworkPolicy.cs
index 9f78b3d..4bcd8b2 100644
--- a/src/Agent.Windows/Network/TaskNetworkPolicy.cs
+++ b/src/Agent.Windows/Network/TaskNetworkPolicy.cs
@@ -1,79 +1,78 @@
+using System.Net;
 using System.Security.Principal;
+using Agent.Admin.Client;
 using Agent.Windows.Recovery;
 using Agent.Windows.Security;
 
 namespace Agent.Windows.Network;
 
-public sealed class TaskNetworkPolicy : IDisposable
+public sealed class TaskNetworkPolicy : IAsyncDisposable
 {
-    private readonly WfpDynamicSession _session;
+    private readonly IElevatedNetworkSession _session;
     private readonly TaskResourceJournal _journal;
     private readonly ISecurityAuditSink _audit;
-    private readonly Guid _entryId;
     private readonly SecurityIdentifier _appContainerSid;
     private readonly Guid? _loopbackEntryId;
-    private readonly bool _loopbackAdded;
     private readonly Guid? _sessionId;
     private bool _disposed;
 
-    private TaskNetworkPolicy(WfpDynamicSession session, TaskResourceJournal journal, ISecurityAuditSink audit, Guid entryId,
-        SecurityIdentifier appContainerSid, Guid? loopbackEntryId, bool loopbackAdded, Guid? sessionId)
+    private TaskNetworkPolicy(IElevatedNetworkSession session, TaskResourceJournal journal, ISecurityAuditSink audit,
+        SecurityIdentifier appContainerSid, Guid? loopbackEntryId, Guid? sessionId)
     {
-        _session = session; _journal = journal; _audit = audit; _entryId = entryId;
-        _appContainerSid = appContainerSid; _loopbackEntryId = loopbackEntryId; _loopbackAdded = loopbackAdded;
+        _session = session;
+        _journal = journal;
+        _audit = audit;
+        _appContainerSid = appContainerSid;
+        _loopbackEntryId = loopbackEntryId;
         _sessionId = sessionId;
     }
 
-    public static TaskNetworkPolicy Apply(TaskResourceJournal journal, SecurityIdentifier appContainerSid, IReadOnlyList<NetworkEndpoint> allowances,
-        ISecurityAuditSink audit, Guid? sessionId = null, bool enableLoopback = true)
+    public static async Task<TaskNetworkPolicy> ApplyAsync(IElevatedNetworkSession session, TaskResourceJournal journal,
+        SecurityIdentifier appContainerSid, string networkAccess, IReadOnlyList<NetworkEndpoint> allowances,
+        ISecurityAuditSink audit, Guid? sessionId, CancellationToken cancellationToken)
     {
-        bool originallyExempt = enableLoopback && LoopbackExemption.Contains(appContainerSid);
-        Guid? loopbackEntryId = enableLoopback ? journal.Prepare("network.loopback_exemption", appContainerSid.Value,
+        bool needsLoopback = allowances.Any(static item => IPAddress.IsLoopback(item.Address));
+        bool originallyExempt = needsLoopback && await session.InspectLoopbackAsync(appContainerSid.Value, cancellationToken).ConfigureAwait(false);
+        Guid? loopbackEntryId = needsLoopback ? journal.Prepare("network.loopback_exemption", appContainerSid.Value,
             new Dictionary<string, string> { ["originally_present"] = originallyExempt.ToString() }) : null;
-        bool loopbackAdded = false;
-        Guid entryId = journal.Prepare("wfp.dynamic", appContainerSid.Value, new Dictionary<string, string>
-        {
-            ["session"] = "dynamic", ["allowances"] = string.Join(',', allowances.Select(static item => $"{item.Address}:{item.Port}"))
-        });
-        WfpDynamicSession session = WfpDynamicSession.Open();
         try
         {
-            if (enableLoopback)
-            {
-                loopbackAdded = LoopbackExemption.Ensure(appContainerSid);
-                journal.MarkApplied(loopbackEntryId!.Value, new Dictionary<string, string> { ["added"] = loopbackAdded.ToString() });
-            }
-            IReadOnlyList<ulong> ids = session.ApplyDefaultBlock(appContainerSid, allowances);
-            journal.MarkApplied(entryId, new Dictionary<string, string> { ["filter_ids"] = string.Join(',', ids) });
+            AdminNetworkEndpoint[] endpoints = allowances.Select(static item => new AdminNetworkEndpoint(item.Address.ToString(), item.Port)).ToArray();
+            await session.ApplyAsync(appContainerSid.Value, networkAccess, endpoints, originallyExempt, cancellationToken).ConfigureAwait(false);
+            if (loopbackEntryId is not null) journal.MarkApplied(loopbackEntryId.Value);
             audit.Write(new(DateTimeOffset.UtcNow, journal.TaskId, "network.rule.apply", "applied", new Dictionary<string, string>
             {
-                ["ipv4_ipv6_default_block"] = "true", ["filter_count"] = ids.Count.ToString(), ["session"] = "dynamic"
+                ["ipv4_ipv6_default_block"] = "true", ["endpoint_count"] = endpoints.Length.ToString(), ["session"] = "dynamic"
             }, SessionId: sessionId));
-            if (enableLoopback) audit.Write(new(DateTimeOffset.UtcNow, journal.TaskId, "network.loopback.apply", loopbackAdded ? "added" : "preserved",
-                new Dictionary<string, string> { ["app_container_sid"] = appContainerSid.Value }, SessionId: sessionId));
-            return new(session, journal, audit, entryId, appContainerSid, loopbackEntryId, loopbackAdded, sessionId);
+            if (loopbackEntryId is not null)
+                audit.Write(new(DateTimeOffset.UtcNow, journal.TaskId, "network.loopback.apply", originallyExempt ? "preserved" : "added",
+                    new Dictionary<string, string> { ["app_container_sid"] = appContainerSid.Value }, SessionId: sessionId));
+            return new(session, journal, audit, appContainerSid, loopbackEntryId, sessionId);
         }
         catch
         {
-            session.Dispose();
-            journal.MarkCleaned(entryId);
-            if (loopbackAdded) LoopbackExemption.Remove(appContainerSid);
-            if (loopbackEntryId is not null) journal.MarkCleaned(loopbackEntryId.Value);
+            bool cleaned = false;
+            try { await session.CleanupAsync(CancellationToken.None).ConfigureAwait(false); cleaned = true; } catch { }
+            if (cleaned && loopbackEntryId is not null) journal.MarkCleaned(loopbackEntryId.Value);
+            try { await session.DisposeAsync().ConfigureAwait(false); } catch { }
             throw;
         }
     }
 
-    public void Dispose()
+    public async ValueTask DisposeAsync()
     {
         if (_disposed) return;
-        _session.RemovePolicy();
-        _session.Dispose();
-        _journal.MarkCleaned(_entryId);
-        if (_loopbackAdded) LoopbackExemption.Remove(_appContainerSid);
-        if (_loopbackEntryId is not null) _journal.MarkCleaned(_loopbackEntryId.Value);
-        _audit.Write(new(DateTimeOffset.UtcNow, _journal.TaskId, "network.rule.remove", "removed", new Dictionary<string, string> { ["residual_filters"] = "0" }, SessionId: _sessionId));
-        if (_loopbackEntryId is not null) _audit.Write(new(DateTimeOffset.UtcNow, _journal.TaskId, "network.loopback.remove", "removed",
-            new Dictionary<string, string> { ["residual_exemptions"] = LoopbackExemption.Contains(_appContainerSid) ? "preexisting" : "0" }, SessionId: _sessionId));
+        try
+        {
+            await _session.CleanupAsync(CancellationToken.None).ConfigureAwait(false);
+            if (_loopbackEntryId is not null) _journal.MarkCleaned(_loopbackEntryId.Value);
+        }
+        finally { await _session.DisposeAsync().ConfigureAwait(false); }
+        _audit.Write(new(DateTimeOffset.UtcNow, _journal.TaskId, "network.rule.remove", "removed",
+            new Dictionary<string, string> { ["residual_filters"] = "0" }, SessionId: _sessionId));
+        if (_loopbackEntryId is not null)
+            _audit.Write(new(DateTimeOffset.UtcNow, _journal.TaskId, "network.loopback.remove", "restored",
+                new Dictionary<string, string> { ["app_container_sid"] = _appContainerSid.Value }, SessionId: _sessionId));
         _disposed = true;
     }
 }