container/Sources/Services/ContainerNetworkService/Client/NetworkClient.swift

//===----------------------------------------------------------------------===//
// Copyright © 2026 Apple Inc. and the container project authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//===----------------------------------------------------------------------===//

import ContainerResource
import ContainerXPC
import ContainerizationError
import ContainerizationExtras
import Foundation

/// A client for interacting with a single network.
public struct NetworkClient: Sendable {
    static let label = "com.apple.container.network"

    public static func machServiceLabel(id: String, plugin: String) -> String {
        "\(Self.label).\(plugin).\(id)"
    }

    private var machServiceLabel: String {
        Self.machServiceLabel(id: id, plugin: plugin)
    }

    let id: String
    let plugin: String

    /// Create a client for a network.
    public init(id: String, plugin: String) {
        self.id = id
        self.plugin = plugin
    }
}

// Runtime Methods
extension NetworkClient {
    public func state() async throws -> NetworkState {
        let request = XPCMessage(route: NetworkRoutes.state.rawValue)
        let client = createClient()

        let response = try await client.send(request)
        let state = try response.state()
        return state
    }

    public func allocate(
        hostname: String,
        macAddress: MACAddress? = nil
    ) async throws -> (attachment: Attachment, additionalData: XPCMessage?) {
        let request = XPCMessage(route: NetworkRoutes.allocate.rawValue)
        request.set(key: NetworkKeys.hostname.rawValue, value: hostname)
        if let macAddress = macAddress {
            request.set(key: NetworkKeys.macAddress.rawValue, value: macAddress.description)
        }

        let client = createClient()

        let response = try await client.send(request)
        let attachment = try response.attachment()
        let additionalData = response.additionalData()
        return (attachment, additionalData)
    }

    /// Open a persistent connection to the network helper.
    ///
    /// The returned session should be reused for `allocate(on:)` calls. The
    /// network helper automatically releases all allocations made over this
    /// session when it closes.
    public func connect() -> XPCClientSession {
        createClient().openSession()
    }

    /// Allocate a network attachment over an existing session.
    ///
    /// Use `connect()` to obtain a session, then pass it here. The session
    /// must remain open for the lifetime of the allocation; closing it
    /// releases the allocation on the network helper automatically.
    public func allocate(
        hostname: String,
        macAddress: MACAddress? = nil,
        on session: XPCClientSession
    ) async throws -> (attachment: Attachment, additionalData: XPCMessage?) {
        let request = XPCMessage(route: NetworkRoutes.allocate.rawValue)
        request.set(key: NetworkKeys.hostname.rawValue, value: hostname)
        if let macAddress = macAddress {
            request.set(key: NetworkKeys.macAddress.rawValue, value: macAddress.description)
        }
        let response = try await session.send(request)
        let attachment = try response.attachment()
        let additionalData = response.additionalData()
        return (attachment, additionalData)
    }

    public func lookup(hostname: String) async throws -> Attachment? {
        let request = XPCMessage(route: NetworkRoutes.lookup.rawValue)
        request.set(key: NetworkKeys.hostname.rawValue, value: hostname)

        let client = createClient()

        let response = try await client.send(request)
        return try response.dataNoCopy(key: NetworkKeys.attachment.rawValue).map {
            try JSONDecoder().decode(Attachment.self, from: $0)
        }
    }

    public func disableAllocator() async throws -> Bool {
        let request = XPCMessage(route: NetworkRoutes.disableAllocator.rawValue)

        let client = createClient()

        let response = try await client.send(request)
        return try response.allocatorDisabled()
    }

    private func createClient() -> XPCClient {
        XPCClient(service: machServiceLabel)
    }
}

extension XPCMessage {
    public func additionalData() -> XPCMessage? {
        guard let additionalData = xpc_dictionary_get_dictionary(self.underlying, NetworkKeys.additionalData.rawValue) else {
            return nil
        }
        return XPCMessage(object: additionalData)
    }

    public func allocatorDisabled() throws -> Bool {
        self.bool(key: NetworkKeys.allocatorDisabled.rawValue)
    }

    public func attachment() throws -> Attachment {
        let data = self.dataNoCopy(key: NetworkKeys.attachment.rawValue)
        guard let data else {
            throw ContainerizationError(.invalidArgument, message: "no network attachment snapshot data in message")
        }
        return try JSONDecoder().decode(Attachment.self, from: data)
    }

    public func hostname() throws -> String {
        let hostname = self.string(key: NetworkKeys.hostname.rawValue)
        guard let hostname else {
            throw ContainerizationError(.invalidArgument, message: "no hostname data in message")
        }
        return hostname
    }

    public func state() throws -> NetworkState {
        let data = self.dataNoCopy(key: NetworkKeys.state.rawValue)
        guard let data else {
            throw ContainerizationError(.invalidArgument, message: "no network snapshot data in message")
        }
        return try JSONDecoder().decode(NetworkState.self, from: data)
    }
}