Improved memory management in ProcThreadAttribute (#19)

* Improved memory management in ProcThreadAttribute * Dont hardcode swift version, just pick the first. You should only have one version installed anyway. * Plug the leaks
FabricMC · Dec 9, 2024 · 7506ed3 · 7506ed3
1 parent 97e561d
commit 7506ed3
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 20 deletions.
diff --git a/windows/Sources/Packager/SwiftRedistributables.swift b/windows/Sources/Packager/SwiftRedistributables.swift
@@ -6,8 +6,14 @@ class SwiftRedistributables {
     try resetDir(out)
 
     let redistributables = try swiftRedistributables()
+    let versionedDir = try redistributables.directoryContents().first
+
+    guard let versionDir = versionedDir else {
+      throw PackagerError("Could not find versioned directory in \(redistributables)")
+    }
+
     let archName = arch == .arm64 ? "arm64" : "amd64"
-    let rtl = redistributables.child("6.0.2").child("rtl.\(archName).msm")
+    let rtl = versionDir.child("rtl.\(archName).msm")
 
     guard rtl.exists() else {
       throw PackagerError("Could not find \(rtl)")

diff --git a/windows/Sources/Sandbox/ProcThreadAttributes.swift b/windows/Sources/Sandbox/ProcThreadAttributes.swift
@@ -3,12 +3,15 @@ import WinSDKExtras
 import WindowsUtils
 
 class ProcThreadAttributeList {
+  let attributes: [ProcThreadAttribute]
   let attributeList: LPPROC_THREAD_ATTRIBUTE_LIST
 
   init(attributes: [ProcThreadAttribute]) throws {
     var attributeList = try ProcThreadAttributeList.createAttributeList(attributeCount: attributes.count)
 
-    for attribute in attributes {
+    self.attributes = attributes
+
+    for attribute in self.attributes {
       try attribute.apply(&attributeList)
     }
 
@@ -52,55 +55,67 @@ protocol ProcThreadAttribute {
 internal func updateProcThreadAttribute<T>(
   attributeList: inout LPPROC_THREAD_ATTRIBUTE_LIST,
   attribute: DWORD,
-  value: inout T,
+  value: UnsafeMutablePointer<T>,
   size: Int
 ) throws {
-  let result = withUnsafeMutablePointer(to: &value) {
-    UpdateProcThreadAttribute(
-      attributeList,
-      0,
-      DWORD_PTR(attribute),
-      $0,
-      SIZE_T(size),
-      nil,
-      nil
-    )
-  }
+  let result = UpdateProcThreadAttribute(
+    attributeList,
+    0,
+    DWORD_PTR(attribute),
+    value,
+    SIZE_T(size),
+    nil,
+    nil
+  )
   guard result else {
     throw Win32Error("UpdateProcThreadAttribute")
   }
 }
 
 class SecurityCapabilitiesProcThreadAttribute: ProcThreadAttribute {
-  var securityCapabilities: SECURITY_CAPABILITIES
+  var securityCapabilities: UnsafeMutablePointer<SECURITY_CAPABILITIES>
   init(container: AppContainer, securityAttributes: UnsafeMutableBufferPointer<SID_AND_ATTRIBUTES>)
   {
-    self.securityCapabilities = SECURITY_CAPABILITIES(
+    self.securityCapabilities = UnsafeMutablePointer<SECURITY_CAPABILITIES>.allocate(capacity: 1)
+    self.securityCapabilities.pointee = SECURITY_CAPABILITIES(
       AppContainerSid: container.sid.value,
       Capabilities: securityAttributes.baseAddress,
       CapabilityCount: DWORD(securityAttributes.count),
       Reserved: 0
     )
   }
 
+  deinit {
+    self.securityCapabilities.deallocate()
+  }
+
   func apply(_ attributeList: inout LPPROC_THREAD_ATTRIBUTE_LIST) throws {
     try updateProcThreadAttribute(
       attributeList: &attributeList,
       attribute: _PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES(),
-      value: &self.securityCapabilities,
+      value: self.securityCapabilities,
       size: MemoryLayout<SECURITY_CAPABILITIES>.size
     )
   }
 }
 
 class LessPrivilegedAppContainerProcThreadAttribute: ProcThreadAttribute {
-  func apply(_ attributeList: inout LPPROC_THREAD_ATTRIBUTE_LIST) throws {
-    var enabled: DWORD = 1
+  var enabled: UnsafeMutablePointer<DWORD>
 
+  init() {
+    self.enabled = UnsafeMutablePointer<DWORD>.allocate(capacity: 1)
+    self.enabled.pointee = 1
+  }
+
+  deinit {
+    self.enabled.deallocate()
+  }
+
+  func apply(_ attributeList: inout LPPROC_THREAD_ATTRIBUTE_LIST) throws {
     try updateProcThreadAttribute(
       attributeList: &attributeList,
       attribute: _PROC_THREAD_ATTRIBUTE_ALL_APPLICATION_PACKAGES_POLICY(),
-      value: &enabled,
+      value: self.enabled,
       size: MemoryLayout<DWORD>.size
     )
   }