From e49d89d63a43664ca8ea9b87acff884bc538c21c Mon Sep 17 00:00:00 2001
From: Luke Howard <lukeh@padl.com>
Date: Fri, 2 Aug 2024 16:50:45 +1000
Subject: [PATCH] API for add/drop multicast membership

---
 Sources/IORingUtils/Extensions.swift | 11 ++++++++-
 Sources/IORingUtils/Socket.swift     | 37 ++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/Sources/IORingUtils/Extensions.swift b/Sources/IORingUtils/Extensions.swift
index 9119803..cb9a275 100644
--- a/Sources/IORingUtils/Extensions.swift
+++ b/Sources/IORingUtils/Extensions.swift
@@ -80,7 +80,7 @@ public extension IORing {
   }
 }
 
-extension UnsafeMutablePointer {
+extension UnsafePointer {
   func propertyBasePointer<Property>(to property: KeyPath<Pointee, Property>)
     -> UnsafePointer<Property>?
   {
@@ -89,6 +89,15 @@ extension UnsafeMutablePointer {
   }
 }
 
+extension UnsafeMutablePointer {
+  func propertyBasePointer<Property>(to property: KeyPath<Pointee, Property>)
+    -> UnsafeMutablePointer<Property>?
+  {
+    guard let offset = MemoryLayout<Pointee>.offset(of: property) else { return nil }
+    return (UnsafeMutableRawPointer(self) + offset).assumingMemoryBound(to: Property.self)
+  }
+}
+
 extension IORing {
   func connect(_ fd: FileDescriptorRepresentable, to address: any SocketAddress) async throws {
     var addressBuffer = [UInt8]()
diff --git a/Sources/IORingUtils/Socket.swift b/Sources/IORingUtils/Socket.swift
index aa87050..e578a10 100644
--- a/Sources/IORingUtils/Socket.swift
+++ b/Sources/IORingUtils/Socket.swift
@@ -142,6 +142,43 @@ public struct Socket: CustomStringConvertible, Equatable, Hashable, Sendable {
     try setBooleanOption(level: CInt(IPPROTO_TCP), option: TCP_NODELAY, to: true)
   }
 
+  private func _addOrDropMembership(
+    _ add: Bool,
+    address: sockaddr_ll,
+    on interfaceIndex: Int
+  ) throws {
+    guard let fileHandle else { throw Errno.badFileDescriptor }
+    var address = address
+    var mreq = packet_mreq()
+    mreq.mr_ifindex = Int32(interfaceIndex)
+    mreq.mr_type = UInt16(PACKET_MR_MULTICAST)
+    mreq.mr_alen = UInt16(address.sll_halen)
+    withUnsafeMutablePointer(to: &mreq.mr_address) { dstAddress in
+      withUnsafePointer(to: &address.sll_addr) { srcAddress in
+        let dstAddressPtr = dstAddress.propertyBasePointer(to: \.0)!
+        let srcAddressPtr = srcAddress.propertyBasePointer(to: \.0)!
+        memcpy(dstAddressPtr, srcAddressPtr, Int(address.sll_halen))
+      }
+    }
+    try Errno.throwingGlobalErrno {
+      setsockopt(
+        fileHandle.fileDescriptor,
+        SOL_PACKET,
+        add ? PACKET_ADD_MEMBERSHIP : PACKET_DROP_MEMBERSHIP,
+        &mreq,
+        socklen_t(MemoryLayout<packet_mreq>.size)
+      )
+    }
+  }
+
+  public func addMulticastMembership(for address: sockaddr_ll, on interfaceIndex: Int) throws {
+    try _addOrDropMembership(true, address: address, on: interfaceIndex)
+  }
+
+  public func dropMulticastMembership(for address: sockaddr_ll, on interfaceIndex: Int) throws {
+    try _addOrDropMembership(false, address: address, on: interfaceIndex)
+  }
+
   public func bind(port: UInt16) throws {
     switch Int32(domain) {
     case AF_INET: