web-terminal/ios/App/WebTermTests/AccessTokenSourceTests.swift

import Foundation
import HostRegistry
import SessionCore
import TestSupport
import Testing
import WireProtocol
@testable import WebTerm

/// C1 · The app's single read path for per-host access tokens, plus the one
/// header-stamping point (`AccessTokenCookie`).
///
/// E1 · The WS cases are the important ones. SessionCore's `tokenProvider` is
/// `() -> String?` — no endpoint, no await — so C1 answered it with "the token
/// every paired host shares", which is *nil* for the deployment the token
/// exists for (a tokened tunnel host next to an open LAN host): the upgrade then
/// omitted the cookie, the server 401'd, and `.failed(.unauthorized)` is
/// terminal — with no in-app remedy, since re-entering the correct token left
/// the fleet just as mixed. The answer is now resolved PER HOST
/// (`wsToken(for:)`), exactly like the HTTP path and like Android's
/// `tokens.tokenFor(endpoint)`.
@Suite("Access-token source")
struct AccessTokenSourceTests {
    private static let tokenA = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
    private static let tokenB = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"

    private struct StoreFailure: Error {}

    private actor ThrowingHostStore: HostStore {
        func loadAll() async throws -> [HostRegistry.Host] { throw StoreFailure() }
        func upsert(_ host: HostRegistry.Host) async throws -> [HostRegistry.Host] {
            throw StoreFailure()
        }
        func remove(id: UUID) async throws -> [HostRegistry.Host] { throw StoreFailure() }
    }

    private func makeHost(_ base: String, token: String?) throws -> HostRegistry.Host {
        let url = try #require(URL(string: base))
        return HostRegistry.Host(
            id: UUID(),
            name: base,
            endpoint: try #require(HostEndpoint(baseURL: url)),
            accessToken: try token.map { try AccessToken(validating: $0) }
        )
    }

    private func makeStore(_ hosts: [HostRegistry.Host]) async throws -> InMemoryHostStore {
        let store = InMemoryHostStore()
        for host in hosts {
            _ = try await store.upsert(host)
        }
        return store
    }

    // MARK: - Per-origin resolution (the HTTP path)

    @Test("按 origin 取到该主机的令牌；其它 origin 与无令牌主机都返回 nil")
    func resolvesTokenByOrigin() async throws {
        let store = try await makeStore([
            try makeHost("http://192.168.1.5:3000", token: Self.tokenA),
            try makeHost("http://192.168.1.9:3000", token: nil),
        ])
        let source = AccessTokenSource(store: store)

        #expect(await source.token(forOrigin: "http://192.168.1.5:3000") == Self.tokenA)
        #expect(await source.token(forOrigin: "http://192.168.1.9:3000") == nil)
        #expect(await source.token(forOrigin: "http://192.168.1.99:3000") == nil)
    }

    @Test("https 默认端口按 originHeader 规范化匹配（不写 :443）")
    func matchesNormalizedHTTPSOrigin() async throws {
        let store = try await makeStore([
            try makeHost("https://mac.tailnet.ts.net", token: Self.tokenA),
        ])
        let source = AccessTokenSource(store: store)

        #expect(await source.token(forOrigin: "https://mac.tailnet.ts.net") == Self.tokenA)
        #expect(await source.token(forOrigin: "https://mac.tailnet.ts.net:443") == nil)
    }

    @Test("存储读取失败 → nil（降级成「无令牌」，不让 App 崩或卡住）")
    func storeFailureDegradesToNoToken() async throws {
        let source = AccessTokenSource(store: ThrowingHostStore())

        #expect(await source.token(forOrigin: "http://192.168.1.5:3000") == nil)
    }

    // MARK: - The per-host WS answer (E1 · replaces `hostIndependentToken`)

    @Test("混合机群：有令牌的那台拿到自己的令牌，没令牌的那台拿 nil（旧解析对两台都给 nil）")
    func resolvesPerHostTokenInAMixedFleet() async throws {
        let tokened = try makeHost("http://192.168.1.5:3000", token: Self.tokenA)
        let open = try makeHost("http://192.168.1.9:3000", token: nil)
        let source = AccessTokenSource(store: try await makeStore([tokened, open]))

        _ = await source.token(forOrigin: tokened.endpoint.originHeader) // 快照热身

        #expect(source.wsToken(for: tokened) == Self.tokenA)
        #expect(source.wsToken(for: open) == nil)
    }

    @Test("两台令牌不同 → 各拿各的，绝不把 A 的令牌发给 B")
    func neverHandsOneHostsTokenToAnother() async throws {
        let hostA = try makeHost("http://192.168.1.5:3000", token: Self.tokenA)
        let hostB = try makeHost("http://192.168.1.9:3000", token: Self.tokenB)
        let source = AccessTokenSource(store: try await makeStore([hostA, hostB]))

        _ = await source.token(forOrigin: hostA.endpoint.originHeader)

        #expect(source.wsToken(for: hostA) == Self.tokenA)
        #expect(source.wsToken(for: hostB) == Self.tokenB)
    }

