Update DebugClient into OpenGraphShims (#160)

Author: Kyle-Ye

Package.resolved

@@ -262,6 +262,7 @@ if attributeGraphCondition {
     }
 } else {
     openGraphShimsTarget.dependencies.append("OpenGraph")
+    package.platforms = [.iOS(.v13), .macOS(.v10_15), .macCatalyst(.v13), .tvOS(.v13), .watchOS(.v5)]
 }
 
 let compatibilityTestCondition = envEnable("OPENGRAPH_COMPATIBILITY_TEST")

Package.swift

@@ -28,6 +28,11 @@
 #include <errno.h>
 #include <assert.h>
 
+OG_EXTERN_C_BEGIN
+// DYLD_INTERPOSE does not work. Directly use the hook one here to match the semantics.
+bool og_variant_has_internal_diagnostics(const char *subsystem);
+OG_EXTERN_C_END
+
 // MARK: DebugServer public API Implementation
 
 OG::DebugServer::DebugServer(OGDebugServerMode mode) {
@@ -260,7 +265,7 @@
     if (
         (mode & OGDebugServerModeValid)
         && !OG::DebugServer::has_shared_server()
-        /*&& os_variant_has_internal_diagnostics("com.apple.AttributeGraph")*/
+        && og_variant_has_internal_diagnostics("org.OpenSwiftUIProject.OpenGraph")
     ) {
         _shared_server = new DebugServer(mode);
     }

Sources/OpenGraphCxx/DebugServer/DebugServer.mm

@@ -0,0 +1,34 @@
+//
+//  interpose.c
+//  OpenGraphCxx
+
+#include <OpenGraph/OGBase.h>
+#include "stdio.h"
+#include "stdbool.h"
+#include "string.h"
+
+#if OG_TARGET_OS_DARWIN
+extern bool os_variant_has_internal_diagnostics(const char *subsystem);
+#endif
+
+bool og_variant_has_internal_diagnostics(const char *subsystem) {
+    if (strcmp(subsystem, "org.OpenSwiftUIProject.OpenGraph") == 0) {
+        return true;
+    } else if (strcmp(subsystem, "com.apple.AttributeGraph") == 0) {
+        return true;
+    } else {
+        #if OG_TARGET_OS_DARWIN
+        return os_variant_has_internal_diagnostics(subsystem);
+        #else
+        return false;
+        #endif
+    }
+}
+
+#if OG_TARGET_OS_DARWIN
+#define DYLD_INTERPOSE(_replacement,_replacee) \
+   __attribute__((used)) static struct{ const void* replacement; const void* replacee; } _interpose_##_replacee \
+            __attribute__ ((section ("__DATA,__interpose"))) = { (const void*)(unsigned long)&_replacement, (const void*)(unsigned long)&_replacee };
+
+DYLD_INTERPOSE(og_variant_has_internal_diagnostics, os_variant_has_internal_diagnostics)
+#endif

Sources/OpenGraphCxx/DebugServer/interpose.c

@@ -0,0 +1,146 @@
+//
+//  DebugServerTests.swift
+//  OpenGraphShims
+
+#if canImport(Darwin)
+public import Foundation
+public import Network
+
+public struct ConnectionUpdates: AsyncSequence {
+    public typealias Element = NWConnection.State
+    
+    private let stream: AsyncStream<NWConnection.State>
+    
+    fileprivate init(stream: AsyncStream<NWConnection.State>) {
+        self.stream = stream
+    }
+    
+    public func makeAsyncIterator() -> AsyncStream<NWConnection.State>.AsyncIterator {
+        stream.makeAsyncIterator()
+    }
+}
+
+@_spi(Debug)
+public final class DebugClient {
+    public enum Command: String, CaseIterable, Hashable {
+        case graphDescription = "graph/description"
+        case profilerStart = "profiler/start"
+        case profilerStop = "profiler/stop"
+        case profilerReset = "profiler/reset"
+        case profilerMark = "profiler/mark"
+    }
+
+    private var connection: NWConnection?
+    private let queue = DispatchQueue(label: "org.openswiftuiproject.opengraph.debugclient")
+
+    public init() {}
+
+    public func connect(to url: URL) -> ConnectionUpdates {
+        guard let host = url.host, let port = url.port else {
+            return ConnectionUpdates(stream: AsyncStream { continuation in
+                continuation.yield(.failed(NWError.posix(.EINVAL)))
+                continuation.finish()
+            })
+        }
+        let nwHost = NWEndpoint.Host(host)
+        let nwPort = NWEndpoint.Port(integerLiteral: UInt16(port))
+        connection = NWConnection(host: nwHost, port: nwPort, using: .tcp)
+        let stream = AsyncStream<NWConnection.State> { continuation in
+            connection?.stateUpdateHandler = { state in
+                continuation.yield(state)
+                if case .cancelled = state {
+                    continuation.finish()
+                }
+            }
+            connection?.start(queue: queue)
+        }
+        return ConnectionUpdates(stream: stream)
+    }
+
+    public func sendMessage(token: UInt32, data: Data) async throws {
+        guard let connection else {
+            throw ClientError.notConnected
+        }
+        let header = DebugServerMessageHeader(
+            token: token,
+            unknown: 0,
+            length: numericCast(data.count),
+            unknown2: 0
+        )
+        let headerData = withUnsafePointer(to: header) {
+            Data(bytes: UnsafeRawPointer($0), count: MemoryLayout<DebugServerMessageHeader>.size)
+        }
+        try await send(data: headerData, on: connection)
+        guard header.length > 0 else {
+            return
+        }
+        try await send(data: data, on: connection)
+    }
+
+    public func receiveMessage() async throws -> (header: DebugServerMessageHeader, data: Data) {
+        guard let connection = connection else {
+            throw ClientError.notConnected
+        }
+        let headerData = try await receive(
+            length: MemoryLayout<DebugServerMessageHeader>.size,
+            from: connection
+        )
+        let header = headerData.withUnsafeBytes { bytes in
+            let buffer = bytes.bindMemory(to: UInt32.self)
+            return DebugServerMessageHeader(
+                token: buffer[0],
+                unknown: buffer[1],
+                length: buffer[2],
+                unknown2: buffer[3]
+            )
+        }
+        guard header.length > 0 else {
+            return (header: header, data: Data())
+        }
+        let payloadData = try await receive(
+            length: numericCast(header.length),
+            from: connection
+        )
+        return (header: header, data: payloadData)
+    }
+
+    public func disconnect() {
+        connection?.cancel()
+        connection = nil
+    }
+
+    private func send(data: Data, on connection: NWConnection) async throws {
+        return try await withCheckedThrowingContinuation { continuation in
+            connection.send(content: data, completion: .contentProcessed { error in
+                if let error {
+                    continuation.resume(throwing: error)
+                } else {
+                    continuation.resume()
+                }
+            })
+        }
+    }
+
+    private func receive(length: Int, from connection: NWConnection) async throws -> Data {
+        return try await withCheckedThrowingContinuation { continuation in
+            connection.receive(minimumIncompleteLength: length, maximumLength: length) { data, _, isComplete, error in
+                if let error {
+                    continuation.resume(throwing: error)
+                } else if let data {
+                    continuation.resume(returning: data)
+                } else {
+                    continuation.resume(throwing: ClientError.noDataReceived)
+                }
+            }
+        }
+    }
+}
+
+enum ClientError: Error {
+    case invalidURL
+    case notConnected
+    case connectionCancelled
+    case noDataReceived
+}
+
+#endif

Sources/OpenGraphShims/DebugClient.swift

@@ -0,0 +1,18 @@
+//
+//  DebugServerMessageHeader.swift
+//  OpenGraphShims
+
+@_spi(Debug)
+public struct DebugServerMessageHeader {
+    public let token: UInt32
+    public let unknown: UInt32
+    public let length: UInt32
+    public let unknown2: UInt32
+
+    public init(token: UInt32, unknown: UInt32, length: UInt32, unknown2: UInt32) {
+        self.token = token
+        self.unknown = unknown
+        self.length = length
+        self.unknown2 = unknown2
+    }
+}

Sources/OpenGraphShims/DebugServerMessageHeader.swift

@@ -15,14 +15,13 @@ struct DebugServerTests {
         #expect(DebugServer.copyURL() == nil)
     }
 
-    // TODO: hook via private API of dyld
     // To make AG start debugServer, we need to pass internal_diagnostics check.
     // In debug mode, we can breakpoint on `_ZN2AG11DebugServer5startEj` and
     // executable `reg write w0 1` after `internal_diagnostics` call.
-    // Or we can disable SIP on the target darwinOS and run `sudo sysctl kern.osvariant_status=xx` to workaround
-    @Test(
-        .disabled(if: compatibilityTestEnabled, "Skip on AG due to internal_diagnostics check"),
-    )
+    // Or we can disable SIP on the target darwinOS and run `sudo sysctl kern.osvariant_status=xx` to workaround.
+    // Or you can add `breakpoint set -n os_variant_has_internal_diagnostics -C "thread return 1"`
+    // to your lldbinit or run it before AGDebugServerStart call.
+    @Test(.disabled(if: compatibilityTestEnabled, "Skip on AG on CI due to internal_diagnostics check"))
     func testMode1() throws {
         let _ = try #require(DebugServer.start(mode: [.valid]))
         let url = try #require(DebugServer.copyURL()) as URL
@@ -33,9 +32,7 @@ struct DebugServerTests {
         DebugServer.stop()
     }
 
-    @Test(
-        .disabled(if: compatibilityTestEnabled, "Skip on AG due to internal_diagnostics check"),
-    )
+    @Test(.disabled(if: compatibilityTestEnabled, "Skip on AG on CI due to internal_diagnostics check"))
     func testMode3() throws {
         let _ = try #require(DebugServer.start(mode: [.valid, .networkInterface]))
         let url = try #require(DebugServer.copyURL()) as URL

Tests/OpenGraphCompatibilityTests/Debug/DebugServerTests.swift

@@ -0,0 +1 @@
+../../../Sources/OpenGraphShims/DebugClient.swift