Skip to content

Commit

Permalink
Replace Starscream with WebSocketKit
Browse files Browse the repository at this point in the history
  • Loading branch information
@tagavari
tagavari committed Aug 6, 2022
1 parent c0e9b16 commit 23db14e
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 66 deletions.
34 changes: 17 additions & 17 deletions AirMessage.xcodeproj/project.pbxproj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
ED44738726A387DD006F5F0E /* KeychainManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED44738626A387DD006F5F0E /* KeychainManager.swift */; };
ED44738926A459CA006F5F0E /* FirebaseAuthHelper.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED44738826A459CA006F5F0E /* FirebaseAuthHelper.swift */; };
ED44738B26A46B68006F5F0E /* DispatchHelper.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED44738A26A46B68006F5F0E /* DispatchHelper.swift */; };
ED456D99289EFE1E00729CF6 /* WebSocketKit in Frameworks */ = {isa = PBXBuildFile; productRef = ED456D98289EFE1E00729CF6 /* WebSocketKit */; };
ED78BDCE2757A4190047EDC4 /* URLSessionCompat.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED78BDCD2757A4190047EDC4 /* URLSessionCompat.swift */; };
ED7CC006285110E9005C6B58 /* CustomQueue.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED7CC005285110E9005C6B58 /* CustomQueue.swift */; };
ED853E072744218900DCF446 /* ReadWriteLock.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED853E062744218900DCF446 /* ReadWriteLock.swift */; };
Expand All @@ -69,7 +70,6 @@
EDADEB1926A0E3D9001DA84A /* StorageManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = EDADEB1826A0E3D9001DA84A /* StorageManager.swift */; };
EDB04DBE273DDFDA00673AE4 /* SQLite in Frameworks */ = {isa = PBXBuildFile; productRef = EDB04DBD273DDFDA00673AE4 /* SQLite */; };
EDC1DAF8274DA24C00D32C6B /* AtomicValue.swift in Sources */ = {isa = PBXBuildFile; fileRef = EDC1DAF7274DA24C00D32C6B /* AtomicValue.swift */; };
EDC1DAFB274DAB7700D32C6B /* Starscream in Frameworks */ = {isa = PBXBuildFile; productRef = EDC1DAFA274DAB7700D32C6B /* Starscream */; };
EDC1DAFE274DABA400D32C6B /* DataProxyConnect.swift in Sources */ = {isa = PBXBuildFile; fileRef = EDC1DAFD274DABA400D32C6B /* DataProxyConnect.swift */; };
EDC1DB03274DBF0A00D32C6B /* ConnectConstants.swift in Sources */ = {isa = PBXBuildFile; fileRef = EDC1DB02274DBF0A00D32C6B /* ConnectConstants.swift */; };
EDC58E042756F59200935C49 /* Certificates in Resources */ = {isa = PBXBuildFile; fileRef = EDC58E032756F59200935C49 /* Certificates */; };
Expand Down Expand Up @@ -544,11 +544,11 @@
isa = PBXFrameworksBuildPhase;
buildActionMask = 2147483647;
files = (
EDC1DAFB274DAB7700D32C6B /* Starscream in Frameworks */,
EDB04DBE273DDFDA00673AE4 /* SQLite in Frameworks */,
ED9C293A2756684200EA07C8 /* Sentry in Frameworks */,
ED98B3F22711D84E0029DD09 /* Zip in Frameworks */,
EDC898662752BF290081BDB7 /* OpenSSL.framework in Frameworks */,
ED456D99289EFE1E00729CF6 /* WebSocketKit in Frameworks */,
ED9FBBC627E13F430037DF16 /* Zlib in Frameworks */,
EDDE47312688C2890029628B /* Swifter in Frameworks */,
ED98B3EF2711D8300029DD09 /* Ink in Frameworks */,
Expand Down Expand Up @@ -1201,9 +1201,9 @@
ED98B3EE2711D8300029DD09 /* Ink */,
ED98B3F12711D84E0029DD09 /* Zip */,
EDB04DBD273DDFDA00673AE4 /* SQLite */,
EDC1DAFA274DAB7700D32C6B /* Starscream */,
ED9C29392756684200EA07C8 /* Sentry */,
ED9FBBC527E13F430037DF16 /* Zlib */,
ED456D98289EFE1E00729CF6 /* WebSocketKit */,
);
productName = AirMessage;
productReference = ED28095125A20332009E7636 /* AirMessage.app */;
Expand Down Expand Up @@ -1283,8 +1283,8 @@
ED98B3ED2711D8300029DD09 /* XCRemoteSwiftPackageReference "ink" */,
ED98B3F02711D84E0029DD09 /* XCRemoteSwiftPackageReference "Zip" */,
EDB04DBC273DDFDA00673AE4 /* XCRemoteSwiftPackageReference "SQLite.swift" */,
EDC1DAF9274DAB7700D32C6B /* XCRemoteSwiftPackageReference "Starscream" */,
ED9C29382756684200EA07C8 /* XCRemoteSwiftPackageReference "sentry-cocoa" */,
ED456D97289EFE1E00729CF6 /* XCRemoteSwiftPackageReference "websocket-kit-osx10" */,
);
productRefGroup = ED28095225A20332009E7636 /* Products */;
projectDirPath = "";
Expand Down Expand Up @@ -1853,6 +1853,14 @@
/* End XCConfigurationList section */

