From 28f3728d0ad1c6218c47dd9fcfccfedae7e9138f Mon Sep 17 00:00:00 2001
From: John Zhou <37914490+johnnzhou@users.noreply.github.com>
Date: Tue, 1 Oct 2024 10:29:17 -0700
Subject: [PATCH] Allow ICMP and HTTP Ping Client to bind to a specific device
 (#12)

---
 Package.swift                                 |  2 +-
 Sources/LCLPing/HTTP/HTTPPingClient.swift     | 29 ++++++-
 Sources/LCLPing/HTTP/NIOHTTPClient.swift      | 20 +++++
 Sources/LCLPing/ICMP/ICMPPingClient.swift     | 77 +++++++++++++++----
 Sources/LCLPing/Models/Errors+LCLPing.swift   |  9 +++
 .../Utilities/LCLPing+ChannelOption.swift     | 20 +++++
 .../Utilities/LCLPing+SocketOption.swift      | 37 +++++++++
 7 files changed, 175 insertions(+), 19 deletions(-)
 create mode 100644 Sources/LCLPing/Utilities/LCLPing+ChannelOption.swift
 create mode 100644 Sources/LCLPing/Utilities/LCLPing+SocketOption.swift

diff --git a/Package.swift b/Package.swift
index 1a26e82..71b6c30 100644
--- a/Package.swift
+++ b/Package.swift
@@ -16,7 +16,7 @@ let package = Package(
     ],
     dependencies: [
         // Dependencies declare other packages that this package depends on.
-        .package(url: "https://github.com/apple/swift-nio.git", from: "2.72.0"),
+        .package(url: "https://github.com/apple/swift-nio.git", from: "2.73.0"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.5.3"),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", from: "2.25.0"),
         .package(url: "https://github.com/apple/swift-collections.git", from: "1.1.1")
diff --git a/Sources/LCLPing/HTTP/HTTPPingClient.swift b/Sources/LCLPing/HTTP/HTTPPingClient.swift
index a01d7a4..86bdd47 100644
--- a/Sources/LCLPing/HTTP/HTTPPingClient.swift
+++ b/Sources/LCLPing/HTTP/HTTPPingClient.swift
@@ -178,6 +178,9 @@ extension HTTPPingClient {
 
         /// The DNS-resolved address according to the URL endpoint
         public let resolvedAddress: SocketAddress
+        
+        /// The outgoing device associated with the given interface name
+        public var device: NIONetworkDevice?
 
         /// Initialize a HTTP Ping Client `Configuration`.
         ///
@@ -190,6 +193,7 @@ extension HTTPPingClient {
         ///     - useServerTimimg: Indicate whether the HTTP Ping Client should take `ServerTiming` attribute
         /// from the reponse header.
         ///     - useURLSession: Indicate whether the HTTP Ping Client should use native URLSession implementation.
+        ///     - deviceName: the interface name for which the outbound data will be sent to.
         ///
         ///     - Throws:
         ///         - httpMissingHost: if URL does not include any host information.
@@ -201,7 +205,8 @@ extension HTTPPingClient {
                     connectionTimeout: TimeAmount = .seconds(5),
                     headers: [String: String] = Configuration.defaultHeaders,
                     useServerTiming: Bool = false,
-                    useURLSession: Bool = false) throws {
+                    useURLSession: Bool = false,
+                    deviceName: String? = nil) throws {
             self.url = url
             self.count = count
             self.readTimeout = readTimeout
@@ -244,6 +249,20 @@ extension HTTPPingClient {
             // NOTE: URLSession is not fully supported in swift-corelibs-foundation
             self.useURLSession = false
             #endif
+            
+            for device in try System.enumerateDevices() {
+                if device.name == deviceName, let address = device.address {
+                    switch (address.protocol, self.resolvedAddress.protocol) {
+                    case (.inet, .inet), (.inet6, .inet6):
+                        self.device = device
+                    default:
+                        continue
+                    }
+                }
+                if self.device != nil {
+                    break
+                }
+            }
         }
 
         /// Initialize a HTTP Ping Client `Configuration`.
@@ -257,6 +276,7 @@ extension HTTPPingClient {
         ///     - useServerTimimg: Indicate whether the HTTP Ping Client should take `ServerTiming` attribute
         /// from the reponse header.
         ///     - useURLSession: Indicate whether the HTTP Ping Client should use native URLSession implementation.
+        ///     - deviceName: the interface name for which the outbound data will be sent to
         ///
         ///     - Throws:
         ///         - httpMissingHost: if URL does not include any host information.
@@ -268,7 +288,8 @@ extension HTTPPingClient {
                     connectionTimeout: TimeAmount = .seconds(5),
                     headers: [String: String] = Configuration.defaultHeaders,
                     useServerTiming: Bool = false,
-                    useURLSession: Bool = false) throws {
+                    useURLSession: Bool = false,
+                    deviceName: String? = nil) throws {
             guard let urlObj = URL(string: url) else {
                 throw PingError.invalidURL(url)
             }
@@ -279,8 +300,8 @@ extension HTTPPingClient {
                           connectionTimeout: connectionTimeout,
                           headers: headers,
                           useServerTiming: useServerTiming,
-                          useURLSession: useURLSession
-                        )
+                          useURLSession: useURLSession,
+                          deviceName: deviceName)
         }
 
         /// Initialize a HTTP Ping Client `Configuration`.
diff --git a/Sources/LCLPing/HTTP/NIOHTTPClient.swift b/Sources/LCLPing/HTTP/NIOHTTPClient.swift
index 054d236..c6fa3eb 100644
--- a/Sources/LCLPing/HTTP/NIOHTTPClient.swift
+++ b/Sources/LCLPing/HTTP/NIOHTTPClient.swift
@@ -198,6 +198,26 @@ final class NIOHTTPClient: Pingable {
                 } catch {
                     return channel.eventLoop.makeFailedFuture(error)
                 }
+                
+                if let device = self.configuration.device {
+                    #if canImport(Darwin)
+                    switch device.address {
+                    case .v4:
+                        return channel.setOption(.ipOption(.ip_bound_if), value: CInt(device.interfaceIndex))
+                    case .v6:
+                        return channel.setOption(.ipv6Option(.ipv6_bound_if), value: CInt(device.interfaceIndex))
+                    case .unixDomainSocket:
+                        self.stateLock.withLock {
+                            self.state = .error
+                        }
+                        return channel.eventLoop.makeFailedFuture(PingError.icmpBindToUnixDomainSocket)
+                    default:
+                        ()
+                    }
+                    #elseif canImport(Glibc)
+                    return (channel as! SocketOptionProvider).setBindToDevice(device.name)
+                    #endif
+                }
 
                 return channel.eventLoop.makeSucceededVoidFuture()
             }
diff --git a/Sources/LCLPing/ICMP/ICMPPingClient.swift b/Sources/LCLPing/ICMP/ICMPPingClient.swift
index 7295843..7898297 100644
--- a/Sources/LCLPing/ICMP/ICMPPingClient.swift
+++ b/Sources/LCLPing/ICMP/ICMPPingClient.swift
@@ -78,23 +78,31 @@ public final class ICMPPingClient: Pingable {
                             }
                         }
                     }
-
-                    for cnt in 0..<self.configuration.count {
+                    
+                    let sendPromise = channel.eventLoop.makePromise(of: Void.self)
+                    sendPromise.futureResult.cascadeFailure(to: self.promise)
+                    
+                    func send(_ cnt: Int) {
+                        if cnt == self.configuration.count {
+                            sendPromise.succeed()
+                            return
+                        }
+                        let el = self.eventLoopGroup.next()
+                        let p = el.makePromise(of: Void.self)
                         logger.debug("Scheduled #\(cnt) request")
                         channel.eventLoop.scheduleTask(in: cnt * self.configuration.interval) {
-                            channel.write(
-                                ICMPPingClient.Request(
-                                    sequenceNum: UInt16(cnt),
-                                    identifier: self.identifier
-                                ),
-                                promise: nil
-                            )
-                        }
+                            channel.writeAndFlush(ICMPPingClient.Request(sequenceNum: UInt16(cnt), identifier: self.identifier), promise: p)
+                        }.futureResult.hop(to: el).cascadeFailure(to: sendPromise)
+                        
+                        p.futureResult.cascadeFailure(to: sendPromise)
+                        send(cnt + 1)
                     }
-
-                    return self.promise.futureResult.flatMap { pingResponse in
+                    
+                    send(0)
+                    
+                    return sendPromise.futureResult.and(self.promise.futureResult).flatMap { (_, pingResponse) in
                         let summary = pingResponse.summarize(host: self.configuration.resolvedAddress)
-                        self.stateLock.withLockVoid {
+                        self.stateLock.withLock {
                             self.state = .finished
                         }
                         return channel.eventLoop.makeSucceededFuture(summary)
@@ -166,6 +174,27 @@ public final class ICMPPingClient: Pingable {
                     }
                     return channel.eventLoop.makeFailedFuture(error)
                 }
+                
+                if let device = self.configuration.device {
+                    #if canImport(Darwin)
+                    switch device.address {
+                    case .v4:
+                        return channel.setOption(.ipOption(.ip_bound_if), value: CInt(device.interfaceIndex))
+                    case .v6:
+                        return channel.setOption(.ipv6Option(.ipv6_bound_if), value: CInt(device.interfaceIndex))
+                    case .unixDomainSocket:
+                        self.stateLock.withLock {
+                            self.state = .error
+                        }
+                        return channel.eventLoop.makeFailedFuture(PingError.icmpBindToUnixDomainSocket)
+                    default:
+                        ()
+                    }
+                    #elseif canImport(Glibc)
+                    return (channel as! SocketOptionProvider).setBindToDevice(device.name)
+                    #endif
+                }
+                
                 return channel.eventLoop.makeSucceededVoidFuture()
             }
     }
@@ -202,13 +231,18 @@ extension ICMPPingClient {
         /// Time, in second, to wait for a reply for each packet sent. Default is 1s.
         public var timeout: TimeAmount
 
+        /// The resolved socket address
         public let resolvedAddress: SocketAddress
+        
+        /// The outgoing device associated with the given interface name
+        public var device: NIONetworkDevice?
 
         public init(endpoint: EndpointTarget,
                     count: Int = 10,
                     interval: TimeAmount = .seconds(1),
                     timeToLive: UInt8 = 64,
-                    timeout: TimeAmount = .seconds(1)
+                    timeout: TimeAmount = .seconds(1),
+                    deviceName: String? = nil
         ) throws {
             self.endpoint = endpoint
             self.count = count
@@ -221,6 +255,21 @@ extension ICMPPingClient {
             case .ipv6(let addr, let port):
                 self.resolvedAddress = try SocketAddress.makeAddressResolvingHost(addr, port: port ?? 0)
             }
+
+            for device in try System.enumerateDevices() {
+                if device.name == deviceName, let address = device.address {
+                    switch (address.protocol, self.endpoint) {
+                    case (.inet, .ipv4), (.inet6, .ipv6):
+                        logger.info("device selcted is \(device)")
+                        self.device = device
+                    default:
+                        continue
+                    }
+                }
+                if self.device != nil {
+                    break
+                }
+            }
         }
     }
 
diff --git a/Sources/LCLPing/Models/Errors+LCLPing.swift b/Sources/LCLPing/Models/Errors+LCLPing.swift
index 0ba0545..e33924d 100644
--- a/Sources/LCLPing/Models/Errors+LCLPing.swift
+++ b/Sources/LCLPing/Models/Errors+LCLPing.swift
@@ -52,6 +52,8 @@ public enum PingError: Error {
     case icmpPointerIndicatesError
     case icmpMissingARequiredOption
     case icmpBadLength
+    case icmpDestinationNotMulticast
+    case icmpBindToUnixDomainSocket
 
     case httpMissingHost
     case httpMissingSchema
@@ -59,6 +61,7 @@ public enum PingError: Error {
     case httpInvalidHandlerState
     case httpMissingResponse
     case httpInvalidURLSessionTask(Int)
+    case httpBindToUnixDomainSocket
 
     case invalidHexFormat
 
@@ -165,6 +168,12 @@ extension PingError: CustomStringConvertible {
             return "Missing HTTP response."
         case .httpInvalidURLSessionTask(let id):
             return "URLSession Task \(id) is invalid."
+        case .icmpDestinationNotMulticast:
+            return "Destination address is not a multicast address."
+        case .icmpBindToUnixDomainSocket:
+            return "Cannot bind to a unix domain socket device."
+        case .httpBindToUnixDomainSocket:
+            return "Cannot bind to a unix domain socket device."
         }
     }
 
diff --git a/Sources/LCLPing/Utilities/LCLPing+ChannelOption.swift b/Sources/LCLPing/Utilities/LCLPing+ChannelOption.swift
new file mode 100644
index 0000000..fd65e42
--- /dev/null
+++ b/Sources/LCLPing/Utilities/LCLPing+ChannelOption.swift
@@ -0,0 +1,20 @@
+//
+// This source file is part of the LCL open source project
+//
+// Copyright (c) 2021-2024 Local Connectivity Lab and the project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE for license information
+// See CONTRIBUTORS for the list of project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+import NIOCore
+
+extension ChannelOption where Self == ChannelOptions.Types.SocketOption {
+    public static func ipv6Option(_ name: NIOBSDSocket.Option) -> Self {
+        .init(level: .ipv6, name: name)
+    }
+}
diff --git a/Sources/LCLPing/Utilities/LCLPing+SocketOption.swift b/Sources/LCLPing/Utilities/LCLPing+SocketOption.swift
new file mode 100644
index 0000000..2e654be
--- /dev/null
+++ b/Sources/LCLPing/Utilities/LCLPing+SocketOption.swift
@@ -0,0 +1,37 @@
+//
+// This source file is part of the LCL open source project
+//
+// Copyright (c) 2021-2024 Local Connectivity Lab and the project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE for license information
+// See CONTRIBUTORS for the list of project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+import NIOCore
+
+extension NIOBSDSocket.Option {
+    #if canImport(Darwin)
+    public static let ip_bound_if: NIOBSDSocket.Option = Self(rawValue: IP_BOUND_IF)
+    public static let ipv6_bound_if: NIOBSDSocket.Option = Self(rawValue: IPV6_BOUND_IF)
+    #elseif canImport(Glibc)
+    public static let so_bindtodevice = Self(rawValue: SO_BINDTODEVICE)
+    #endif
+}
+
+extension SocketOptionProvider {
+    #if canImport(Glibc)
+    /// Sets the socket option SO_BINDTODEVICE to `value`.
+    ///
+    /// - parameters:
+    ///     - value: The value to set SO_BINDTODEVICE to.
+    /// - returns: An `EventLoopFuture` that fires when the option has been set,
+    ///     or if an error has occurred.
+    public func setBindToDevice(_ value: String) -> EventLoopFuture<Void> {
+        self.unsafeSetSocketOption(level: .socket, name: .so_bindtodevice, value: value)
+    }
+    #endif
+}