    @Test("快照还没热（刚启动就开终端）→ 回落到该主机记录，绝不因竞态漏掉 Cookie")
    func fallsBackToTheHostRecordBeforeAnyStoreRead() async throws {
        let host = try makeHost("http://192.168.1.5:3000", token: Self.tokenA)
        let source = AccessTokenSource(store: try await makeStore([host]))

        // No `token(forOrigin:)` call yet — the snapshot is empty.
        #expect(source.wsToken(for: host) == Self.tokenA)
    }

    @Test("令牌轮换后：下一次连接用存储里的新值，而不是打开终端时的旧值")
    func picksUpARotatedTokenOnTheNextConnect() async throws {
        let opened = try makeHost("http://192.168.1.5:3000", token: Self.tokenA)
        let store = try await makeStore([opened])
        let source = AccessTokenSource(store: store)
        let rotated = HostRegistry.Host(
            id: opened.id, name: opened.name, endpoint: opened.endpoint,
            accessToken: try AccessToken(validating: Self.tokenB)
        )
        _ = try await store.upsert(rotated)

        _ = await source.token(forOrigin: opened.endpoint.originHeader) // 任一 HTTP 请求即刷新

        #expect(source.wsToken(for: opened) == Self.tokenB)
    }

    @Test("存储读取失败 → 回落到主机记录（打开时的值），不是别的主机的令牌")
    func wsTokenSurvivesAStoreFailure() async throws {
        let host = try makeHost("http://192.168.1.5:3000", token: Self.tokenA)
        let source = AccessTokenSource(store: ThrowingHostStore())

        #expect(await source.token(forOrigin: host.endpoint.originHeader) == nil)
        #expect(source.wsToken(for: host) == Self.tokenA)
    }

    @Test("完全没有令牌的机群 → nil（LAN 零配置逐字不变，不带 Cookie）")
    func tokenlessFleetYieldsNoCookie() async throws {
        let host = try makeHost("http://192.168.1.5:3000", token: nil)
        let source = AccessTokenSource(store: try await makeStore([host]))

        _ = await source.token(forOrigin: host.endpoint.originHeader)

        #expect(source.wsToken(for: host) == nil)
    }
}

/// E1 · The wiring that the per-host answer depends on: a terminal's WS
/// transport is built for THE host it dials, not once for the whole app.
@MainActor
@Suite("Per-host WS transport wiring")
struct PerHostTermTransportTests {
    /// `@Sendable`-safe recorder for the hosts the factory was asked about.
    private final class HostRecorder: @unchecked Sendable {
        private let lock = NSLock()
        private var hosts: [HostRegistry.Host] = []

        func record(_ host: HostRegistry.Host) {
            lock.withLock { hosts.append(host) }
        }

        var recorded: [HostRegistry.Host] { lock.withLock { hosts } }
    }

    private func makeHost(_ base: String, token: String?) throws -> HostRegistry.Host {
        let url = try #require(URL(string: base))
        return HostRegistry.Host(
            id: UUID(), name: base,
            endpoint: try #require(HostEndpoint(baseURL: url)),
            accessToken: try token.map { try AccessToken(validating: $0) }
        )
    }

    private func makeEnvironment(
        shared: FakeTransport,
        factory: (@Sendable (HostRegistry.Host) -> any TermTransport)? = nil
    ) throws -> AppEnvironment {
        let defaults = try #require(UserDefaults(suiteName: "PerHostTermTransportTests"))
        return AppEnvironment(
            hostStore: InMemoryHostStore(),
            lastSessionStore: UserDefaultsLastSessionStore(defaults: defaults),
            http: FakeHTTPTransport(),
            termTransport: shared,
            probe: PairingViewModel.Probe(verifyHost: { _, _ in .failure(.timeout) }),
            termTransportFactory: factory
        )
    }

    @Test("未注入工厂 → 回落到共享传输（既有 App 层测试装配零回归）")
    func fallsBackToTheSharedTransport() throws {
        let shared = FakeTransport()
        let environment = try makeEnvironment(shared: shared)

        let resolved = environment.makeTermTransport(for: try makeHost("http://192.168.1.5:3000", token: nil))

        #expect(resolved as? FakeTransport === shared)
    }

    @Test("注入工厂 → 每台主机各建一个传输，且工厂拿到的就是那台主机（含其令牌）")
    func buildsOneTransportPerHost() throws {
        let recorder = HostRecorder()
        let environment = try makeEnvironment(
            shared: FakeTransport(),
            factory: { host in
                recorder.record(host)
                return FakeTransport()
            }
        )
        let tokened = try makeHost("http://192.168.1.5:3000", token: String(repeating: "a", count: 32))
        let open = try makeHost("http://192.168.1.9:3000", token: nil)

        _ = environment.makeTermTransport(for: tokened)
        _ = environment.makeTermTransport(for: open)

        #expect(recorder.recorded.map(\.id) == [tokened.id, open.id])
        #expect(recorder.recorded.first?.accessToken == tokened.accessToken)
    }

