Swift IPv4 add support for outbound fragments

Sources/SwiftNetwork/Protocols/FrameArray.swift

@@ -96,10 +96,11 @@ public struct FrameArray: ~Copyable {
         }
     }
 
-    public enum FrameIterationResult {
+    public enum FrameIterationResult: ~Copyable {
         case continueIterating
         case stopIterating
         case removeFrameAndContinue
+        case replaceWithFramesAndContinue(FrameArray)
     }
 
     @_optimize(speed)
@@ -108,7 +109,7 @@ public struct FrameArray: ~Copyable {
         var index = 0
         while index < count {
             let result = enumerator(&frames[index])
-            switch result {
+            switch consume result {
             case .continueIterating:
                 index += 1
                 continue
@@ -118,6 +119,17 @@ public struct FrameArray: ~Copyable {
                 frames.remove(at: index)
                 count -= 1
             // Don't increment index
+            case .replaceWithFramesAndContinue(var newFrames):
+                frames.remove(at: index)
+                count -= 1
+                let insertCount = newFrames.count
+                var insertIndex = index
+                while let newFrame = newFrames.popFirst() {
+                    frames.insert(newFrame, at: insertIndex)
+                    insertIndex += 1
+                }
+                index += insertCount
+                count += insertCount
             }
         }
     }

Sources/SwiftNetwork/Protocols/IPProtocol.swift

@@ -938,8 +938,104 @@ public struct IPProtocol: NetworkProtocol {
                     if dscpValue != 0 {
                         tos |= (dscpValue << 2)  // IPTOS_DSCP_SHIFT
                     }
-                    // TODO: Handle fragmentation cases differently
-                    let offset: UInt16 = 0x4000  // Don't Fragment
+
+                    let enableFragmentation: Bool
+                    if let fragmentationOverride = frame.fragmentationOverride {
+                        enableFragmentation = fragmentationOverride
+                    } else {
+                        enableFragmentation = self.flags.enableFragmentation
+                    }
+
+                    // Payload is the unclaimed bytes beyond the IPv4 header.
+                    let payloadLength = frame.unclaimedLength - IPv4Instance.headerLength
+                    // MTU minus the header gives the correct fragment payload room
+                    let mtu = self.pathProperties.mtu
+                    var maxPayloadPerFragment = 0
+                    if mtu > IPv4Instance.headerLength {
+                        maxPayloadPerFragment = mtu - IPv4Instance.headerLength
+                    }
+                    // Handle fragmentation if payloadLength is greater than maxPayloadPerFragment and enableFragmentation is enabled
+                    if enableFragmentation && maxPayloadPerFragment > 0 && payloadLength > maxPayloadPerFragment {
+                        // MTU-splitting path: fragment the oversized datagram.
+                        var randomNumber = SystemRandomNumberGenerator()
+                        let identifier = UInt16(truncatingIfNeeded: randomNumber.next())
+                        // Align fragment payload to blocks of 8 bytes - RFC 791.
+                        let fragmentRoom = maxPayloadPerFragment - (maxPayloadPerFragment % 8)
+                        guard fragmentRoom > 0 else {
+                            frame.finalize(success: false)
+                            return .removeFrameAndContinue
+                        }
+                        var cursor = 0
+                        var fragmentationSucceeded = true
+                        var fragmentFrames = FrameArray()
+                        while cursor < payloadLength {
+                            // Determine if last or hold large the chunk length is
+                            let remaining = payloadLength - cursor
+                            let isLast = remaining <= fragmentRoom
+                            let chunkLength = isLast ? remaining : fragmentRoom
+                            // Create the fragment frame with this chunk length
+                            var fragmentFrame = Frame(count: IPv4Instance.headerLength + chunkLength)
+                            // MF bit is always set except for the last fragment
+                            let ipOff = UInt16(isLast ? 0 : 0x2000) | UInt16(cursor / 8)
+                            let fragmentTotalLength = UInt16(IPv4Instance.headerLength + chunkLength)
+                            let result = Serializer.serialize(&fragmentFrame, claim: false) {
+                                write throws(SerializationError) in
+                                try write.uint8(versionAndHeaderLength)
+                                try write.uint8(tos)
+                                try write.uint16NetworkByteOrder(fragmentTotalLength)
+                                try write.uint16NetworkByteOrder(identifier)
+                                try write.uint16NetworkByteOrder(ipOff)
+                                try write.uint8(self.ttl)
+                                try write.uint8(self.ipProtocolNumber)
+                                try write.uint16(0)  // Checksum
+                                try write.uint32(localAddressValue)
+                                try write.uint32(remoteAddressValue)
+                            }
+                            guard result.isValid else {
+                                Logger.proto.error("Serializing IPv4 fragment failed with result: \(result)")
+                                fragmentFrame.finalize(success: false)
+                                fragmentationSucceeded = false
+                                break
+                            }
+                            let copied = frame.copyInto(
+                                &fragmentFrame,
+                                atOffset: IPv4Instance.headerLength,
+                                fromOffset: IPv4Instance.headerLength + cursor,
+                                length: chunkLength
+                            )
+                            guard copied == chunkLength else {
+                                fragmentFrame.finalize(success: false)
+                                fragmentationSucceeded = false
+                                break
+                            }
+                            do throws(ChecksumError) {
+                                if self.flags.csumOffload {
+                                    fragmentFrame.checksumOffloadFlags = ChecksumFlags.ip.rawValue
+                                } else {
+                                    let checksumValue = try fragmentFrame.ipChecksum(offset: 0, length: 20)
+                                    self.setChecksumValue(frame: &fragmentFrame, value: checksumValue)
+                                }
+                            } catch {
+                                Logger.proto.error("Failed to compute IPv4 fragment checksum")
+                                fragmentFrame.finalize(success: false)
+                                fragmentationSucceeded = false
+                                break
+                            }
+                            self.counters.txPackets += 1
+                            fragmentFrames.add(frame: fragmentFrame)
+                            cursor += chunkLength
+                        }
+                        frame.finalize(success: fragmentationSucceeded)
+                        if fragmentationSucceeded {
+                            return .replaceWithFramesAndContinue(fragmentFrames)
+                        }
+                        // This is a case where something went wrong on fragmentation and we need to remove any fragments that were created
+                        fragmentFrames.finalizeAllFramesAsFailed()
+                        return .removeFrameAndContinue
+                    }
+
+                    // No fragmentation, standard outbound path
+                    let offset: UInt16 = 0x4000  // Don't Fragment (IP_DF)
                     let identifier: UInt16 = 0
 
                     let result = Serializer.serialize(&frame, claim: false) { write throws(SerializationError) in
@@ -959,16 +1055,6 @@ public struct IPProtocol: NetworkProtocol {
                         frame.finalize(success: false)
                         return .removeFrameAndContinue
                     }
-                    var enableFragmentation = false
-                    if let fragmentationOverride = frame.fragmentationOverride, fragmentationOverride == true {
-                        enableFragmentation = true
-                    } else if self.flags.enableFragmentation || frame.fragmentationOverride == nil {
-                        enableFragmentation = true
-                    }
-
-                    if enableFragmentation {
-                        // TODO: Fragmentation
-                    }
 
                     do throws(ChecksumError) {
                         if self.flags.corruptChecksums {

Tests/SwiftNetworkTests/SwiftNetworkIPTests.swift

@@ -905,6 +905,127 @@ final class SwiftNetworkIPTests: NetTestCase {
         XCTAssertNil(readBytes, "Unexpectedly received a packet from overlapping IPv6 fragments (RFC 5722 violation)")
     }
 
+    func testIPv4OutboundFragmentation() {
+        // Send a 29-byte payload on a path with MTU set to 40 and fragmentation enabled.
+        // IPv4 header is 20 bytes, so max payload per fragment = 20 bytes, aligned to 8 = 16 bytes.
+        // Expected: 2 fragments — payload[0..<16] with MF=1, payload[16..<29] with MF=0.
+        let mtu = 40
+        let payload: [UInt8] = [
+            0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
+            0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10,
+            0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18,
+            0x19, 0x1A, 0x1B, 0x1C, 0x1D,
+        ]
+        let parameters = Parameters()
+        let expectation = XCTestExpectation()
+        let context = parameters.context
+        context.async {
+            defer { expectation.fulfill() }
+            // Set directInterface so networkIsSatisfied enables mtu
+            var path = PathProperties(parameters: parameters)
+            path.directInterface = Interface(
+                index: 1,
+                name: "lo0",
+                type: .loopback,
+                subtype: .other,
+                mtu: mtu
+            )
+            path.effectiveMTU = UInt32(mtu)
+
+            let reference = IPProtocol.instance(context: parameters.context)
+            let ipOptions = IPProtocol.options()
+            // Enable fragmentation through IPOptions
+            ipOptions.flags = IPProtocol.IPOptions.Flags(rawValue: ipOptions.flags.rawValue)
+                .union(.fragmentationEnabledOverridden)
+                .union(.fragmentationEnabled)
+            ipOptions.setLogID(prefix: "C", parent: "1", protocolLogIDNumber: 2)
+            ipOptions.setProtocolInstance(reference)
+            parameters.defaultStack.internet = .ip(ipOptions)
+
+            let udpOptions = UDPProtocol.options()
+            udpOptions.noMetadata = true
+            udpOptions.setLogID(prefix: "C", parent: "1", protocolLogIDNumber: 1)
+            parameters.defaultStack.transport = .udp(udpOptions)
+
+            let localEndpoint = Endpoint(address: IPv4Address(SwiftNetworkIPTests.localIPv4Address)!, port: 0)
+            let remoteEndpoint = Endpoint(address: IPv4Address(SwiftNetworkIPTests.remoteIPv4Address)!, port: 0)
+            let ipLinkage = OutboundDatagramLinkage(reference: reference)
+            guard
+                let upperHarness = DatagramUpperHarness(
+                    identifier: "Client",
+                    local: localEndpoint,
+                    remote: remoteEndpoint,
+                    parameters: parameters,
+                    path: path,
+                    context: parameters.context,
+                    lowerProtocol: ipLinkage
+                )
+            else {
+                XCTFail("Failed to attach IP to upper harness")
+                return
+            }
+            let lowerHarness = DatagramLowerHarness(identifier: "Client", context: parameters.context)
+            do {
+                try reference.attachLowerDatagramProtocol(
+                    lowerHarness.reference,
+                    remote: remoteEndpoint,
+                    local: localEndpoint,
+                    parameters: parameters,
+                    path: path
+                )
+            } catch {
+                XCTFail("Failed to attach IP to lower harness")
+                return
+            }
+            upperHarness.start { connected in XCTAssertTrue(connected) }
+            _ = upperHarness.write(payload)
+
+            var fragments: [[UInt8]] = []
+            while lowerHarness.hasOutboundPackets {
+                if let packets = lowerHarness.extractLastOutboundPacket() {
+                    fragments.append(packets)
+                }
+            }
+            // Expect exactly 2 fragments: 16-byte payload + 13-byte payload
+            XCTAssertEqual(fragments.count, 2, "Expected exactly 2 IPv4 fragments")
+            guard fragments.count == 2,
+                let fragmentOne = fragments.first,
+                let fragmentTwo = fragments.last
+            else { return }
+
+            // Both must be at least 20 bytes (IPv4 header)
+            XCTAssertGreaterThanOrEqual(fragmentOne.count, 20)
+            XCTAssertGreaterThanOrEqual(fragmentTwo.count, 20)
+            guard fragmentOne.count >= 20, fragmentTwo.count >= 20 else { return }
+
+            // Both fragments must share the same identifier (bytes 4–5 in the header)
+            // This is generatated randomly, so make sure both fragments carry the same id
+            let id1 = UInt16(fragmentOne[4]) << 8 | UInt16(fragmentOne[5])
+            let id2 = UInt16(fragmentTwo[4]) << 8 | UInt16(fragmentTwo[5])
+            XCTAssertEqual(id1, id2, "All fragments must share the same IP identifier")
+
+            // Fragment 1: MF=1 (This is the first fragment so MF=1 needs to be set), offset=0
+            let flagsOffset1 = UInt16(fragmentOne[6]) << 8 | UInt16(fragmentOne[7])
+            XCTAssertNotEqual(flagsOffset1 & 0x2000, 0, "Fragment 1 must have MF=1")
+            XCTAssertEqual(flagsOffset1 & 0x1FFF, 0, "Fragment 1 must have offset=0")
+            // Fragment 1 payload must be the first 16 bytes (first 16 of the 29 bytes payload)
+            XCTAssertEqual(fragmentOne.count, 20 + 16, "Fragment 1 must be header + 16 payload bytes")
+            XCTAssertEqual(Array(fragmentOne[20...]), Array(payload[0..<16]))
+
+            // Fragment 2: MF=0 (last fragment), offset=2 (16 bytes / 8 = 2 units)
+            let flagsOffset2 = UInt16(fragmentTwo[6]) << 8 | UInt16(fragmentTwo[7])
+            XCTAssertEqual(flagsOffset2 & 0x2000, 0, "Fragment 2 must have MF=0")
+            XCTAssertEqual(flagsOffset2 & 0x1FFF, 2, "Fragment 2 must have offset=2 (16 bytes)")
+            // Fragment 2 payload must be the remaining 13 bytes (last 13 of the 29 bytes payload)
+            XCTAssertEqual(fragmentTwo.count, 20 + 13, "Fragment 2 must be header + 13 payload bytes")
+            XCTAssertEqual(Array(fragmentTwo[20...]), Array(payload[16...]))
+
+            upperHarness.stop()
+            upperHarness.teardown()
+        }
+        wait(for: [expectation], timeout: 10.0)
+    }
+
     // Sets up a minimal IP harness to test different fragment and reassembly conditions
     private func processIPFragment(
         packets: [[UInt8]],