Implementing Robust WebSocket Communication in Swift

In today’s interconnected world, real-time communication is crucial for many applications. WebSockets provide a powerful solution for bidirectional, full-duplex communication between clients and servers. In this article, we’ll explore a robust implementation of a WebSocket handler in Swift, designed for use in iOS applications.

The WebSocketHandler Protocol

Let’s start by examining the WebSocketHandler protocol:

public protocol WebSocketHandler {
  func connect(request: NetworkRequest, completion: @escaping (Result<Void, EngineError.WebSocketError>) -> Void)
  func disconnect(completion: @escaping (Result<Void, EngineError.WebSocketError>) -> Void)
  func send<T>(message: T, completion: @escaping (Result<Void, EngineError.WebSocketError>) -> Void)  where T : Encodable
  func receive<T>(completion: @escaping (Result<T, EngineError.WebSocketError>) -> Void)  where T : Decodable
}

This protocol defines four essential methods for WebSocket communication:

1. connect: Establishes a WebSocket connection
2. disconnect: Closes the WebSocket connection
3. send: Sends a message through the WebSocket
4. receive: Receives a message from the WebSocket

Each method uses a completion handler with a Result type, allowing for proper error handling.

The DefaultWebSocketHandler Implementation

Now, let’s dive into the DefaultWebSocketHandler class, which implements the WebSocketHandler protocol:

public class DefaultWebSocketHandler: WebSocketHandler {
  private var webSocketTask: URLSessionWebSocketTask?
  private let logger: Logger
  private let responseHandler: HTTPResponseHandler
  
  public init(logger: Logger = DefaultLogger.shared,
              responseHandler: HTTPResponseHandler = DefaultHTTPResponseHandler()) {
    self.logger = logger
    self.responseHandler = responseHandler
  }
  
  // ... (implementation of protocol methods)
}

This class uses URLSessionWebSocketTask to manage the WebSocket connection. It also incorporates a logger for debugging and a response handler for decoding received messages.

Connecting to a WebSocket

The connect method establishes a WebSocket connection:

