288 lines
9.5 KiB
Swift
288 lines
9.5 KiB
Swift
import Foundation
|
|
import Network
|
|
|
|
enum LoopbackStartResult {
|
|
case success(UInt16)
|
|
case failure
|
|
}
|
|
|
|
/**
|
|
* One-shot HTTP listener bound to 127.0.0.1 that hands the provider's redirect to the app scheme.
|
|
*
|
|
* Security properties:
|
|
* - Loopback only (`requiredLocalEndpoint`), so nothing on the local network can reach it.
|
|
* - Only a well-formed `GET /auth/callback` with a matching Host reaches the redirect step; any
|
|
* other request is answered with an error and does not consume the session, so the real callback
|
|
* can still arrive.
|
|
* - The response is a 302 to `usage://oauth/<provider>`; the code is passed through as-is with
|
|
* URLComponents, never interpolated into a header string or an HTML page.
|
|
* - First accepted callback ends the session: the listener stops and later requests are refused.
|
|
*/
|
|
final class LoopbackServer {
|
|
private static let requestHeadTerminator = Data("\r\n\r\n".utf8)
|
|
private static let maxRequestBytes = 16 * 1024
|
|
private static let readyTimeoutSeconds: TimeInterval = 3
|
|
private static let connectionIdleSeconds: TimeInterval = 20
|
|
private static let maximumLifetimeSeconds: TimeInterval = 1800
|
|
|
|
private let queue = DispatchQueue(label: "com.nextroad-dev.aiusage.oauth-loopback.server")
|
|
private let provider: String
|
|
private var listener: NWListener?
|
|
private var connections: [ObjectIdentifier: NWConnection] = [:]
|
|
private var boundPort: UInt16 = 0
|
|
private var finished = false
|
|
private var reported = false
|
|
private var lifetime: DispatchWorkItem?
|
|
|
|
init(provider: String) {
|
|
self.provider = provider
|
|
}
|
|
|
|
func start(port: UInt16, completion: @escaping (LoopbackStartResult) -> Void) {
|
|
queue.async {
|
|
let parameters = NWParameters.tcp
|
|
parameters.allowLocalEndpointReuse = true
|
|
// Restrict the bind to 127.0.0.1: the browser on this device reaches it, the network cannot.
|
|
parameters.requiredLocalEndpoint = NWEndpoint.hostPort(
|
|
host: .ipv4(.loopback),
|
|
port: NWEndpoint.Port(rawValue: port) ?? .any
|
|
)
|
|
|
|
let listener: NWListener
|
|
do {
|
|
listener = try NWListener(using: parameters)
|
|
} catch {
|
|
completion(.failure)
|
|
return
|
|
}
|
|
|
|
self.listener = listener
|
|
listener.newConnectionHandler = { [weak self] connection in
|
|
self?.accept(connection)
|
|
}
|
|
listener.stateUpdateHandler = { [weak self] state in
|
|
guard let self, !self.reported, !self.finished else { return }
|
|
switch state {
|
|
case .ready:
|
|
self.reported = true
|
|
self.boundPort = listener.port?.rawValue ?? port
|
|
self.scheduleLifetime()
|
|
completion(.success(self.boundPort))
|
|
case .failed:
|
|
self.reported = true
|
|
self.finishFlow()
|
|
completion(.failure)
|
|
default:
|
|
break
|
|
}
|
|
}
|
|
listener.start(queue: self.queue)
|
|
|
|
// A listener can also stay in `.waiting` when the address is not usable; never hang on it.
|
|
self.queue.asyncAfter(deadline: .now() + Self.readyTimeoutSeconds) { [weak self] in
|
|
guard let self, !self.reported, !self.finished else { return }
|
|
self.reported = true
|
|
self.finishFlow()
|
|
completion(.failure)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Stops accepting and releases every open connection. Safe to call more than once.
|
|
func stop() {
|
|
queue.async {
|
|
self.finishFlow()
|
|
for connection in self.connections.values {
|
|
connection.cancel()
|
|
}
|
|
self.connections.removeAll()
|
|
}
|
|
}
|
|
|
|
// MARK: - Session lifecycle
|
|
|
|
private func finishFlow() {
|
|
guard !finished else { return }
|
|
finished = true
|
|
lifetime?.cancel()
|
|
lifetime = nil
|
|
listener?.cancel()
|
|
listener = nil
|
|
}
|
|
|
|
private func scheduleLifetime() {
|
|
let item = DispatchWorkItem { [weak self] in self?.stop() }
|
|
lifetime = item
|
|
queue.asyncAfter(deadline: .now() + Self.maximumLifetimeSeconds, execute: item)
|
|
}
|
|
|
|
private func accept(_ connection: NWConnection) {
|
|
guard !finished else {
|
|
connection.cancel()
|
|
return
|
|
}
|
|
connections[ObjectIdentifier(connection)] = connection
|
|
connection.stateUpdateHandler = { [weak self] state in
|
|
switch state {
|
|
case .failed, .cancelled:
|
|
self?.release(connection)
|
|
default:
|
|
break
|
|
}
|
|
}
|
|
connection.start(queue: queue)
|
|
read(connection, buffer: Data())
|
|
queue.asyncAfter(deadline: .now() + Self.connectionIdleSeconds) { [weak self] in
|
|
self?.release(connection)
|
|
}
|
|
}
|
|
|
|
private func release(_ connection: NWConnection) {
|
|
connections.removeValue(forKey: ObjectIdentifier(connection))
|
|
connection.cancel()
|
|
}
|
|
|
|
// MARK: - HTTP
|
|
|
|
private func read(_ connection: NWConnection, buffer: Data) {
|
|
connection.receive(minimumIncompleteLength: 1, maximumLength: Self.maxRequestBytes) {
|
|
[weak self] data, _, isComplete, error in
|
|
guard let self else { return }
|
|
var accumulated = buffer
|
|
if let data, !data.isEmpty {
|
|
accumulated.append(data)
|
|
}
|
|
if let terminator = accumulated.range(of: Self.requestHeadTerminator) {
|
|
self.handle(connection, head: accumulated.subdata(in: 0..<terminator.lowerBound))
|
|
return
|
|
}
|
|
if accumulated.count > Self.maxRequestBytes || isComplete || error != nil {
|
|
self.respond(connection, status: 400, location: nil)
|
|
return
|
|
}
|
|
self.read(connection, buffer: accumulated)
|
|
}
|
|
}
|
|
|
|
private func handle(_ connection: NWConnection, head: Data) {
|
|
guard let text = String(data: head, encoding: .utf8) else {
|
|
respond(connection, status: 400, location: nil)
|
|
return
|
|
}
|
|
let lines = text.components(separatedBy: "\r\n")
|
|
let requestLine = (lines.first ?? "").components(separatedBy: " ").filter { !$0.isEmpty }
|
|
guard requestLine.count >= 2, requestLine[0] == "GET" else {
|
|
respond(connection, status: 405, location: nil)
|
|
return
|
|
}
|
|
guard let target = URLComponents(string: requestLine[1]), target.path == "/auth/callback" else {
|
|
respond(connection, status: 404, location: nil)
|
|
return
|
|
}
|
|
guard hostMatchesListener(lines) else {
|
|
respond(connection, status: 400, location: nil)
|
|
return
|
|
}
|
|
|
|
let items = target.queryItems ?? []
|
|
let code = singleValue(items, "code")
|
|
let state = singleValue(items, "state")
|
|
let error = singleValue(items, "error")
|
|
// Anything without a complete code/state signal is ignored rather than consumed: the real
|
|
// callback must still be able to finish the sign-in.
|
|
guard error != nil || (code != nil && state != nil) else {
|
|
respond(connection, status: 400, location: nil)
|
|
return
|
|
}
|
|
|
|
var redirect = URLComponents()
|
|
redirect.scheme = "usage"
|
|
redirect.host = "oauth"
|
|
redirect.path = "/" + provider
|
|
var query: [URLQueryItem] = []
|
|
if let code {
|
|
query.append(URLQueryItem(name: "code", value: code))
|
|
}
|
|
if let state {
|
|
query.append(URLQueryItem(name: "state", value: state))
|
|
}
|
|
if let error {
|
|
query.append(URLQueryItem(name: "error", value: error))
|
|
if let description = singleValue(items, "error_description") {
|
|
query.append(URLQueryItem(name: "error_description", value: description))
|
|
}
|
|
}
|
|
redirect.queryItems = query
|
|
guard let location = redirect.string else {
|
|
respond(connection, status: 400, location: nil)
|
|
return
|
|
}
|
|
|
|
finishFlow()
|
|
respond(connection, status: 302, location: location)
|
|
}
|
|
|
|
private func hostMatchesListener(_ lines: [String]) -> Bool {
|
|
let prefix = "host:"
|
|
guard
|
|
let header = lines.dropFirst().first(where: { $0.lowercased().hasPrefix(prefix) }),
|
|
boundPort != 0
|
|
else {
|
|
return false
|
|
}
|
|
let value = header.dropFirst(prefix.count).trimmingCharacters(in: .whitespaces).lowercased()
|
|
return value == "127.0.0.1:\(boundPort)" || value == "localhost:\(boundPort)"
|
|
}
|
|
|
|
/// Exactly one value for the parameter; repeated parameters are treated as invalid.
|
|
private func singleValue(_ items: [URLQueryItem], _ name: String) -> String? {
|
|
let values = items.filter { $0.name == name }.map { $0.value ?? "" }
|
|
return values.count == 1 ? values[0] : nil
|
|
}
|
|
|
|
private func respond(_ connection: NWConnection, status: Int, location: String?) {
|
|
var head = "HTTP/1.1 \(status) \(statusText(status))\r\n"
|
|
head += "Cache-Control: no-store\r\n"
|
|
head += "Connection: close\r\n"
|
|
if let location {
|
|
head += "Location: \(location)\r\n"
|
|
head += "Content-Length: 0\r\n\r\n"
|
|
send(connection, Data(head.utf8))
|
|
return
|
|
}
|
|
let body = "Invalid callback request."
|
|
head += "Content-Type: text/plain; charset=utf-8\r\n"
|
|
head += "Content-Length: \(body.utf8.count)\r\n\r\n"
|
|
var payload = Data(head.utf8)
|
|
payload.append(Data(body.utf8))
|
|
send(connection, payload)
|
|
}
|
|
|
|
private func send(_ connection: NWConnection, _ payload: Data) {
|
|
connection.send(content: payload, completion: .contentProcessed { _ in })
|
|
// Keep reading until the peer closes so the response is not reset by unread request data.
|
|
drain(connection)
|
|
}
|
|
|
|
private func drain(_ connection: NWConnection) {
|
|
connection.receive(minimumIncompleteLength: 1, maximumLength: Self.maxRequestBytes) {
|
|
[weak self] _, _, isComplete, error in
|
|
guard let self else { return }
|
|
if isComplete || error != nil {
|
|
self.release(connection)
|
|
return
|
|
}
|
|
self.drain(connection)
|
|
}
|
|
}
|
|
|
|
private func statusText(_ status: Int) -> String {
|
|
switch status {
|
|
case 302: return "Found"
|
|
case 404: return "Not Found"
|
|
case 405: return "Method Not Allowed"
|
|
default: return "Bad Request"
|
|
}
|
|
}
|
|
}
|