| 1 | import Foundation |
| 2 | |
| 3 | // The live half of the transport: a WebSocket to /ws?token=…, receiving the |
| 4 | // hub's fan-out frames (hub.go wsOut and friends) and sending "msg" frames |
| 5 | // back. Everything on the wire is ciphertext blobs + routing metadata — |
| 6 | // encryption/decryption happens in whoever uses the socket (ChatFeedModel), |
| 7 | // never here. |
| 8 | public final class ChatSocket: NSObject, URLSessionWebSocketDelegate { |
| 9 | // One hub frame, decoded leniently: the hub sends several shapes (msg / |
| 10 | // edit / delete / typing) that share `type` + `thread_id`; fields |
| 11 | // the shape doesn't carry just stay nil. Reactions are now encrypted events |
| 12 | // (type="reaction") inside "msg" frames. dm_closed / dm_reopened (#578) |
| 13 | // carry only type + thread_id and need no fields of their own. |
| 14 | public struct Frame: Decodable { |
| 15 | public let type: String |
| 16 | public let thread_id: String? |
| 17 | public let id: Int64? // message id (msg/edit; the deleted id for delete) |
| 18 | public let created_at: Int64? // seconds |
| 19 | public let blob: String? |
| 20 | public let edited_at: Int64? |
| 21 | public let ephemeral: Bool? // relay-only traffic (presence) — never stored |
| 22 | public let user_id: String? // "presence" / "label" frames: whose state changed |
| 23 | public let presence: String? // "presence" frames: "online" | "away" | "offline" |
| 24 | public let last_read_at: Int64? // "read" frames: our watermark, moved on another device |
| 25 | public let kind: String? // "apps_changed" frames: "created" | "updated" | "removed" |
| 26 | // The wire text this frame decoded from, attached after decode (never |
| 27 | // a server field — CodingKeys excludes it): debug surfaces show what |
| 28 | // actually arrived, fields the lenient decode ignores included. |
| 29 | public var raw: String? = nil |
| 30 | |
| 31 | private enum CodingKeys: String, CodingKey { |
| 32 | case type, thread_id, id, created_at, blob, edited_at, ephemeral, |
| 33 | user_id, presence, last_read_at, kind |
| 34 | } |
| 35 | } |
| 36 | |
| 37 | /// Delivered on the main thread. |
| 38 | public var onFrame: ((Frame) -> Void)? |
| 39 | /// Connection state changes ("connected" / "reconnecting…"), main thread. |
| 40 | public var onStatus: ((String) -> Void)? |
| 41 | |
| 42 | /// Whether the app is frontmost — set once per platform at startup (iOS: |
| 43 | /// applicationState, Mac: NSApp.isActive). While true, every open socket |
| 44 | /// heartbeats an "active" frame (the same beat the web sends every 30s |
| 45 | /// while visible): the server uses it to suppress push to a person who's |
| 46 | /// reading, and to resolve online-vs-idle presence (presence.go). Read on |
| 47 | /// the main thread (the ping timer's), like everything else here. |
| 48 | public static var appIsActive: () -> Bool = { true } |
| 49 | |
| 50 | private let url: URL |
| 51 | private let origin: String |
| 52 | private var task: URLSessionWebSocketTask? |
| 53 | private var pingTimer: Timer? |
| 54 | // Own session, so the delegate's didOpen callback can report "connected" |
| 55 | // when the handshake actually completes (resume() alone proves nothing). |
| 56 | // The session retains its delegate — stop() invalidates it to break the |
| 57 | // cycle, and start() mints a fresh one. |
| 58 | private var session: URLSession? |
| 59 | // Bumped on every start()/stop(); callbacks from a previous generation's |
| 60 | // task (its error, its reconnect timer) see a stale value and stand down. |
| 61 | private var generation = 0 |
| 62 | // The connect→handshake signpost, open between connect() and didOpen. |
| 63 | // Main-thread state, like everything else here — the delegate queue is |
| 64 | // background (frame decode), so delegate callbacks hop to main first. |
| 65 | private var connectSpan: Perf.Span? |
| 66 | |
| 67 | // baseURL is the instance origin (https://host[:port]); the socket rides |
| 68 | // ws(s)://host/ws?token=…. The Origin header must be set explicitly — the |
| 69 | // hub's Accept applies the browser same-origin check, and a native client |
| 70 | // sends no Origin by default. |
| 71 | public init?(baseURL: URL, token: String) { |
| 72 | var comps = URLComponents(url: baseURL, resolvingAgainstBaseURL: false)! |
| 73 | comps.scheme = baseURL.scheme == "https" ? "wss" : "ws" |
| 74 | comps.path = "/ws" |
| 75 | comps.queryItems = [URLQueryItem(name: "token", value: token)] |
| 76 | guard let url = comps.url else { return nil } |
| 77 | self.url = url |
| 78 | var o = "\(baseURL.scheme ?? "https")://\(baseURL.host ?? "")" |
| 79 | if let port = baseURL.port { o += ":\(port)" } |
| 80 | self.origin = o |
| 81 | super.init() |
| 82 | } |
| 83 | |
| 84 | public func start() { |
| 85 | generation += 1 |
| 86 | connect(gen: generation) |
| 87 | } |
| 88 | |
| 89 | public func stop() { |
| 90 | generation += 1 |
| 91 | connectSpan?.end(record: false) // mid-handshake stop isn't a slow connect |
| 92 | connectSpan = nil |
| 93 | pingTimer?.invalidate() |
| 94 | pingTimer = nil |
| 95 | task?.cancel(with: .goingAway, reason: nil) |
| 96 | task = nil |
| 97 | session?.invalidateAndCancel() // breaks the session→delegate retain |
| 98 | session = nil |
| 99 | } |
| 100 | |
| 101 | /// Whether a live task is up — the send paths check this synchronously |
| 102 | /// (keep-the-typed-text-on-a-down-socket), then encrypt + send off main. |
| 103 | public var isLive: Bool { task?.state == .running } |
| 104 | |
| 105 | /// Send one already-encoded JSON frame. Returns false when the socket has |
| 106 | /// no live task — the caller keeps what it was sending (the web's emit |
| 107 | /// resolves false on a down socket for the same reason: never silently |
| 108 | /// drop what was typed). Thread-safe: the underlying WebSocket task |
| 109 | /// accepts sends from any thread, so the off-main crypto lane calls this |
| 110 | /// directly. |
| 111 | @discardableResult |
| 112 | public func send(_ jsonText: String) -> Bool { |
| 113 | guard let task, task.state == .running else { return false } |
| 114 | task.send(.string(jsonText)) { _ in } |
| 115 | return true |
| 116 | } |
| 117 | |
| 118 | deinit { stop() } |
| 119 | |
| 120 | // "connected" arrives via the delegate when the handshake actually |
| 121 | // completes — resuming the task only means we've started trying. The |
| 122 | // delegate queue is a background queue now, so hop to main before touching |
| 123 | // main-owned state (`task`, the span) or the consumer's callback. |
| 124 | public func urlSession(_ session: URLSession, webSocketTask: URLSessionWebSocketTask, |
| 125 | didOpenWithProtocol protocol: String?) { |
| 126 | DispatchQueue.main.async { [weak self] in |
| 127 | guard let self, webSocketTask === self.task else { return } |
| 128 | self.connectSpan?.end() |
| 129 | self.connectSpan = nil |
| 130 | self.onStatus?("connected") |
| 131 | } |
| 132 | } |
| 133 | |
| 134 | // Frames land (and their JSON decodes) here, off main — every socket on the |
| 135 | // instance receives EVERY thread's frames, so a busy group otherwise decodes |
| 136 | // the same burst on the main thread once per open/parked model (#465). |
| 137 | private let receiveQueue: OperationQueue = { |
| 138 | let q = OperationQueue() |
| 139 | q.maxConcurrentOperationCount = 1 |
| 140 | q.name = "chat-socket-receive" |
| 141 | return q |
| 142 | }() |
| 143 | |
| 144 | private func connect(gen: Int) { |
| 145 | guard gen == generation else { return } |
| 146 | // A dangling span from a failed attempt closes unrecorded — only a |
| 147 | // completed handshake is an honest connect duration. |
| 148 | connectSpan?.end(record: false) |
| 149 | connectSpan = Perf.begin("socket.connect", slowMs: 3000) |
| 150 | var req = URLRequest(url: url) |
| 151 | req.setValue(origin, forHTTPHeaderField: "Origin") |
| 152 | if session == nil { |
| 153 | session = URLSession(configuration: .default, delegate: self, delegateQueue: receiveQueue) |
| 154 | } |
| 155 | let task = session!.webSocketTask(with: req) |
| 156 | self.task = task |
| 157 | task.resume() |
| 158 | receive(task, gen: gen) |
| 159 | |
| 160 | // Protocol-level pings keep NATs/proxies from reaping an idle socket |
| 161 | // (the server's websocket library answers them without surfacing a frame). |
| 162 | // While the app is frontmost, ride the same beat with an "active" frame — |
| 163 | // presence + push suppression (see appIsActive above). |
| 164 | pingTimer?.invalidate() |
| 165 | pingTimer = Timer.scheduledTimer(withTimeInterval: 25, repeats: true) { [weak self] _ in |
| 166 | self?.task?.sendPing { _ in } |
| 167 | // No beat needed on connect itself — the server counts connecting |
| 168 | // as activity (hub.go handleWS). |
| 169 | if Self.appIsActive() { self?.send(#"{"type":"active"}"#) } |
| 170 | } |
| 171 | } |
| 172 | |
| 173 | private func receive(_ task: URLSessionWebSocketTask, gen: Int) { |
| 174 | task.receive { [weak self] result in |
| 175 | // Decode on the delegate queue (background, serial — order kept); |
| 176 | // only the generation check + delivery hop to main. |
| 177 | var frame: Frame? |
| 178 | if case .success(.string(let text)) = result { |
| 179 | frame = try? JSONDecoder().decode(Frame.self, from: Data(text.utf8)) |
| 180 | frame?.raw = text |
| 181 | } |
| 182 | let failed: Bool |
| 183 | if case .failure = result { failed = true } else { failed = false } |
| 184 | DispatchQueue.main.async { |
| 185 | guard let self, gen == self.generation else { return } |
| 186 | if failed { |
| 187 | // Dropped (sleep, network blip, server restart) — reconnect and |
| 188 | // let the consumer refetch what it missed. |
| 189 | Perf.event("socket.reconnect") |
| 190 | self.onStatus?("reconnecting…") |
| 191 | DispatchQueue.main.asyncAfter(deadline: .now() + 3) { [weak self] in |
| 192 | self?.connect(gen: gen) |
| 193 | } |
| 194 | return |
| 195 | } |
| 196 | if let frame { self.onFrame?(frame) } |
| 197 | self.receive(task, gen: gen) |
| 198 | } |
| 199 | } |
| 200 | } |
| 201 | } |
| 202 | |