@@ -46,8 +46,8 @@ public class PinghubClient {
4646
4747 /// Initializes the client with an OAuth2 token.
4848 ///
49- public convenience init ( token: String ) {
50- let socket = starscreamSocket ( url: PinghubClient . endpoint, token: token)
49+ public convenience init ( token: String , endpoint : URL ? = nil ) {
50+ let socket = starscreamSocket ( url: endpoint ?? PinghubClient . endpoint, token: token)
5151 self . init ( socket: socket)
5252 }
5353
@@ -83,14 +83,17 @@ public class PinghubClient {
8383 guard let client = self else {
8484 return
8585 }
86- let error = error as NSError ?
87- let filteredError : NSError ? = error. flatMap ( { error in
88- // Filter out normal disconnects that we initiated
89- if error. domain == WebSocket . ErrorDomain && error. code == Int ( CloseCode . normal. rawValue) {
90- return nil
91- }
92- return error
93- } )
86+
87+ let filteredError : Error ?
88+
89+ if ( error as? WSError ) ? . code == UInt16 ( CloseCode . normal. rawValue) {
90+ filteredError = nil
91+ } else if ( error as? ConnectionClosed ) ? . code == . normalClosure {
92+ filteredError = nil
93+ } else {
94+ filteredError = error
95+ }
96+
9497 client. delegate? . pinghubDidDisconnect ( client, error: filteredError)
9598 }
9699 socket. onData = { [ weak self] data in
@@ -218,13 +221,85 @@ internal protocol Socket: AnyObject {
218221// MARK: - Starscream
219222
220223private func starscreamSocket( url: URL , token: String ) -> Socket {
221- var request = URLRequest ( url: PinghubClient . endpoint )
224+ var request = URLRequest ( url: url )
222225 request. setValue ( " Bearer \( token) " , forHTTPHeaderField: " Authorization " )
223- return WebSocket ( request: request)
226+ return StarscreamWebSocket ( socket: WebSocket ( request: request) )
227+ }
228+
229+ private class StarscreamWebSocket : Socket {
230+
231+ var onConnect : ( ( ) -> Void ) ?
232+ var onDisconnect : ( ( Error ? ) -> Void ) ?
233+ var onText : ( ( String ) -> Void ) ?
234+ var onData : ( ( Data ) -> Void ) ?
235+
236+ let socket : WebSocket
237+
238+ init ( socket: WebSocket ) {
239+ self . socket = socket
240+
241+ socket. onEvent = { [ weak self] event in
242+ guard let self else { return }
243+
244+ switch event {
245+ case . connected:
246+ self . onConnect ? ( )
247+ case let . disconnected( _, code) :
248+ self . onDisconnect ? ( ConnectionClosed ( code: Int ( code) ) )
249+ case . peerClosed:
250+ // Remote peer has closed the network connection. See `ConnectionState.peerClosed`.
251+ self . onDisconnect ? ( ConnectionClosed ( code: . normal) )
252+ case let . text( text) :
253+ self . onText ? ( text)
254+ case let . binary( data) :
255+ self . onData ? ( data)
256+ case let . error( error) :
257+ self . onDisconnect ? ( error)
258+ case . pong, . ping, . viabilityChanged, . reconnectSuggested, . cancelled:
259+ break
260+ }
261+ }
262+ }
263+
264+ func connect( ) {
265+ socket. connect ( )
266+ }
267+
268+ func disconnect( ) {
269+ socket. disconnect ( )
270+ }
271+
224272}
225273
226- extension WebSocket : Socket {
227- @objc func disconnect( ) {
228- disconnect ( forceTimeout: nil )
274+ private struct ConnectionClosed : Error {
275+ var code : URLSessionWebSocketTask . CloseCode
276+
277+ init ( code: Int ) {
278+ self . code = . init( rawValue: code) ?? . protocolError
279+ }
280+
281+ init ( code: URLSessionWebSocketTask . CloseCode ) {
282+ self . code = code
283+ }
284+
285+ init ( code: Starscream . CloseCode ) {
286+ switch code {
287+ case . normal:
288+ self . code = . normalClosure
289+ case . goingAway:
290+ self . code = . goingAway
291+ case . protocolError:
292+ self . code = . protocolError
293+ case . protocolUnhandledType:
294+ self . code = . unsupportedData
295+ case . noStatusReceived:
296+ self . code = . noStatusReceived
297+ case . encoding:
298+ self . code = . invalidFramePayloadData
299+ case . policyViolated:
300+ self . code = . policyViolation
301+ case . messageTooBig:
302+ self . code = . messageTooBig
303+ }
229304 }
230305}
0 commit comments