daltoniam / Starscream

Websockets in swift for iOS and OSX
Apache License 2.0
8.28k stars 1.2k forks source link

Calling `disconnect` on a WebSocket using a `FoundationTransport` never indicates when the disconnect is completed. #869

Open designatednerd opened 3 years ago

designatednerd commented 3 years ago

Tangentially related to #853, when I call disconnect on my web socket, it calls through to engine.stop(). In its completion block, that calls through to self?.forceStop(), which then calls through to the transport's disconnect method, which disconnects everything but doesn't seem to have anything that calls back and tells the socket the disconnect has completed.

I suspect what's happening is that because the delegate for the two streams is set to nil before the streams are closed, all the delegate stuff that would normally get called when the sockets close is not getting called.

Is this considered: a) expected behavior, and if we manually call disconnect we should just assume it eventually disconnected? b) a bug, and there should be some kind of indication that the socket disconnected that goes back up through the delegates to the didReceive(event: WebSocketEvent, client: WebSocket) method? c) Something else I'm not thinking of? 👽

nahung89 commented 3 years ago

To put more detail, there are two different flows comparing WSEngine and NativeEngine. When I look into NativeEngine, I realize that it creates new request every time start(request:) is called.

In the other hand, WSEngine must wait the disconnection is complete (as you said) in order to start over again.

From that point, I come up with a solution to create new WebSocket instance every time I perform a disconnect / connect. From that point I don't have to know when old WebSocket instance is disconnected completely.

public typealias WebSocketEventCallback = (WebSocketEvent) -> Void
public protocol WebSocketServiceInterface: AnyObject {
    var baseURL: URL { get }
    var cookies: [HTTPCookie] { get }
    var isConnected: Bool { get }
    var eventCallback: WebSocketEventCallback? { get set }

    func connect(endpoint: WebSocketEndpoint) -> Cancellable
    func disconnect()
    func write(data: Data, completion: (() -> Void)?)
    func write(string: String, completion: (() -> Void)?)
}

public final class WebSocketService: WebSocketServiceInterface {
    public let baseURL: URL
    public let cookies: [HTTPCookie]

    public var eventCallback: WebSocketEventCallback?

    public private(set) var isConnected = false

    private var socket: WebSocket?

    public init(baseURL: URL, cookies: [HTTPCookie]) {
        self.baseURL = baseURL
        self.cookies = cookies
    }

    public func connect(endpoint: WebSocketEndpoint) -> Cancellable {
        disconnect()

        let request = createRequest(endpoint)
        socket = WebSocket(request: request, useCustomEngine: false)
        socket?.onEvent = { [unowned self] event in
            self.handle(event: event)
        }

        socket?.connect()

        return CancelToken(action: { [weak socket] in
            socket?.disconnect()
        })
    }

    public func disconnect() {
        socket?.disconnect()
        socket = nil
        isConnected = false
    }

    public func write(data: Data, completion: (() -> Void)?) {
        guard let socket = socket else {
            completion?()
            return
        }
        socket.write(data: data, completion: completion)
    }

    public func write(string: String, completion: (() -> Void)?) {
        guard let socket = socket else {
            completion?()
            return
        }
        socket.write(string: string, completion: completion)
    }
}

// MARK: - Helpers

extension WebSocketService {
    private func createRequest(_ endpoint: WebSocketEndpoint) -> URLRequest {
        var request = URLRequest(url: baseURL.appendingPathComponent(endpoint.path))

        for (key, value) in endpoint.headers {
            request.setValue(value, forHTTPHeaderField: key)
        }

        let cookiesValue = cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
        request.setValue(cookiesValue, forHTTPHeaderField: "cookie")

        return request
    }

    private func handle(event: Starscream.WebSocketEvent) {
        Logger.shared.debug(msg: "Websocket received event: \(event)")

        switch event {
        case .connected:
            isConnected = true
            eventCallback?(.connected)

        case let .disconnected(reason, code):
            isConnected = false
            eventCallback?(.disconnected(.raw(reason: reason, code: Int(code))))

        case .cancelled:
            isConnected = false
            eventCallback?(.cancelled)

        case let .error(error):
            isConnected = false
            if let error = error {
                eventCallback?(.error(.unknown(error)))
            } else {
                eventCallback?(.error(.notSpecific))
            }

        case let .text(string):
            eventCallback?(.text(string))

        case let .binary(data):
            eventCallback?(.data(data))

        case .ping,
             .pong,
             .viabilityChanged,
             .reconnectSuggested:
            break
        }
    }
}