feat: pass agent updates to UI (#35)

Author: ethanndickson

Coder Desktop/Coder Desktop/Preview Content/PreviewVPN.swift

@@ -4,21 +4,30 @@ import SwiftUI
 @MainActor
 final class PreviewVPN: Coder_Desktop.VPNService {
     @Published var state: Coder_Desktop.VPNServiceState = .disabled
-    @Published var agents: [Coder_Desktop.Agent] = [
-        Agent(id: UUID(), name: "dogfood2", status: .error, copyableDNS: "asdf.coder", workspaceName: "dogfood2"),
-        Agent(id: UUID(), name: "testing-a-very-long-name", status: .okay, copyableDNS: "asdf.coder",
-              workspaceName: "testing-a-very-long-name"),
-        Agent(id: UUID(), name: "opensrc", status: .warn, copyableDNS: "asdf.coder", workspaceName: "opensrc"),
-        Agent(id: UUID(), name: "gvisor", status: .off, copyableDNS: "asdf.coder", workspaceName: "gvisor"),
-        Agent(id: UUID(), name: "example", status: .off, copyableDNS: "asdf.coder", workspaceName: "example"),
-        Agent(id: UUID(), name: "dogfood2", status: .error, copyableDNS: "asdf.coder", workspaceName: "dogfood2"),
-        Agent(id: UUID(), name: "testing-a-very-long-name", status: .okay, copyableDNS: "asdf.coder",
-              workspaceName: "testing-a-very-long-name"),
-        Agent(id: UUID(), name: "opensrc", status: .warn, copyableDNS: "asdf.coder", workspaceName: "opensrc"),
-        Agent(id: UUID(), name: "gvisor", status: .off, copyableDNS: "asdf.coder", workspaceName: "gvisor"),
-        Agent(id: UUID(), name: "example", status: .off, copyableDNS: "asdf.coder", workspaceName: "example"),
+    @Published var agents: [UUID: Coder_Desktop.Agent] = [
+        UUID(): Agent(id: UUID(), name: "dev", status: .error, copyableDNS: "asdf.coder", wsName: "dogfood2",
+                      wsID: UUID()),
+        UUID(): Agent(id: UUID(), name: "dev", status: .okay, copyableDNS: "asdf.coder",
+                      wsName: "testing-a-very-long-name", wsID: UUID()),
+        UUID(): Agent(id: UUID(), name: "dev", status: .warn, copyableDNS: "asdf.coder", wsName: "opensrc",
+                      wsID: UUID()),
+        UUID(): Agent(id: UUID(), name: "dev", status: .off, copyableDNS: "asdf.coder", wsName: "gvisor",
+                      wsID: UUID()),
+        UUID(): Agent(id: UUID(), name: "dev", status: .off, copyableDNS: "asdf.coder", wsName: "example",
+                      wsID: UUID()),
+        UUID(): Agent(id: UUID(), name: "dev", status: .error, copyableDNS: "asdf.coder", wsName: "dogfood2",
+                      wsID: UUID()),
+        UUID(): Agent(id: UUID(), name: "dev", status: .okay, copyableDNS: "asdf.coder",
+                      wsName: "testing-a-very-long-name", wsID: UUID()),
+        UUID(): Agent(id: UUID(), name: "dev", status: .warn, copyableDNS: "asdf.coder", wsName: "opensrc",
+                      wsID: UUID()),
+        UUID(): Agent(id: UUID(), name: "dev", status: .off, copyableDNS: "asdf.coder", wsName: "gvisor",
+                      wsID: UUID()),
+        UUID(): Agent(id: UUID(), name: "dev", status: .off, copyableDNS: "asdf.coder", wsName: "example",
+                      wsID: UUID()),
     ]
     let shouldFail: Bool
+    let longError = "This is a long error to test the UI with long error messages"
 
     init(shouldFail: Bool = false) {
         self.shouldFail = shouldFail
@@ -35,10 +44,10 @@ final class PreviewVPN: Coder_Desktop.VPNService {
             do {
                 try await Task.sleep(for: .seconds(5))
             } catch {
-                state = .failed(.longTestError)
+                state = .failed(.internalError(longError))
                 return
             }
-            state = shouldFail ? .failed(.longTestError) : .connected
+            state = shouldFail ? .failed(.internalError(longError)) : .connected
         }
         defer { startTask = nil }
         await startTask?.value
@@ -57,7 +66,7 @@ final class PreviewVPN: Coder_Desktop.VPNService {
             do {
                 try await Task.sleep(for: .seconds(5))
             } catch {
-                state = .failed(.longTestError)
+                state = .failed(.internalError(longError))
                 return
             }
             state = .disabled

Coder Desktop/Coder Desktop/VPNService.swift

@@ -2,14 +2,12 @@ import NetworkExtension
 import os
 import SwiftUI
 import VPNLib
-import VPNXPC
 
 @MainActor
 protocol VPNService: ObservableObject {
     var state: VPNServiceState { get }
-    var agents: [Agent] { get }
+    var agents: [UUID: Agent] { get }
     func start() async
-    // Stop must be idempotent
     func stop() async
     func configureTunnelProviderProtocol(proto: NETunnelProviderProtocol?)
 }
@@ -26,12 +24,9 @@ enum VPNServiceError: Error, Equatable {
     case internalError(String)
     case systemExtensionError(SystemExtensionState)
     case networkExtensionError(NetworkExtensionState)
-    case longTestError
 
     var description: String {
         switch self {
-        case .longTestError:
-            "This is a long error to test the UI with long errors"
         case let .internalError(description):
             "Internal Error: \(description)"
         case let .systemExtensionError(state):
@@ -47,6 +42,7 @@ final class CoderVPNService: NSObject, VPNService {
     var logger = Logger(subsystem: Bundle.main.bundleIdentifier!, category: "vpn")
     lazy var xpc: VPNXPCInterface = .init(vpn: self)
     var terminating = false
+    var workspaces: [UUID: String] = [:]
 
     @Published var tunnelState: VPNServiceState = .disabled
     @Published var sysExtnState: SystemExtensionState = .uninstalled
@@ -61,7 +57,7 @@ final class CoderVPNService: NSObject, VPNService {
         return tunnelState
     }
 
-    @Published var agents: [Agent] = []
+    @Published var agents: [UUID: Agent] = [:]
 
     // systemExtnDelegate holds a reference to the SystemExtensionDelegate so that it doesn't get
     // garbage collected while the OSSystemExtensionRequest is in flight, since the OS framework
@@ -74,6 +70,16 @@ final class CoderVPNService: NSObject, VPNService {
         Task {
             await loadNetworkExtension()
         }
+        NotificationCenter.default.addObserver(
+            self,
+            selector: #selector(vpnDidUpdate(_:)),
+            name: .NEVPNStatusDidChange,
+            object: nil
+        )
+    }
+
+    deinit {
+        NotificationCenter.default.removeObserver(self)
     }
 
     func start() async {
@@ -84,16 +90,14 @@ final class CoderVPNService: NSObject, VPNService {
             return
         }
 
+        await enableNetworkExtension()
         // this ping is somewhat load bearing since it causes xpc to init
         xpc.ping()
-        tunnelState = .connecting
-        await enableNetworkExtension()
         logger.debug("network extension enabled")
     }
 
     func stop() async {
         guard tunnelState == .connected else { return }
-        tunnelState = .disconnecting
         await disableNetworkExtension()
         logger.info("network extension stopped")
     }
@@ -131,31 +135,97 @@ final class CoderVPNService: NSObject, VPNService {
     }
 
     func onExtensionPeerUpdate(_ data: Data) {
-        // TODO: handle peer update
         logger.info("network extension peer update")
         do {
-            let msg = try Vpn_TunnelMessage(serializedBytes: data)
+            let msg = try Vpn_PeerUpdate(serializedBytes: data)
             debugPrint(msg)
+            applyPeerUpdate(with: msg)
         } catch {
             logger.error("failed to decode peer update \(error)")
         }
     }
 
-    func onExtensionStart() {
-        logger.info("network extension reported started")
-        tunnelState = .connected
-    }
+    func applyPeerUpdate(with update: Vpn_PeerUpdate) {
+        // Delete agents
+        update.deletedAgents
+            .compactMap { UUID(uuidData: $0.id) }
+            .forEach { agentID in
+                agents[agentID] = nil
+            }
+        update.deletedWorkspaces
+            .compactMap { UUID(uuidData: $0.id) }
+            .forEach { workspaceID in
+                workspaces[workspaceID] = nil
+                for (id, agent) in agents where agent.wsID == workspaceID {
+                    agents[id] = nil
+                }
+            }
 
-    func onExtensionStop() {
-        logger.info("network extension reported stopped")
-        tunnelState = .disabled
-        if terminating {
-            NSApp.reply(toApplicationShouldTerminate: true)
+        // Update workspaces
+        for workspaceProto in update.upsertedWorkspaces {
+            if let workspaceID = UUID(uuidData: workspaceProto.id) {
+                workspaces[workspaceID] = workspaceProto.name
+            }
+        }
+
+        for agentProto in update.upsertedAgents {
+            guard let agentID = UUID(uuidData: agentProto.id) else {
+                continue
+            }
+            guard let workspaceID = UUID(uuidData: agentProto.workspaceID) else {
+                continue
+            }
+            let workspaceName = workspaces[workspaceID] ?? "Unknown Workspace"
+            let newAgent = Agent(
+                id: agentID,
+                name: agentProto.name,
+                // If last handshake was not within last five minutes, the agent is unhealthy
+                status: agentProto.lastHandshake.date > Date.now.addingTimeInterval(-300) ? .okay : .off,
+                copyableDNS: agentProto.fqdn.first ?? "UNKNOWN",
+                wsName: workspaceName,
+                wsID: workspaceID
+            )
+
+            // An existing agent with the same name, belonging to the same workspace
+            // is from a previous workspace build, and should be removed.
+            agents
+                .filter { $0.value.name == agentProto.name && $0.value.wsID == workspaceID }
+                .forEach { agents[$0.key] = nil }
+
+            agents[agentID] = newAgent
         }
     }
+}
 
-    func onExtensionError(_ error: NSError) {
-        logger.error("network extension reported error: \(error)")
-        tunnelState = .failed(.internalError(error.localizedDescription))
+extension CoderVPNService {
+    @objc private func vpnDidUpdate(_ notification: Notification) {
+        guard let connection = notification.object as? NETunnelProviderSession else {
+            return
+        }
+        switch connection.status {
+        case .disconnected:
+            if terminating {
+                NSApp.reply(toApplicationShouldTerminate: true)
+            }
+            connection.fetchLastDisconnectError { err in
+                self.tunnelState = if let err {
+                    .failed(.internalError(err.localizedDescription))
+                } else {
+                    .disabled
+                }
+            }
+        case .connecting:
+            tunnelState = .connecting
+        case .connected:
+            tunnelState = .connected
+        case .reasserting:
+            tunnelState = .connecting
+        case .disconnecting:
+            tunnelState = .disconnecting
+        case .invalid:
+            tunnelState = .failed(.networkExtensionError(.unconfigured))
+        @unknown default:
+            tunnelState = .disabled
+        }
     }
 }

Coder Desktop/Coder Desktop/Views/Agent.swift

@@ -1,18 +1,27 @@
 import SwiftUI
 
-struct Agent: Identifiable, Equatable {
+struct Agent: Identifiable, Equatable, Comparable {
     let id: UUID
     let name: String
     let status: AgentStatus
     let copyableDNS: String
-    let workspaceName: String
+    let wsName: String
+    let wsID: UUID
+
+    // Agents are sorted by status, and then by name
+    static func < (lhs: Agent, rhs: Agent) -> Bool {
+        if lhs.status != rhs.status {
+            return lhs.status < rhs.status
+        }
+        return lhs.wsName.localizedCompare(rhs.wsName) == .orderedAscending
+    }
 }
 
-enum AgentStatus: Equatable {
-    case okay
-    case warn
-    case error
-    case off
+enum AgentStatus: Int, Equatable, Comparable {
+    case okay = 0
+    case warn = 1
+    case error = 2
+    case off = 3
 
     public var color: Color {
         switch self {
@@ -22,16 +31,20 @@ enum AgentStatus: Equatable {
         case .off: .gray
         }
     }
+
+    static func < (lhs: AgentStatus, rhs: AgentStatus) -> Bool {
+        lhs.rawValue < rhs.rawValue
+    }
 }
 
 struct AgentRowView: View {
-    let workspace: Agent
+    let agent: Agent
     let baseAccessURL: URL
     @State private var nameIsSelected: Bool = false
     @State private var copyIsSelected: Bool = false
 
     private var fmtWsName: AttributedString {
-        var formattedName = AttributedString(workspace.name)
+        var formattedName = AttributedString(agent.wsName)
         formattedName.foregroundColor = .primary
         var coderPart = AttributedString(".coder")
         coderPart.foregroundColor = .gray
@@ -41,7 +54,7 @@ struct AgentRowView: View {
 
     private var wsURL: URL {
         // TODO: CoderVPN currently only supports owned workspaces
-        baseAccessURL.appending(path: "@me").appending(path: workspace.workspaceName)
+        baseAccessURL.appending(path: "@me").appending(path: agent.wsName)
     }
 
     var body: some View {
@@ -50,10 +63,10 @@ struct AgentRowView: View {
                 HStack(spacing: Theme.Size.trayPadding) {
                     ZStack {
                         Circle()
-                            .fill(workspace.status.color.opacity(0.4))
+                            .fill(agent.status.color.opacity(0.4))
                             .frame(width: 12, height: 12)
                         Circle()
-                            .fill(workspace.status.color.opacity(1.0))
+                            .fill(agent.status.color.opacity(1.0))
                             .frame(width: 7, height: 7)
                     }
                     Text(fmtWsName).lineLimit(1).truncationMode(.tail)
@@ -69,7 +82,7 @@ struct AgentRowView: View {
             }.buttonStyle(.plain)
             Button {
                 // TODO: Proper clipboard abstraction
-                NSPasteboard.general.setString(workspace.copyableDNS, forType: .string)
+                NSPasteboard.general.setString(agent.copyableDNS, forType: .string)
             } label: {
                 Image(systemName: "doc.on.doc")
                     .symbolVariant(.fill)

Coder Desktop/Coder Desktop/Views/Agents.swift

@@ -10,11 +10,12 @@ struct Agents<VPN: VPNService, S: Session>: View {
 
     var body: some View {
         Group {
-            // Workspaces List
+            // Agents List
             if vpn.state == .connected {
-                let visibleData = viewAll ? vpn.agents : Array(vpn.agents.prefix(defaultVisibleRows))
-                ForEach(visibleData, id: \.id) { workspace in
-                    AgentRowView(workspace: workspace, baseAccessURL: session.baseAccessURL!)
+                let sortedAgents = vpn.agents.values.sorted()
+                let visibleData = viewAll ? sortedAgents[...] : sortedAgents.prefix(defaultVisibleRows)
+                ForEach(visibleData, id: \.id) { agent in
+                    AgentRowView(agent: agent, baseAccessURL: session.baseAccessURL!)
                         .padding(.horizontal, Theme.Size.trayMargin)
                 }
                 if vpn.agents.count > defaultVisibleRows {

Coder Desktop/Coder Desktop/Views/Util.swift

@@ -12,3 +12,22 @@ final class Inspection<V> {
         }
     }
 }
+
+extension UUID {
+    var uuidData: Data {
+        withUnsafePointer(to: uuid) {
+            Data(bytes: $0, count: MemoryLayout.size(ofValue: uuid))
+        }
+    }
+
+    init?(uuidData: Data) {
+        guard uuidData.count == 16 else {
+            return nil
+        }
+        var uuid: uuid_t = (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
+        withUnsafeMutableBytes(of: &uuid) {
+            $0.copyBytes(from: uuidData)
+        }
+        self.init(uuid: uuid)
+    }
+}