public func connect(request: NetworkRequest, completion: @escaping (Result<Void, EngineError.WebSocketError>) -> Void) {
  guard webSocketTask == nil else {
    completion(.failure(.connectionFailed("WebSocket is already connected")))
    return
  }
  
  let urlSession = URLSession(configuration:.default, delegate: nil, delegateQueue: OperationQueue())
  do {
    let urlRequest = try request.urlRequest()
    webSocketTask = urlSession.webSocketTask(with: urlRequest)
    webSocketTask?.resume()
    
    // Simple connection check
    receive { [weak self] result in
      switch result {
      case.success:
        self?.logger.log(file: #file, function: #function, line: #line, logLevel:.info, "WebSocket connected")
        completion(.success(()))
      case.failure(let error):
        self?.logger.log(file: #file, function: #function, line: #line, logLevel:.error, error)
        completion(.failure(error))
      }
    }
  } catch {
    completion(.failure(.connectionFailed(error.localizedDescription)))
  }
}

This method creates a URLSessionWebSocketTask and attempts to connect. It then performs a simple connection check by trying to receive a message, confirming the connection's success or failure.

Sending Messages

The send method encodes and sends messages through the WebSocket:

public func send<T>(message: T, completion: @escaping (Result<Void, EngineError.WebSocketError>) -> Void) where T: Encodable {
  guard let webSocketTask = webSocketTask else {
    completion(.failure(.messageSendFailed("WebSocket is not connected")))
    return
  }
  
  do {
    let data = try JSONEncoder().encode(message)
    let message = URLSessionWebSocketTask.Message.data(data)
    webSocketTask.send(message) { [weak self] error in
      if let error = error {
        self?.logger.log(file: #file, function: #function, line: #line, logLevel:.error, error)
        completion(.failure(.messageSendFailed(error.localizedDescription)))
      } else {
        self?.logger.log(file: #file, function: #function, line: #line, logLevel:.info, "WebSocket message sent")
        completion(.success(()))
      }
    }
  } catch {
    let error = EngineError.WebSocketError.messageSendFailed("Failed to encode message")
    logger.log(file: #file, function: #function, line: #line, logLevel:.error, error)
    completion(.failure(error))
  }
}

This method first checks if the WebSocket is connected, then encodes the message and sends it using the webSocketTask.send method.

Receiving Messages

The receive method handles incoming messages:

public func receive<T>(completion: @escaping (Result<T, EngineError.WebSocketError>) -> Void) where T: Decodable {
  guard let webSocketTask = webSocketTask else {
    completion(.failure(.messageReceiveFailed("WebSocket is not connected")))
    return
  }
  
  webSocketTask.receive { [weak self] result in
    switch result {
    case .success(let message):
      switch message {
      case .data(let data):
        do {
          let typedMessage = try self?.responseHandler.decode(data: data, to: T.self)
          self?.logger.log(file: #file, function: #function, line: #line, logLevel: .info, "WebSocket message received")
          if let typedMessage = typedMessage {
            completion(.success(typedMessage))
          } else {
            let error = EngineError.WebSocketError.messageReceiveFailed("Failed to decode received message")
            self?.logger.log(file: #file, function: #function, line: #line, logLevel: .error, error)
            completion(.failure(error))
          }
        } catch {
          let error = EngineError.WebSocketError.messageReceiveFailed("Failed to decode received message: \(error.localizedDescription)")
          self?.logger.log(file: #file, function: #function, line: #line, logLevel: .error, error)
          completion(.failure(error))
        }
      case .string(let string):
        let error = EngineError.WebSocketError.messageReceiveFailed("Received unexpected string message: \(string)")
        self?.logger.log(file: #file, function: #function, line: #line, logLevel: .error, error)
        completion(.failure(error))
      @unknown default:
        let error = EngineError.WebSocketError.messageReceiveFailed("Received unknown message type")
        self?.logger.log(file: #file, function: #function, line: #line, logLevel: .error, error)
        completion(.failure(error))
      }
    case .failure(let error):
      self?.logger.log(file: #file, function: #function, line: #line, logLevel: .error, error)
      completion(.failure(.messageReceiveFailed(error.localizedDescription)))
    }
  }
}

This method receives messages from the WebSocket, decodes them using the responseHandler, and returns the result through the completion handler.

Disconnecting from the WebSocket

The disconnect method disconnect from the WebSocket connection:

public func disconnect(completion: @escaping (Result<Void, EngineError.WebSocketError>) -> Void) {
    guard let webSocketTask = webSocketTask else {
      completion(.failure(.connectionFailed("WebSocket is not connected")))
      return
    }
    
    webSocketTask.cancel(with:.goingAway, reason: nil)
    self.webSocketTask = nil
    logger.log(file: #file, function: #function, line: #line, logLevel:.info, "WebSocket disconnected")
    completion(.success(()))
  }

This example will demonstrate connecting to a WebSocket, sending a message, and receiving messages.

import Foundation

// First, let's define a simple message structure
struct ChatMessage: Codable {
    let sender: String
    let content: String
}

class ChatViewModel {
    private let webSocketHandler: WebSocketHandler
    
    init(webSocketHandler: WebSocketHandler = DefaultWebSocketHandler()) {
        self.webSocketHandler = webSocketHandler
    }
    
    func connectToChat() {
        // Create a network request for the WebSocket connection
        let request = NetworkRequest(url: URL(string: "wss://example.com/chat")!)
        
        webSocketHandler.connect(request: request) { result in
            switch result {
            case .success:
                print("Connected to chat successfully")
                self.startListeningForMessages()
            case .failure(let error):
                print("Failed to connect to chat: \(error)")
            }
        }
    }
    
    func sendMessage(_ message: String) {
        let chatMessage = ChatMessage(sender: "User", content: message)
        
        webSocketHandler.send(message: chatMessage) { result in
            switch result {
            case .success:
                print("Message sent successfully")
            case .failure(let error):
                print("Failed to send message: \(error)")
            }
        }
    }
    
    private func startListeningForMessages() {
        self.receiveMessage()
    }
    
    private func receiveMessage() {
        webSocketHandler.receive { [weak self] (result: Result<ChatMessage, EngineError.WebSocketError>) in
            switch result {
            case .success(let message):
                print("Received message: \(message.sender): \(message.content)")
                // Process the received message (e.g., update UI)
                
                // Continue listening for the next message
                self?.receiveMessage()
            case .failure(let error):
                print("Error receiving message: \(error)")
                // Handle the error (e.g., attempt reconnection)
            }
        }
    }
    
    func disconnect() {
        webSocketHandler.disconnect { result in
            switch result {
            case .success:
                print("Disconnected from chat successfully")
            case .failure(let error):
                print("Failed to disconnect from chat: \(error)")
            }
        }
    }
}

// Usage example
let chatViewModel = ChatViewModel()

// Connect to the chat
chatViewModel.connectToChat()

// Send a message
chatViewModel.sendMessage("Hello, everyone!")

// ... Later, when done with the chat
chatViewModel.disconnect()

In this example, we’ve created a ChatViewModel class that utilizes the DefaultWebSocketHandler to manage a chat connection. Here's a breakdown of the main components:

1. ChatMessage struct: A simple structure representing a chat message, conforming to Codable for easy encoding and decoding.
2. ChatViewModel class: This class encapsulates the chat functionality using the WebSocketHandler.
3. connectToChat(): Establishes a connection to the WebSocket chat server.
4. sendMessage(_:): Sends a chat message through the WebSocket.
5. startListeningForMessages() and receiveMessage(): These methods set up a continuous loop to receive incoming messages.
6. disconnect(): Closes the WebSocket connection when the chat session is finished.

Conclusion

This WebSocket handler implementation provides a robust foundation for real-time communication in Swift applications. It offers error handling, logging, and support for encodable and decodable message types. By using this handler, developers can easily integrate WebSocket functionality into their iOS apps, enabling real-time features and enhancing user experience.

Remember to handle potential edge cases, implement reconnection logic, and consider security measures when working with WebSockets in production environments. Happy coding!