/* Begin XCRemoteSwiftPackageReference section */
ED456D97289EFE1E00729CF6 /* XCRemoteSwiftPackageReference "websocket-kit-osx10" */ = {
isa = XCRemoteSwiftPackageReference;
repositoryURL = "https://github.com/tagavari/websocket-kit-osx10.git";
requirement = {
branch = osx10;
kind = branch;
};
};
ED98B3ED2711D8300029DD09 /* XCRemoteSwiftPackageReference "ink" */ = {
isa = XCRemoteSwiftPackageReference;
repositoryURL = "https://github.com/johnsundell/ink.git";
Expand Down Expand Up @@ -1885,14 +1893,6 @@
minimumVersion = 0.13.0;
};
};
EDC1DAF9274DAB7700D32C6B /* XCRemoteSwiftPackageReference "Starscream" */ = {
isa = XCRemoteSwiftPackageReference;
repositoryURL = "https://github.com/daltoniam/Starscream.git";
requirement = {
kind = upToNextMinorVersion;
minimumVersion = 4.0.0;
};
};
EDDE472F2688C2890029628B /* XCRemoteSwiftPackageReference "swifter" */ = {
isa = XCRemoteSwiftPackageReference;
repositoryURL = "https://github.com/httpswift/swifter.git";
Expand All @@ -1904,6 +1904,11 @@
/* End XCRemoteSwiftPackageReference section */

/* Begin XCSwiftPackageProductDependency section */
ED456D98289EFE1E00729CF6 /* WebSocketKit */ = {
isa = XCSwiftPackageProductDependency;
package = ED456D97289EFE1E00729CF6 /* XCRemoteSwiftPackageReference "websocket-kit-osx10" */;
productName = WebSocketKit;
};
ED98B3EE2711D8300029DD09 /* Ink */ = {
isa = XCSwiftPackageProductDependency;
package = ED98B3ED2711D8300029DD09 /* XCRemoteSwiftPackageReference "ink" */;
Expand All @@ -1928,11 +1933,6 @@
package = EDB04DBC273DDFDA00673AE4 /* XCRemoteSwiftPackageReference "SQLite.swift" */;
productName = SQLite;
};
EDC1DAFA274DAB7700D32C6B /* Starscream */ = {
isa = XCSwiftPackageProductDependency;
package = EDC1DAF9274DAB7700D32C6B /* XCRemoteSwiftPackageReference "Starscream" */;
productName = Starscream;
};
EDDE47302688C2890029628B /* Swifter */ = {
isa = XCSwiftPackageProductDependency;
package = EDDE472F2688C2890029628B /* XCRemoteSwiftPackageReference "swifter" */;
Expand Down
120 changes: 71 additions & 49 deletions AirMessage/Connection/Connect/DataProxyConnect.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,9 @@
//

import Foundation
import Starscream
import NIO
import NIOWebSocket
import WebSocketKit
import Sentry

class DataProxyConnect: DataProxy {
Expand All @@ -32,6 +34,7 @@ class DataProxyConnect: DataProxy {
//The max num of attempts before capping the delay time - not before giving up
private static let connectionRecoveryCountMax = 8

private let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
private let processingQueue = DispatchQueue(label: Bundle.main.bundleIdentifier! + ".proxy.connect.processing", qos: .utility)

private var isActive = false
Expand All @@ -42,6 +45,9 @@ class DataProxyConnect: DataProxy {
}

deinit {
//Shut down the event loop
try? eventLoopGroup.syncShutdownGracefully()

//Ensure the server proxy isn't running when we go out of scope
assert(!isActive, "DataProxyConnect was deinitialized while active")

Expand Down Expand Up @@ -76,16 +82,39 @@ class DataProxyConnect: DataProxy {
var components = URLComponents(string: Bundle.main.infoDictionary!["CONNECT_ENDPOINT"] as! String)!
components.queryItems = queryItems

//Create the request
var urlRequest = URLRequest(url: components.url!)
urlRequest.addValue("app", forHTTPHeaderField: "Origin")
var headers = HTTPHeaders()
headers.add(name: "Origin", value: "app")

//Create the WebSocket connection
let socket = WebSocket(request: urlRequest)
socket.delegate = self
socket.callbackQueue = processingQueue
socket.connect()
webSocket = socket
WebSocket.connect(to: components.url!, headers: headers, configuration: WebSocketClient.Configuration(), on: eventLoopGroup, onUpgrade: { [weak self] webSocket in
//Report open event
self?.processingQueue.async { [weak self] in
guard let self = self else { return }
self.webSocket = webSocket
self.onWSConnect()
}

//Set listener
webSocket.onBinary { [weak self] _, byteBuffer in
self?.processingQueue.async { [weak self] in
self?.onWSReceive(data: Data(byteBuffer.readableBytesView))
}
}
webSocket.onClose.whenComplete { [weak self] result in
self?.processingQueue.async { [weak self] in
switch result {
case .failure(let error):
self?.onWSError(error: error)
case .success(_):
self?.onWSDisconnect(withCode: webSocket.closeCode ?? WebSocketErrorCode.normalClosure)
}
}
}
}).whenFailure { [weak self] error in
self?.processingQueue.async { [weak self] in
self?.onWSError(error: error)
}
}

isActive = true
}
Expand All @@ -106,7 +135,7 @@ class DataProxyConnect: DataProxy {
NotificationNames.postUpdateConnectionCount(0)

//Disconnect
socket.disconnect()
_ = socket.close()
delegate?.dataProxy(self, didStopWithState: .stopped, isRecoverable: false)
}

Expand Down Expand Up @@ -153,7 +182,13 @@ class DataProxyConnect: DataProxy {
packer.pack(data: secureData)

//Send the message
socket.write(data: packer.data, completion: onSent)
if let onSent = onSent {
let promise = eventLoopGroup.next().makePromise(of: Void.self)
socket.send([UInt8](packer.data), promise: promise)
promise.futureResult.whenSuccess(onSent)
} else {
socket.send([UInt8](packer.data))
}
}

func send(pushNotification data: Data, version: Int) {
Expand All @@ -169,7 +204,7 @@ class DataProxyConnect: DataProxy {
packer.pack(data: data)

//Send the message
socket.write(data: packer.data, completion: nil)
socket.send([UInt8](packer.data))
}

func disconnect(client: ClientConnection) {
Expand All @@ -185,7 +220,7 @@ class DataProxyConnect: DataProxy {
packer.pack(int: clientID)

//Send the message
socket.write(data: packer.data, completion: nil)
socket.send([UInt8](packer.data))

//Remove the client
removeClient(clientID: clientID)
Expand Down Expand Up @@ -228,7 +263,7 @@ class DataProxyConnect: DataProxy {
assertDispatchQueue(processingQueue)

//Ping
webSocket?.write(ping: Data())
webSocket?.sendPing()
}

//MARK: Handshake Timer
Expand Down Expand Up @@ -261,7 +296,7 @@ class DataProxyConnect: DataProxy {
assertDispatchQueue(processingQueue)

//Disconnect
webSocket?.disconnect()
_ = webSocket?.close()

//Clean up
stopHandshakeTimer()
Expand Down Expand Up @@ -314,8 +349,8 @@ class DataProxyConnect: DataProxy {
startHandshakeTimer()
}

private func onWSDisconnect(reason: String, code: UInt16) {
LogManager.log("Connection to Connect relay lost: \(code) / \(reason)", level: .info)
private func onWSDisconnect(withCode code: WebSocketErrorCode) {
LogManager.log("Connection to Connect relay lost: \(code)", level: .info)

//Stop the ping timer
stopPingTimer()
Expand All @@ -331,20 +366,26 @@ class DataProxyConnect: DataProxy {
//Map the error code
let localError: ServerState
switch code {
case 1006 /* Abnormal close */, CloseCode.normal.rawValue:
case .normalClosure:
localError = .errorInternet
case CloseCode.protocolError.rawValue, CloseCode.policyViolated.rawValue:
case .protocolError, .policyViolation:
localError = .errorConnBadRequest
case ConnectCloseCode.incompatibleProtocol.rawValue:
localError = .errorConnOutdated
case ConnectCloseCode.accountValidation.rawValue:
localError = .errorConnValidation
case ConnectCloseCode.serverTokenRefresh.rawValue:
localError = .errorConnToken
case ConnectCloseCode.noActivation.rawValue:
localError = .errorConnActivation
case ConnectCloseCode.otherLocation.rawValue:
localError = .errorConnAccountConflict
case .unknown(let rawCode):
//Custom AirMessage codes
switch rawCode {
case ConnectCloseCode.incompatibleProtocol.rawValue:
localError = .errorConnOutdated
case ConnectCloseCode.accountValidation.rawValue:
localError = .errorConnValidation
case ConnectCloseCode.serverTokenRefresh.rawValue:
localError = .errorConnToken
case ConnectCloseCode.noActivation.rawValue:
localError = .errorConnActivation
case ConnectCloseCode.otherLocation.rawValue:
localError = .errorConnAccountConflict
default:
localError = .errorExternal
}
default:
localError = .errorExternal
}
Expand Down Expand Up @@ -480,7 +521,7 @@ class DataProxyConnect: DataProxy {
}

//Disconnect
onWSDisconnect(reason: "", code: CloseCode.normal.rawValue)
onWSDisconnect(withCode: .normalClosure)
}

//MARK: Message handling
Expand Down Expand Up @@ -534,22 +575,3 @@ class DataProxyConnect: DataProxy {
delegate?.dataProxy(self, didDisconnectClient: client, totalCount: connectionCount)
}
}

//MARK: WebSocket Delegate

extension DataProxyConnect: WebSocketDelegate {
func didReceive(event: WebSocketEvent, client: WebSocket) {
switch event {
case .connected(_): onWSConnect()
case .disconnected(let reason, let code): onWSDisconnect(reason: reason, code: code)
case .binary(let data): onWSReceive(data: data)
case .error(let error): onWSError(error: error)
case .cancelled: onWSError(error: nil)
case .viabilityChanged(let viable):
if !viable {
onWSError(error: nil)
}
default: break
}
}
}

0 comments on commit 23db14e

Please sign in to comment.