An Introduction to WebSockets

Connecting to the Server

Your app doesn’t do much good if it can’t connect to the server to send the user’s questions. Your next step is to build that connection.

Open WebSocketController.swift and locate connect(). Replace // TODO: Implement with the following code to start the socket connection:

    self.socket = session.webSocketTask(with: 
      URL(string: "ws://localhost:8080/socket")!)
    self.listen()
    self.socket.resume()

This will set self.socket to a URLSessionWebSocketTask, start your listening cycle and resume the task. If you’ve used URLSession, this will look familiar.

Next, you’ll tell URLSession to receive new messages in listen. Replace the TODO comment with:

    // 1
    self.socket.receive { [weak self] (result) in
      guard let self = self else { return }
      // 2
      switch result {
      case .failure(let error):
        print(error)
        // 3
        let alert = Alert(
            title: Text("Unable to connect to server!"),
            dismissButton: .default(Text("Retry")) {
              self.alert = nil
              self.socket.cancel(with: .goingAway, reason: nil)
              self.connect()
            }
        )
        self.alert = alert
        return
      case .success(let message):
        // 4
        switch message {
        case .data(let data):
          self.handle(data)
        case .string(let str):
          guard let data = str.data(using: .utf8) else { return }
          self.handle(data)
        @unknown default:
          break
        }
      }
      // 5
      self.listen()
    }

Here’s what’s going on with this snippet:

1. First, you add a callback to URLSessionWebSocketTask.receive. When the server receives a message, it will execute this closure. self is weakly captured so before executing the rest of the closure, the guard ensures self is not nil.
2. The value passed into the closure is of type Result<URLSessionWebSocketTask.Message, Error>. The switch lets you handle both successful messages and errors.
3. In case of an error, create a SwiftUI Alert, which you’ll present to the user later.
4. In case of a successful message, pass the raw Data to your handle.
5. URLSessionWebSocketTask.receive registers a one-time callback. So after the callback executes, re-register the callback to keep listening for new messages.

Now your server is ready to receive messages. However, you still have a bit of work to do to let the server know what kinds of messages it will receive.

Creating the Data Structure

Before handling any data, you need to define your data structure. You’ll do that now.

Start by creating a new Swift file called WebSocketTypes.swift containing the following code:

import Foundation

enum QnAMessageType: String, Codable {
  // Client to server types
  case newQuestion
  // Server to client types
  case questionResponse, handshake, questionAnswer
}

struct QnAMessageSinData: Codable {
  let type: QnAMessageType
}

This is the base layer of your data structure. QnAMessageType declares the different types of messages, both client-to.server and server-to-client.

Every message you send between your client and server will contain a type field so the receiver knows what type of message to decode.

Now, create the actual message types on top of this layer. At the bottom of WebSocketTypes.swift, add the following:

struct QnAHandshake: Codable {
  let id: UUID
}

struct NewQuestionMessage: Codable {
  var type: QnAMessageType = .newQuestion
  let id: UUID
  let content: String
}

class NewQuestionResponse: Codable, Comparable {
  let success: Bool
  let message: String
  let id: UUID?
  var answered: Bool
  let content: String
  let createdAt: Date?
  
  static func < (lhs: NewQuestionResponse, rhs: NewQuestionResponse) -> Bool {
    guard let lhsDate = lhs.createdAt,
      let rhsDate = rhs.createdAt else {
        return false
    }

    return lhsDate < rhsDate
  }
  
  static func == (lhs: NewQuestionResponse, rhs: NewQuestionResponse) -> Bool {
    lhs.id == rhs.id
  }
}

struct QuestionAnsweredMessage: Codable {
  let questionId: UUID
}

All the above types, except for NewQuestionMessage, are server-to-client types and, thus, don’t have the type field set.

Handling the Data

Now you’ve set up your data structure, it’s time to let the server know how to handle the questions it receives from connecting clients.

Open WebSocketController.swift. Now, you can now implement the rest of the methods, starting with handle(_:).

Replace the TODO with the following code:

    do {
      // 1
      let sinData = try decoder.decode(QnAMessageSinData.self, from: data)
      // 2
      switch sinData.type {
      case .handshake:
        // 3
        print("Shook the hand")
        let message = try decoder.decode(QnAHandshake.self, from: data)
        self.id = message.id
      // 4
      case .questionResponse:
        try self.handleQuestionResponse(data)
      case .questionAnswer:
        try self.handleQuestionAnswer(data)
      default:
        break
      }
    } catch {
      print(error)
    }

This code does the following:

1. Decodes the QnAMessageSinData type you just declared. This will only decode the type field.
2. Switches on the decoded type field and does type-specific handling.
3. For the handshake type, decodes the full message and stores the ID.
4. For a question response and answer, passes the data along to specific implementation functions.

You’re making great progress! You probably want to try it out, but before you can test the full system, there are two things you need to do:

1. The server needs to send a custom handshake to the client to identify clients later.
2. The client needs to send new questions to the server.

Those will be the next two steps. You’ll start with the server.

Setting up the Server Handshake

In your websocket-backend project, open QuestionsController.swift and replace the contents of webSocket(req:socket:) with the following:

self.wsController.connect(socket)

At this point, your route collection will no longer control the socket connection.

Now, open WebSocketController.swift to implement connect(_:). Replace the TODO comment with the following:

    // 1
    let uuid = UUID()
    self.lock.withLockVoid {
      self.sockets[uuid] = ws
    }
    // 2
    ws.onBinary { [weak self] ws, buffer in
      guard let self = self,
        let data = buffer.getData(
          at: buffer.readerIndex, length: buffer.readableBytes) else {
            return
      }

      self.onData(ws, data)
    }
    // 3
    ws.onText { [weak self] ws, text in
      guard let self = self,
        let data = text.data(using: .utf8) else {
          return
      }

      self.onData(ws, data)
    }
    // 4
    self.send(message: QnAHandshake(id: uuid), to: .socket(ws))

This does the following:

1. Generates a random UUID for every socket connection to identify each one, then stores the connection in a dictionary based on this UUID.
2. When receiving binary data, it stores the data in SwiftNIO’s ByteBuffer. Using getData, you get a Data instance from the buffer, then pass it on to be handled.
3. For text data, you use String.data to get a Data instance and, again, pass it on to be handled.
4. Finally, you send a QnAHandshake message to the new connection, containing the UUID you generated above.

The QnAHandshake type here is the same one you created in the iOS app earlier. Look at the server-side data model in WebSocketTypes.swift. It’s the same as the iOS side, except the types that already have their type field set are reversed.

Next, back in WebSocketController.swift, put the following code in send(message:to:):

    logger.info("Sending \(T.self) to \(sendOption)")
    do {
      // 1
      let sockets: [WebSocket] = self.lock.withLock {
        switch sendOption {
        case .id(let id):
          return [self.sockets[id]].compactMap { $0 }
        case .socket(let socket):
          return [socket]
        }
      }
      
      // 2
      let encoder = JSONEncoder()
      let data = try encoder.encode(message)
      
      // 3
      sockets.forEach {
        $0.send(raw: data, opcode: .binary)
      }
    } catch {
      logger.report(error: error)
    }

The following happens here:

1. Based on WebSocketSendOption, you select a socket from the connected sockets.
2. You encode the message you want to send using JSONEncoder.
3. You send the encoded data to all selected sockets using the binary opcode.