824 lines
35 KiB
Swift
824 lines
35 KiB
Swift
import Foundation
|
|
|
|
#if canImport(FoundationNetworking)
|
|
import FoundationNetworking
|
|
#endif
|
|
#if canImport(Network)
|
|
import Network
|
|
#endif
|
|
|
|
#if os(Linux)
|
|
import Glibc
|
|
#endif
|
|
/// A loopback token-injecting reverse proxy for Claude's API traffic.
|
|
///
|
|
/// The problem it solves: Nucleic keeps the rotating refresh token broker-only and hands a session
|
|
/// only a *static* `CLAUDE_CODE_OAUTH_TOKEN`, fixed for the life of the `claude` process. Access
|
|
/// tokens are short-lived, so a single long turn can outlive its token and 401 mid-turn — and
|
|
/// `claude` can't self-refresh because it never held the refresh token. This proxy closes that gap
|
|
/// without loosening the security posture: `claude` is pointed at `http://127.0.0.1:<port>` via
|
|
/// `ANTHROPIC_BASE_URL`, and on **every request** the proxy overwrites the `Authorization` header
|
|
/// with a broker-fresh access token (``ClaudeCredentialBroker/freshAccessToken()``) before
|
|
/// forwarding to `https://api.anthropic.com`. The refresh token still never leaves the broker; a
|
|
/// mid-turn request simply picks up a freshly-rotated bearer.
|
|
///
|
|
/// The HTTP translation (parse request → rewrite auth → serialize the upstream response back, with
|
|
/// SSE streamed through untouched) is pure and unit-tested through an injected ``Upstream``; the
|
|
/// listeners (`AF_UNIX` for the container relay, loopback TCP for host runs) mirror
|
|
/// ``MCPApprovalServer``'s transport-agnostic accept loop.
|
|
public actor ClaudeTokenProxy {
|
|
/// Supplies the current access token to inject, refreshing through Nucleic's own OAuth when near
|
|
/// expiry. `nil` result ⇒ no subscription login (API-key mode or signed out); the proxy then
|
|
/// forwards whatever `Authorization` the client sent, so an API key still works.
|
|
public typealias TokenProvider = @Sendable () async -> String?
|
|
|
|
/// One upstream round-trip, streamed. Injected so the HTTP translation is testable without the
|
|
/// network; the default drives `URLSession` against `api.anthropic.com`.
|
|
public typealias Upstream = @Sendable (_ request: UpstreamRequest) async throws -> UpstreamResponse
|
|
|
|
/// The forwarded request as the proxy reconstructs it from the inbound connection.
|
|
public struct UpstreamRequest: Sendable {
|
|
public var method: String
|
|
/// Path + query, e.g. `/v1/messages?beta=true`.
|
|
public var path: String
|
|
/// Header order preserved; `Authorization` already rewritten to the fresh bearer (or dropped
|
|
/// when there is no bearer and the client sent none).
|
|
public var headers: [(name: String, value: String)]
|
|
public var body: Data
|
|
}
|
|
|
|
/// The upstream response, its body streamed so SSE (`text/event-stream`) flows to the client as
|
|
/// it arrives rather than buffering a long-lived turn.
|
|
public struct UpstreamResponse: Sendable {
|
|
public var status: Int
|
|
public var headers: [(name: String, value: String)]
|
|
public var body: AsyncThrowingStream<Data, Error>
|
|
public init(
|
|
status: Int, headers: [(name: String, value: String)],
|
|
body: AsyncThrowingStream<Data, Error>
|
|
) {
|
|
self.status = status
|
|
self.headers = headers
|
|
self.body = body
|
|
}
|
|
}
|
|
|
|
private let tokenProvider: TokenProvider
|
|
private let upstream: Upstream
|
|
|
|
public init(tokenProvider: @escaping TokenProvider, upstream: @escaping Upstream = ClaudeTokenProxy.liveUpstream) {
|
|
self.tokenProvider = tokenProvider
|
|
self.upstream = upstream
|
|
}
|
|
|
|
// MARK: - Request handling (pure; unit-tested over in-memory byte streams)
|
|
|
|
/// Serve one client connection: read one HTTP/1.1 request, rewrite its auth, forward it, and
|
|
/// stream the response back. One request per connection (`Connection: close`) — simplest correct
|
|
/// framing for a mix of short JSON calls and long SSE streams, and `claude`'s pooled client just
|
|
/// opens a fresh connection for the next request.
|
|
func serve(reader: ByteReader, writer: ByteWriter) async {
|
|
defer { Task { await writer.close() } }
|
|
guard let request = await Self.readRequest(from: reader) else {
|
|
await Self.writeSimpleResponse(400, "Bad Request", to: writer)
|
|
return
|
|
}
|
|
// Overwrite Authorization with a broker-fresh bearer. When the broker has no token (API-key
|
|
// mode / signed out) leave whatever the client sent, so an `ANTHROPIC_API_KEY` path is
|
|
// untouched. `x-api-key` (API-key auth) is likewise left as the client set it.
|
|
var headers = request.headers
|
|
if let bearer = await tokenProvider(), !bearer.isEmpty {
|
|
headers.removeAll { $0.name.caseInsensitiveCompare("authorization") == .orderedSame }
|
|
headers.append((name: "Authorization", value: "Bearer \(bearer)"))
|
|
}
|
|
let forwarded = UpstreamRequest(
|
|
method: request.method, path: request.path, headers: headers, body: request.body)
|
|
|
|
let response: UpstreamResponse
|
|
do {
|
|
response = try await upstream(forwarded)
|
|
} catch {
|
|
claudeAuthLog.error("Claude token proxy upstream failed: \(String(describing: error), privacy: .public)")
|
|
await Self.writeSimpleResponse(502, "Bad Gateway", to: writer)
|
|
return
|
|
}
|
|
await Self.writeResponse(response, to: writer)
|
|
}
|
|
|
|
/// Parse one HTTP/1.1 request: request line, headers, then a `Content-Length` body (the only
|
|
/// framing `claude`'s client sends — JSON bodies with an explicit length, or no body). Returns
|
|
/// nil on a malformed head or a truncated body.
|
|
static func readRequest(from reader: ByteReader) async -> ParsedRequest? {
|
|
let head = await readUntilDoubleCRLF(reader)
|
|
guard let (headText, leftover) = head else { return nil }
|
|
var lines = headText.split(separator: "\r\n", omittingEmptySubsequences: false).map(String.init)
|
|
guard !lines.isEmpty else { return nil }
|
|
let requestLine = lines.removeFirst().split(separator: " ").map(String.init)
|
|
guard requestLine.count >= 2 else { return nil }
|
|
let method = requestLine[0]
|
|
let path = requestLine[1]
|
|
|
|
var headers: [(name: String, value: String)] = []
|
|
var contentLength = 0
|
|
for line in lines where !line.isEmpty {
|
|
guard let colon = line.firstIndex(of: ":") else { continue }
|
|
let name = String(line[..<colon]).trimmingCharacters(in: .whitespaces)
|
|
let value = String(line[line.index(after: colon)...]).trimmingCharacters(in: .whitespaces)
|
|
// Drop hop-by-hop / host-specific headers we re-derive; keep everything else (anthropic
|
|
// beta, version, content-type, user-agent…) so the upstream sees an unchanged request.
|
|
let lower = name.lowercased()
|
|
if lower == "host" || lower == "connection" || lower == "proxy-connection" { continue }
|
|
if lower == "content-length", let n = Int(value) { contentLength = n }
|
|
headers.append((name: name, value: value))
|
|
}
|
|
|
|
var body = leftover
|
|
while body.count < contentLength {
|
|
guard let chunk = await reader.read(), !chunk.isEmpty else { return nil }
|
|
body.append(chunk)
|
|
}
|
|
if body.count > contentLength { body = body.prefix(contentLength) } // no pipelining
|
|
return ParsedRequest(method: method, path: path, headers: headers, body: body)
|
|
}
|
|
|
|
/// A parsed inbound request (distinct from the auth-rewritten ``UpstreamRequest``).
|
|
struct ParsedRequest: Sendable, Equatable {
|
|
var method: String
|
|
var path: String
|
|
var headers: [(name: String, value: String)]
|
|
var body: Data
|
|
static func == (l: ParsedRequest, r: ParsedRequest) -> Bool {
|
|
l.method == r.method && l.path == r.path && l.body == r.body
|
|
&& l.headers.count == r.headers.count
|
|
&& zip(l.headers, r.headers).allSatisfy { $0.name == $1.name && $0.value == $1.value }
|
|
}
|
|
}
|
|
|
|
/// Serialize an upstream response back to the client: status line, headers (hop-by-hop and stale
|
|
/// representation framing stripped), `Connection: close`, then the body chunks as they arrive.
|
|
/// URLSession transparently decodes compressed upstream bodies, so forwarding the original
|
|
/// `Content-Encoding` would make Claude's fetch decode those bytes a second time (`ZlibError`).
|
|
static func writeResponse(_ response: UpstreamResponse, to writer: ByteWriter) async {
|
|
var head = "HTTP/1.1 \(response.status) \(reasonPhrase(response.status))\r\n"
|
|
for (name, value) in response.headers {
|
|
let lower = name.lowercased()
|
|
if lower == "content-length" || lower == "transfer-encoding"
|
|
|| lower == "content-encoding" || lower == "connection" { continue }
|
|
head += "\(name): \(value)\r\n"
|
|
}
|
|
head += "Connection: close\r\n\r\n"
|
|
await writer.write(Data(head.utf8))
|
|
do {
|
|
for try await chunk in response.body where !chunk.isEmpty {
|
|
await writer.write(chunk)
|
|
}
|
|
} catch {
|
|
// The upstream stream faulted mid-body; we've already sent the head, so the best we can
|
|
// do is close (below, via the caller's defer). The client sees a truncated stream.
|
|
claudeAuthLog.notice("Claude token proxy response stream ended early: \(String(describing: error), privacy: .public)")
|
|
}
|
|
}
|
|
|
|
static func writeSimpleResponse(_ status: Int, _ phrase: String, to writer: ByteWriter) async {
|
|
let body = Data(phrase.utf8)
|
|
let head = "HTTP/1.1 \(status) \(phrase)\r\nContent-Length: \(body.count)\r\nConnection: close\r\n\r\n"
|
|
await writer.write(Data(head.utf8))
|
|
await writer.write(body)
|
|
}
|
|
|
|
/// Read from `reader` until the header terminator (`\r\n\r\n`), returning the header text (without
|
|
/// the terminator) and any bytes already read past it (the start of the body). Nil on EOF first.
|
|
private static func readUntilDoubleCRLF(_ reader: ByteReader) async -> (head: String, leftover: Data)? {
|
|
var buffer = Data()
|
|
let terminator = Data("\r\n\r\n".utf8)
|
|
// Bound the header size so a nonsense client can't grow the buffer without limit.
|
|
let maxHead = 64 * 1024
|
|
while buffer.count <= maxHead {
|
|
if let range = buffer.range(of: terminator) {
|
|
let head = String(decoding: buffer[..<range.lowerBound], as: UTF8.self)
|
|
let leftover = Data(buffer[range.upperBound...])
|
|
return (head, leftover)
|
|
}
|
|
guard let chunk = await reader.read(), !chunk.isEmpty else { return nil }
|
|
buffer.append(chunk)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
private static func reasonPhrase(_ status: Int) -> String {
|
|
switch status {
|
|
case 200: return "OK"
|
|
case 201: return "Created"
|
|
case 204: return "No Content"
|
|
case 400: return "Bad Request"
|
|
case 401: return "Unauthorized"
|
|
case 403: return "Forbidden"
|
|
case 429: return "Too Many Requests"
|
|
case 500: return "Internal Server Error"
|
|
case 502: return "Bad Gateway"
|
|
case 503: return "Service Unavailable"
|
|
default: return status < 400 ? "OK" : "Error"
|
|
}
|
|
}
|
|
|
|
// MARK: - Live upstream (URLSession → api.anthropic.com)
|
|
|
|
/// Where forwarded requests go. Overridable via `ANTHROPIC_PROXY_UPSTREAM` for local testing;
|
|
/// defaults to the real API base.
|
|
public static let upstreamBase: String =
|
|
ProcessInfo.processInfo.environment["ANTHROPIC_PROXY_UPSTREAM"].flatMap { $0.isEmpty ? nil : $0 }
|
|
?? "https://api.anthropic.com"
|
|
|
|
/// The default `Upstream`: forward to `api.anthropic.com` and stream the response body. Uses
|
|
/// `URLSession.bytes(for:)` so an SSE stream is relayed incrementally rather than buffered.
|
|
public static let liveUpstream: Upstream = { request in
|
|
guard let url = URL(string: upstreamBase + request.path) else {
|
|
throw ProxyError.badUpstreamURL(request.path)
|
|
}
|
|
var urlRequest = URLRequest(url: url)
|
|
urlRequest.httpMethod = request.method
|
|
for (name, value) in request.headers {
|
|
urlRequest.setValue(value, forHTTPHeaderField: name)
|
|
}
|
|
// Claude's fetch advertises gzip/br. URLSession decodes compressed responses before exposing
|
|
// their bytes but can retain the upstream Content-Encoding header, which would make the
|
|
// downstream fetch decompress the decoded stream again. Request identity as the primary
|
|
// defense; writeResponse also drops stale encoding metadata defensively.
|
|
urlRequest.setValue("identity", forHTTPHeaderField: "Accept-Encoding")
|
|
if !request.body.isEmpty { urlRequest.httpBody = request.body }
|
|
|
|
#if canImport(FoundationNetworking)
|
|
// FoundationNetworking has no async `bytes(for:)`. A data delegate preserves the SSE
|
|
// stream instead of buffering an entire Claude turn before returning it.
|
|
let (http, stream) = try await StreamingUpstreamTask.start(urlRequest)
|
|
return UpstreamResponse(
|
|
status: http.statusCode, headers: ClaudeTokenProxy.headerPairs(http), body: stream)
|
|
#else
|
|
let (bytes, response) = try await URLSession.shared.bytes(for: urlRequest)
|
|
let http = response as? HTTPURLResponse
|
|
let stream = AsyncThrowingStream<Data, Error> { continuation in
|
|
let task = Task {
|
|
do {
|
|
var buffer = Data()
|
|
buffer.reserveCapacity(16 * 1024)
|
|
for try await byte in bytes {
|
|
buffer.append(byte)
|
|
// Flush on newline (SSE frames end in \n\n) or when the buffer fills, so a
|
|
// streamed event reaches the client promptly without a per-byte write storm.
|
|
if byte == 0x0A || buffer.count >= 16 * 1024 {
|
|
continuation.yield(buffer)
|
|
buffer.removeAll(keepingCapacity: true)
|
|
}
|
|
}
|
|
if !buffer.isEmpty { continuation.yield(buffer) }
|
|
continuation.finish()
|
|
} catch {
|
|
continuation.finish(throwing: error)
|
|
}
|
|
}
|
|
continuation.onTermination = { _ in task.cancel() }
|
|
}
|
|
return UpstreamResponse(
|
|
status: http?.statusCode ?? 502,
|
|
headers: ClaudeTokenProxy.headerPairs(http), body: stream)
|
|
#endif
|
|
}
|
|
|
|
private static func headerPairs(_ http: HTTPURLResponse?) -> [(name: String, value: String)] {
|
|
guard let http else { return [] }
|
|
return http.allHeaderFields.compactMap { key, value in
|
|
guard let name = key as? String else { return nil }
|
|
return (name: name, value: String(describing: value))
|
|
}
|
|
}
|
|
|
|
public enum ProxyError: Error, Equatable {
|
|
case badUpstreamURL(String)
|
|
case listenerFailed(String)
|
|
}
|
|
|
|
// MARK: - Listeners
|
|
//
|
|
// Darwin uses Network.framework for host TCP and a BSD AF_UNIX listener for container relays.
|
|
// Linux uses a BSD loopback listener: all runner chats share one credential broker, rather than
|
|
// letting several Claude processes race Anthropic's rotating refresh token.
|
|
|
|
#if canImport(Network)
|
|
private var tcpListener: NWListener?
|
|
private var unixListenFD: Int32?
|
|
private var unixAcceptSource: DispatchSourceRead?
|
|
private var boundUnixPath: String?
|
|
|
|
/// Start a loopback TCP listener for a host run. Returns the bound ephemeral port; point `claude`
|
|
/// at `http://127.0.0.1:<port>` via `ANTHROPIC_BASE_URL`.
|
|
@discardableResult
|
|
public func start(host: String = "127.0.0.1") async throws -> UInt16 {
|
|
let params = NWParameters.tcp
|
|
params.requiredLocalEndpoint = NWEndpoint.hostPort(host: NWEndpoint.Host(host), port: .any)
|
|
let listener = try NWListener(using: params)
|
|
listener.newConnectionHandler = { [weak self] connection in
|
|
connection.start(queue: .global(qos: .userInitiated))
|
|
let conn = NWSocketConn(connection)
|
|
Task { await self?.serve(reader: conn, writer: conn) }
|
|
}
|
|
let resume = ProxyOnce()
|
|
let port: UInt16 = try await withCheckedThrowingContinuation { cont in
|
|
listener.stateUpdateHandler = { state in
|
|
switch state {
|
|
case .ready:
|
|
if let p = listener.port?.rawValue, resume.claim() { cont.resume(returning: p) }
|
|
case .failed(let error):
|
|
if resume.claim() { cont.resume(throwing: error) }
|
|
case .cancelled:
|
|
if resume.claim() { cont.resume(throwing: ProxyError.listenerFailed("cancelled")) }
|
|
default:
|
|
break
|
|
}
|
|
}
|
|
listener.start(queue: .global(qos: .userInitiated))
|
|
}
|
|
tcpListener = listener
|
|
return port
|
|
}
|
|
|
|
/// Start an `AF_UNIX` listener at `path` (the host end of a container's relayed proxy socket).
|
|
/// Mirrors ``MCPApprovalServer``'s real-socket bind so the containerization framework's host-side
|
|
/// relay can connect to it.
|
|
public func start(unixSocketPath path: String) async throws {
|
|
let fm = FileManager.default
|
|
try? fm.createDirectory(
|
|
atPath: (path as NSString).deletingLastPathComponent, withIntermediateDirectories: true)
|
|
try? fm.removeItem(atPath: path)
|
|
|
|
let fd = socket(AF_UNIX, SOCK_STREAM, 0)
|
|
guard fd >= 0 else {
|
|
throw ProxyError.listenerFailed("socket(): \(String(cString: strerror(errno)))")
|
|
}
|
|
_ = fcntl(fd, F_SETFD, FD_CLOEXEC)
|
|
|
|
var addr = sockaddr_un()
|
|
addr.sun_family = sa_family_t(AF_UNIX)
|
|
let pathBytes = Array(path.utf8)
|
|
let capacity = MemoryLayout.size(ofValue: addr.sun_path)
|
|
guard pathBytes.count < capacity else {
|
|
close(fd)
|
|
throw ProxyError.listenerFailed("socket path too long: \(path)")
|
|
}
|
|
withUnsafeMutablePointer(to: &addr.sun_path) { raw in
|
|
raw.withMemoryRebound(to: CChar.self, capacity: capacity) { dst in
|
|
for (i, byte) in pathBytes.enumerated() { dst[i] = CChar(bitPattern: byte) }
|
|
dst[pathBytes.count] = 0
|
|
}
|
|
}
|
|
let bindRC = withUnsafePointer(to: &addr) { ptr in
|
|
ptr.withMemoryRebound(to: sockaddr.self, capacity: 1) {
|
|
Darwin.bind(fd, $0, socklen_t(MemoryLayout<sockaddr_un>.size))
|
|
}
|
|
}
|
|
guard bindRC == 0 else {
|
|
let err = String(cString: strerror(errno))
|
|
close(fd)
|
|
throw ProxyError.listenerFailed("bind(\(path)): \(err)")
|
|
}
|
|
chmod(path, 0o600)
|
|
guard listen(fd, SOMAXCONN) == 0 else {
|
|
let err = String(cString: strerror(errno))
|
|
close(fd)
|
|
try? fm.removeItem(atPath: path)
|
|
throw ProxyError.listenerFailed("listen(\(path)): \(err)")
|
|
}
|
|
|
|
let source = DispatchSource.makeReadSource(
|
|
fileDescriptor: fd, queue: .global(qos: .userInitiated))
|
|
source.setEventHandler { [weak self] in
|
|
let client = accept(fd, nil, nil)
|
|
guard client >= 0 else { return }
|
|
_ = fcntl(client, F_SETFD, FD_CLOEXEC)
|
|
let conn = UnixSocketConn(fd: client)
|
|
Task { await self?.serve(reader: conn, writer: conn) }
|
|
}
|
|
source.setCancelHandler { close(fd) }
|
|
source.resume()
|
|
|
|
unixListenFD = fd
|
|
unixAcceptSource = source
|
|
boundUnixPath = path
|
|
}
|
|
|
|
/// Tear down every listener (session teardown). Idempotent.
|
|
public func stop() {
|
|
tcpListener?.cancel()
|
|
tcpListener = nil
|
|
unixAcceptSource?.cancel() // cancel handler closes the fd
|
|
unixAcceptSource = nil
|
|
unixListenFD = nil
|
|
if let path = boundUnixPath { try? FileManager.default.removeItem(atPath: path) }
|
|
boundUnixPath = nil
|
|
}
|
|
#elseif os(Linux)
|
|
private var linuxTCPListenFD: Int32?
|
|
private var linuxTCPAcceptSource: DispatchSourceRead?
|
|
|
|
/// Linux/headless loopback listener. This is intentionally host-only: the cloud container is
|
|
/// already the sandbox, so there is no second AF_UNIX container-relay hop on this platform.
|
|
@discardableResult
|
|
public func start(host: String = "127.0.0.1") async throws -> UInt16 {
|
|
let fd = Glibc.socket(AF_INET, Int32(SOCK_STREAM.rawValue), 0)
|
|
guard fd >= 0 else {
|
|
throw ProxyError.listenerFailed("socket(): \(String(cString: strerror(errno)))")
|
|
}
|
|
_ = fcntl(fd, F_SETFD, FD_CLOEXEC)
|
|
var one: Int32 = 1
|
|
_ = setsockopt(
|
|
fd, SOL_SOCKET, SO_REUSEADDR, &one, socklen_t(MemoryLayout<Int32>.size))
|
|
|
|
var inAddress = in_addr()
|
|
guard inet_pton(AF_INET, host, &inAddress) == 1 else {
|
|
Glibc.close(fd)
|
|
throw ProxyError.listenerFailed("invalid IPv4 listen host: \(host)")
|
|
}
|
|
var address = sockaddr_in()
|
|
address.sin_family = sa_family_t(AF_INET)
|
|
address.sin_port = in_port_t(0).bigEndian
|
|
address.sin_addr = inAddress
|
|
let bindResult = withUnsafePointer(to: &address) { pointer in
|
|
pointer.withMemoryRebound(to: sockaddr.self, capacity: 1) {
|
|
Glibc.bind(fd, $0, socklen_t(MemoryLayout<sockaddr_in>.size))
|
|
}
|
|
}
|
|
guard bindResult == 0 else {
|
|
let message = String(cString: strerror(errno))
|
|
Glibc.close(fd)
|
|
throw ProxyError.listenerFailed("bind(\(host)): \(message)")
|
|
}
|
|
guard Glibc.listen(fd, SOMAXCONN) == 0 else {
|
|
let message = String(cString: strerror(errno))
|
|
Glibc.close(fd)
|
|
throw ProxyError.listenerFailed("listen(): \(message)")
|
|
}
|
|
|
|
var bound = sockaddr_in()
|
|
var boundLength = socklen_t(MemoryLayout<sockaddr_in>.size)
|
|
let nameResult = withUnsafeMutablePointer(to: &bound) { pointer in
|
|
pointer.withMemoryRebound(to: sockaddr.self, capacity: 1) {
|
|
Glibc.getsockname(fd, $0, &boundLength)
|
|
}
|
|
}
|
|
guard nameResult == 0 else {
|
|
let message = String(cString: strerror(errno))
|
|
Glibc.close(fd)
|
|
throw ProxyError.listenerFailed("getsockname(): \(message)")
|
|
}
|
|
let port = UInt16(bigEndian: bound.sin_port)
|
|
|
|
let flags = fcntl(fd, F_GETFL, 0)
|
|
if flags >= 0 { _ = fcntl(fd, F_SETFL, flags | O_NONBLOCK) }
|
|
let source = DispatchSource.makeReadSource(
|
|
fileDescriptor: fd, queue: .global(qos: .userInitiated))
|
|
source.setEventHandler { [weak self] in
|
|
while true {
|
|
let client = Glibc.accept(fd, nil, nil)
|
|
if client >= 0 {
|
|
_ = fcntl(client, F_SETFD, FD_CLOEXEC)
|
|
let connection = LinuxSocketConn(fd: client)
|
|
Task { await self?.serve(reader: connection, writer: connection) }
|
|
continue
|
|
}
|
|
if errno == EINTR { continue }
|
|
if errno == EAGAIN || errno == EWOULDBLOCK { break }
|
|
break
|
|
}
|
|
}
|
|
source.setCancelHandler { Glibc.close(fd) }
|
|
source.resume()
|
|
|
|
linuxTCPListenFD = fd
|
|
linuxTCPAcceptSource = source
|
|
return port
|
|
}
|
|
|
|
/// Tear down the loopback listener. Idempotent.
|
|
public func stop() {
|
|
linuxTCPAcceptSource?.cancel()
|
|
linuxTCPAcceptSource = nil
|
|
linuxTCPListenFD = nil
|
|
}
|
|
#endif
|
|
}
|
|
|
|
#if canImport(FoundationNetworking)
|
|
/// FoundationNetworking's delegate bridge for one streamed upstream request. Linux does not expose
|
|
/// `URLSession.bytes(for:)`; yielding `didReceive data` chunks keeps Claude's SSE/tool stream live.
|
|
private final class StreamingUpstreamTask: NSObject, URLSessionDataDelegate, @unchecked Sendable {
|
|
typealias Body = AsyncThrowingStream<Data, Error>
|
|
|
|
private let lock = NSLock()
|
|
private let bodyContinuation: Body.Continuation
|
|
let body: Body
|
|
private var responseContinuation: CheckedContinuation<HTTPURLResponse, Error>?
|
|
private var task: URLSessionDataTask?
|
|
private var session: URLSession?
|
|
private var finished = false
|
|
private var terminalError: Error?
|
|
|
|
override init() {
|
|
let pair = Body.makeStream()
|
|
body = pair.stream
|
|
bodyContinuation = pair.continuation
|
|
super.init()
|
|
bodyContinuation.onTermination = { [weak self] termination in
|
|
guard case .cancelled = termination else { return }
|
|
self?.cancel()
|
|
}
|
|
}
|
|
|
|
static func start(_ request: URLRequest) async throws -> (HTTPURLResponse, Body) {
|
|
let bridge = StreamingUpstreamTask()
|
|
let response = try await bridge.response(for: request)
|
|
return (response, bridge.body)
|
|
}
|
|
|
|
private func response(for request: URLRequest) async throws -> HTTPURLResponse {
|
|
try await withTaskCancellationHandler {
|
|
try await withCheckedThrowingContinuation { continuation in
|
|
let session = URLSession(
|
|
configuration: .ephemeral, delegate: self, delegateQueue: nil)
|
|
let task = session.dataTask(with: request)
|
|
|
|
lock.lock()
|
|
if finished {
|
|
let error = terminalError ?? CancellationError()
|
|
lock.unlock()
|
|
session.invalidateAndCancel()
|
|
continuation.resume(throwing: error)
|
|
return
|
|
}
|
|
self.responseContinuation = continuation
|
|
self.session = session
|
|
self.task = task
|
|
lock.unlock()
|
|
task.resume()
|
|
}
|
|
} onCancel: {
|
|
self.cancel()
|
|
}
|
|
}
|
|
|
|
func urlSession(
|
|
_ session: URLSession,
|
|
dataTask: URLSessionDataTask,
|
|
didReceive response: URLResponse,
|
|
completionHandler: @escaping @Sendable (URLSession.ResponseDisposition) -> Void
|
|
) {
|
|
guard let http = response as? HTTPURLResponse else {
|
|
completionHandler(.cancel)
|
|
complete(ClaudeTokenProxy.ProxyError.listenerFailed("non-HTTP upstream response"))
|
|
return
|
|
}
|
|
|
|
lock.lock()
|
|
guard !finished else {
|
|
lock.unlock()
|
|
completionHandler(.cancel)
|
|
return
|
|
}
|
|
let continuation = responseContinuation
|
|
responseContinuation = nil
|
|
lock.unlock()
|
|
continuation?.resume(returning: http)
|
|
completionHandler(.allow)
|
|
}
|
|
|
|
func urlSession(
|
|
_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data
|
|
) {
|
|
if !data.isEmpty { bodyContinuation.yield(data) }
|
|
}
|
|
|
|
func urlSession(
|
|
_ session: URLSession,
|
|
task: URLSessionTask,
|
|
didCompleteWithError error: (any Error)?
|
|
) {
|
|
complete(error)
|
|
}
|
|
|
|
private func cancel() {
|
|
lock.lock()
|
|
let task = self.task
|
|
lock.unlock()
|
|
complete(CancellationError())
|
|
task?.cancel()
|
|
}
|
|
|
|
private func complete(_ error: Error?) {
|
|
lock.lock()
|
|
guard !finished else { lock.unlock(); return }
|
|
finished = true
|
|
terminalError = error
|
|
let response = responseContinuation
|
|
responseContinuation = nil
|
|
let session = self.session
|
|
self.session = nil
|
|
task = nil
|
|
lock.unlock()
|
|
|
|
if let error {
|
|
response?.resume(throwing: error)
|
|
bodyContinuation.finish(throwing: error)
|
|
} else if let response {
|
|
let error = ClaudeTokenProxy.ProxyError.listenerFailed("upstream ended before response")
|
|
response.resume(throwing: error)
|
|
bodyContinuation.finish(throwing: error)
|
|
} else {
|
|
bodyContinuation.finish()
|
|
}
|
|
session?.finishTasksAndInvalidate()
|
|
}
|
|
}
|
|
#endif
|
|
|
|
#if canImport(Network)
|
|
/// A resume-once guard for the `NWListener` state handler (`.ready` may be followed by `.failed`).
|
|
private final class ProxyOnce: @unchecked Sendable {
|
|
private let lock = NSLock()
|
|
private var claimed = false
|
|
func claim() -> Bool {
|
|
lock.lock(); defer { lock.unlock() }
|
|
if claimed { return false }
|
|
claimed = true
|
|
return true
|
|
}
|
|
}
|
|
|
|
/// ``ByteReader`` + ``ByteWriter`` over a Network.framework connection (host loopback TCP).
|
|
private final class NWSocketConn: ByteReader, ByteWriter, @unchecked Sendable {
|
|
private let connection: NWConnection
|
|
init(_ connection: NWConnection) { self.connection = connection }
|
|
func read() async -> Data? {
|
|
await withCheckedContinuation { cont in
|
|
connection.receive(minimumIncompleteLength: 1, maximumLength: 64 * 1024) {
|
|
data, _, isComplete, error in
|
|
if let data, !data.isEmpty { cont.resume(returning: data) }
|
|
else if isComplete || error != nil { cont.resume(returning: nil) }
|
|
else { cont.resume(returning: Data()) }
|
|
}
|
|
}
|
|
}
|
|
func write(_ data: Data) async {
|
|
await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
|
|
connection.send(content: data, completion: .contentProcessed { _ in cont.resume() })
|
|
}
|
|
}
|
|
func close() async { connection.cancel() }
|
|
}
|
|
|
|
/// ``ByteReader`` + ``ByteWriter`` over an accepted `AF_UNIX` stream socket (the relayed container
|
|
/// transport). fd-lifetime discipline mirrors ``MCPApprovalServer``'s `UnixSocketByteConn`: `close`
|
|
/// only `shutdown(2)`s; the descriptor is released once in `deinit`, after the last kernel call.
|
|
private final class UnixSocketConn: ByteReader, ByteWriter, @unchecked Sendable {
|
|
private let fd: Int32
|
|
private let lock = NSLock()
|
|
private var closed = false
|
|
init(fd: Int32) {
|
|
self.fd = fd
|
|
var one: Int32 = 1
|
|
setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &one, socklen_t(MemoryLayout<Int32>.size))
|
|
}
|
|
deinit { Darwin.close(fd) }
|
|
private var isClosed: Bool { lock.lock(); defer { lock.unlock() }; return closed }
|
|
func read() async -> Data? {
|
|
guard !isClosed else { return nil }
|
|
return await withCheckedContinuation { cont in
|
|
DispatchQueue.global(qos: .userInitiated).async {
|
|
var buffer = [UInt8](repeating: 0, count: 64 * 1024)
|
|
while true {
|
|
let n = buffer.withUnsafeMutableBytes { Darwin.read(self.fd, $0.baseAddress, 64 * 1024) }
|
|
if n > 0 { cont.resume(returning: Data(buffer[0..<n])) }
|
|
else if n == 0 { cont.resume(returning: nil) }
|
|
else if errno == EINTR { continue }
|
|
else { cont.resume(returning: nil) }
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
func write(_ data: Data) async {
|
|
guard !isClosed else { return }
|
|
await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
|
|
DispatchQueue.global(qos: .userInitiated).async {
|
|
data.withUnsafeBytes { raw in
|
|
guard let base = raw.baseAddress else { return }
|
|
var offset = 0
|
|
while offset < raw.count {
|
|
let n = Darwin.send(self.fd, base + offset, raw.count - offset, Int32(MSG_NOSIGNAL_COMPAT))
|
|
if n > 0 { offset += n }
|
|
else if n < 0 && errno == EINTR { continue }
|
|
else { break }
|
|
}
|
|
}
|
|
cont.resume()
|
|
}
|
|
}
|
|
}
|
|
func close() async { shutdownOnce() }
|
|
|
|
/// Synchronous so it can hold the lock (NSLock is unavailable from an async context).
|
|
private func shutdownOnce() {
|
|
lock.lock(); defer { lock.unlock() }
|
|
guard !closed else { return }
|
|
closed = true
|
|
shutdown(fd, Int32(SHUT_RDWR))
|
|
}
|
|
}
|
|
|
|
// macOS has no MSG_NOSIGNAL; SO_NOSIGPIPE (set on the fd above) covers it, so send flags are 0.
|
|
private let MSG_NOSIGNAL_COMPAT: Int32 = 0
|
|
#endif
|
|
|
|
#if os(Linux)
|
|
/// ``ByteReader`` + ``ByteWriter`` over an accepted Linux loopback stream socket. Blocking
|
|
/// syscalls stay on a background queue; the descriptor is closed exactly once from `deinit`.
|
|
private final class LinuxSocketConn: ByteReader, ByteWriter, @unchecked Sendable {
|
|
private let fd: Int32
|
|
private let lock = NSLock()
|
|
private var closed = false
|
|
|
|
init(fd: Int32) { self.fd = fd }
|
|
deinit { Glibc.close(fd) }
|
|
|
|
private var isClosed: Bool {
|
|
lock.lock()
|
|
defer { lock.unlock() }
|
|
return closed
|
|
}
|
|
|
|
func read() async -> Data? {
|
|
guard !isClosed else { return nil }
|
|
return await withCheckedContinuation { continuation in
|
|
DispatchQueue.global(qos: .userInitiated).async {
|
|
var buffer = [UInt8](repeating: 0, count: 64 * 1024)
|
|
while true {
|
|
let count = buffer.withUnsafeMutableBytes {
|
|
Glibc.read(self.fd, $0.baseAddress, 64 * 1024)
|
|
}
|
|
if count > 0 {
|
|
continuation.resume(returning: Data(buffer[0..<count]))
|
|
} else if count == 0 {
|
|
continuation.resume(returning: nil)
|
|
} else if errno == EINTR {
|
|
continue
|
|
} else {
|
|
continuation.resume(returning: nil)
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func write(_ data: Data) async {
|
|
guard !isClosed else { return }
|
|
await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
|
|
DispatchQueue.global(qos: .userInitiated).async {
|
|
data.withUnsafeBytes { bytes in
|
|
guard let base = bytes.baseAddress else { return }
|
|
var offset = 0
|
|
while offset < bytes.count {
|
|
let count = Glibc.send(
|
|
self.fd, base + offset, bytes.count - offset, Int32(MSG_NOSIGNAL))
|
|
if count > 0 { offset += count }
|
|
else if count < 0 && errno == EINTR { continue }
|
|
else { break }
|
|
}
|
|
}
|
|
continuation.resume()
|
|
}
|
|
}
|
|
}
|
|
|
|
func close() async { shutdownOnce() }
|
|
|
|
private func shutdownOnce() {
|
|
lock.lock()
|
|
defer { lock.unlock() }
|
|
guard !closed else { return }
|
|
closed = true
|
|
_ = Glibc.shutdown(fd, Int32(SHUT_RDWR))
|
|
}
|
|
}
|
|
#endif
|
|
|
|
// MARK: - Byte-stream seams (in-memory in tests; socket-backed in the listeners)
|
|
|
|
/// A one-directional byte source. `read()` returns the next available bytes, or nil at EOF.
|
|
protocol ByteReader: Sendable {
|
|
func read() async -> Data?
|
|
}
|
|
|
|
/// A one-directional byte sink.
|
|
protocol ByteWriter: Sendable {
|
|
func write(_ data: Data) async
|
|
func close() async
|
|
}
|