    @Test("TerminalSessionController 用自己那台主机建传输（终端 401 的根因就在这里）")
    func controllerBuildsItsTransportForItsOwnHost() throws {
        let recorder = HostRecorder()
        let environment = try makeEnvironment(
            shared: FakeTransport(),
            factory: { host in
                recorder.record(host)
                return FakeTransport()
            }
        )
        let host = try makeHost("http://192.168.1.5:3000", token: String(repeating: "b", count: 32))

        _ = TerminalSessionController(
            host: host, sessionId: nil, environment: environment,
            onPendingChanged: { _, _ in }
        )

        #expect(recorder.recorded.map(\.id) == [host.id])
    }
}

/// C1 · The App layer's one place where a token becomes a header.
@Suite("Access-token cookie stamping")
struct AccessTokenCookieTests {
    private static let token = "abcdefghijklmnopqrstuvwxyz012345"

    private func request(_ url: String) throws -> URLRequest {
        URLRequest(url: try #require(URL(string: url)))
    }

    @Test("请求 URL → originHeader 形式的 origin（路径/查询串被忽略）")
    func derivesOriginFromRequestURL() throws {
        let withPath = try request("http://192.168.1.5:3000/live-sessions/abc/preview?x=1")

        #expect(AccessTokenCookie.origin(of: withPath) == "http://192.168.1.5:3000")
    }

    @Test("https 默认端口不写 :443（与服务器的 URL 规范化一致）")
    func omitsDefaultHTTPSPort() throws {
        let secure = try request("https://mac.tailnet.ts.net/live-sessions")

        #expect(AccessTokenCookie.origin(of: secure) == "https://mac.tailnet.ts.net")
    }

    @Test("非 http(s) / 无 host 的请求 → nil（不可能是本 App 的主机）")
    func rejectsNonHTTPRequests() throws {
        #expect(AccessTokenCookie.origin(of: try request("ftp://192.168.1.5/x")) == nil)
        #expect(AccessTokenCookie.origin(of: URLRequest(url: URL(fileURLWithPath: "/tmp"))) == nil)
    }

    @Test("盖上 Cookie: webterm_auth=<t>，且返回新请求（不原地改调用方的请求）")
    func stampsCookieImmutably() throws {
        let original = try request("http://192.168.1.5:3000/live-sessions")

        let stamped = AccessTokenCookie.stamped(original, token: Self.token)

        #expect(stamped.value(forHTTPHeaderField: "Cookie") == "webterm_auth=\(Self.token)")
        #expect(original.value(forHTTPHeaderField: "Cookie") == nil)
    }

    @Test("已经带 Cookie 的请求原样通过（配对探针的候选令牌不能被覆盖）")
    func neverOverwritesAnExistingCookie() throws {
        var candidate = try request("http://192.168.1.5:3000/live-sessions")
        candidate.setValue("webterm_auth=candidate-token-value", forHTTPHeaderField: "Cookie")

        let stamped = AccessTokenCookie.stamped(candidate, token: Self.token)

        #expect(stamped.value(forHTTPHeaderField: "Cookie") == "webterm_auth=candidate-token-value")
    }

    @Test("Cookie 名与服务器 AUTH_COOKIE_NAME 逐字一致")
    func cookieNameMatchesTheServer() {
        #expect(AccessTokenCookie.name == "webterm_auth")
    }

    // MARK: - The transport choke point itself (no network involved)

    @Test("传输层给每个请求按 origin 盖上令牌 Cookie（RO GET 也一样）")
    func transportStampsEveryRequest() async throws {
        let transport = URLSessionHTTPTransport(
            identityProvider: { nil },
            tokenForOrigin: { origin in
                origin == "http://192.168.1.5:3000" ? Self.token : nil
            }
        )
        let readOnly = try request("http://192.168.1.5:3000/live-sessions")

        let stamped = await transport.authenticated(readOnly)

        #expect(stamped.value(forHTTPHeaderField: "Cookie") == "webterm_auth=\(Self.token)")
    }

    @Test("没有该 origin 的令牌 → 请求逐字不变（LAN 零配置主机不受影响）")
    func transportLeavesTokenlessOriginsAlone() async throws {
        let transport = URLSessionHTTPTransport(
            identityProvider: { nil }, tokenForOrigin: { _ in nil }
        )
        let plain = try request("http://192.168.1.9:3000/live-sessions")

        let result = await transport.authenticated(plain)

        #expect(result.value(forHTTPHeaderField: "Cookie") == nil)
        #expect(result.allHTTPHeaderFields == plain.allHTTPHeaderFields)
    }

    @Test("请求已带 Cookie（配对探针的候选令牌）→ 传输层不覆盖")
    func transportNeverOverwritesTheProbeCandidate() async throws {
        let transport = URLSessionHTTPTransport(
            identityProvider: { nil }, tokenForOrigin: { _ in Self.token }
        )
        var candidate = try request("http://192.168.1.5:3000/live-sessions")
        candidate.setValue("webterm_auth=candidate", forHTTPHeaderField: "Cookie")

        let result = await transport.authenticated(candidate)

        #expect(result.value(forHTTPHeaderField: "Cookie") == "webterm_auth=candidate")
    }
}