Skip to content

Commit 810544e

Browse files
authored
Add RawSocketBootstrap (#2320)
1 parent 00341c9 commit 810544e

File tree

10 files changed

+786
-3
lines changed

10 files changed

+786
-3
lines changed

Sources/NIOCore/BSDSocketAPI.swift

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,13 @@ extension NIOBSDSocket.Option {
267267
/// Control multicast time-to-live.
268268
public static let ip_multicast_ttl: NIOBSDSocket.Option =
269269
NIOBSDSocket.Option(rawValue: IP_MULTICAST_TTL)
270+
271+
/// The IPv4 layer generates an IP header when sending a packet
272+
/// unless the ``ip_hdrincl`` socket option is enabled on the socket.
273+
/// When it is enabled, the packet must contain an IP header. For
274+
/// receiving, the IP header is always included in the packet.
275+
public static let ip_hdrincl: NIOBSDSocket.Option =
276+
NIOBSDSocket.Option(rawValue: IP_HDRINCL)
270277
}
271278

272279
// IPv6 Options

Sources/NIOCore/ChannelOption.swift

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,11 @@ public struct ChannelOptions {
276276
public static let socketOption = { (name: NIOBSDSocket.Option) -> Types.SocketOption in
277277
.init(level: .socket, name: name)
278278
}
279+
280+
/// - seealso: `SocketOption`.
281+
public static let ipOption = { (name: NIOBSDSocket.Option) -> Types.SocketOption in
282+
.init(level: .ip, name: name)
283+
}
279284

280285
/// - seealso: `SocketOption`.
281286
public static let tcpOption = { (name: NIOBSDSocket.Option) -> Types.SocketOption in

Sources/NIOPosix/BSDSocketAPICommon.swift

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,14 @@ extension NIOBSDSocket.SocketType {
8383
internal static let stream: NIOBSDSocket.SocketType =
8484
NIOBSDSocket.SocketType(rawValue: SOCK_STREAM)
8585
#endif
86+
87+
#if os(Linux)
88+
internal static let raw: NIOBSDSocket.SocketType =
89+
NIOBSDSocket.SocketType(rawValue: CInt(SOCK_RAW.rawValue))
90+
#else
91+
internal static let raw: NIOBSDSocket.SocketType =
92+
NIOBSDSocket.SocketType(rawValue: SOCK_RAW)
93+
#endif
8694
}
8795

8896
// IPv4 Options
Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the SwiftNIO open source project
4+
//
5+
// Copyright (c) 2022 Apple Inc. and the SwiftNIO project authors
6+
// Licensed under Apache License v2.0
7+
//
8+
// See LICENSE.txt for license information
9+
// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
10+
//
11+
// SPDX-License-Identifier: Apache-2.0
12+
//
13+
//===----------------------------------------------------------------------===//
14+
import NIOCore
15+
16+
/// A `RawSocketBootstrap` is an easy way to interact with IP based protocols other then TCP and UDP.
17+
///
18+
/// Example:
19+
///
20+
/// ```swift
21+
/// let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
22+
/// defer {
23+
/// try! group.syncShutdownGracefully()
24+
/// }
25+
/// let bootstrap = RawSocketBootstrap(group: group)
26+
/// .channelInitializer { channel in
27+
/// channel.pipeline.addHandler(MyChannelHandler())
28+
/// }
29+
/// let channel = try! bootstrap.bind(host: "127.0.0.1", ipProtocol: .icmp).wait()
30+
/// /* the Channel is now ready to send/receive IP packets */
31+
///
32+
/// try channel.closeFuture.wait() // Wait until the channel un-binds.
33+
/// ```
34+
///
35+
/// The `Channel` will operate on `AddressedEnvelope<ByteBuffer>` as inbound and outbound messages.
36+
public final class NIORawSocketBootstrap {
37+
38+
private let group: EventLoopGroup
39+
private var channelInitializer: Optional<ChannelInitializerCallback>
40+
@usableFromInline
41+
internal var _channelOptions: ChannelOptions.Storage
42+
43+
/// Create a `RawSocketBootstrap` on the `EventLoopGroup` `group`.
44+
///
45+
/// The `EventLoopGroup` `group` must be compatible, otherwise the program will crash. `RawSocketBootstrap` is
46+
/// compatible only with `MultiThreadedEventLoopGroup` as well as the `EventLoop`s returned by
47+
/// `MultiThreadedEventLoopGroup.next`. See `init(validatingGroup:)` for a fallible initializer for
48+
/// situations where it's impossible to tell ahead of time if the `EventLoopGroup` is compatible or not.
49+
///
50+
/// - parameters:
51+
/// - group: The `EventLoopGroup` to use.
52+
public convenience init(group: EventLoopGroup) {
53+
guard NIOOnSocketsBootstraps.isCompatible(group: group) else {
54+
preconditionFailure("RawSocketBootstrap is only compatible with MultiThreadedEventLoopGroup and " +
55+
"SelectableEventLoop. You tried constructing one with \(group) which is incompatible.")
56+
}
57+
self.init(validatingGroup: group)!
58+
}
59+
60+
/// Create a `RawSocketBootstrap` on the `EventLoopGroup` `group`, validating that `group` is compatible.
61+
///
62+
/// - parameters:
63+
/// - group: The `EventLoopGroup` to use.
64+
public init?(validatingGroup group: EventLoopGroup) {
65+
guard NIOOnSocketsBootstraps.isCompatible(group: group) else {
66+
return nil
67+
}
68+
self._channelOptions = ChannelOptions.Storage()
69+
self.group = group
70+
self.channelInitializer = nil
71+
}
72+
73+
/// Initialize the bound `Channel` with `initializer`. The most common task in initializer is to add
74+
/// `ChannelHandler`s to the `ChannelPipeline`.
75+
///
76+
/// - parameters:
77+
/// - handler: A closure that initializes the provided `Channel`.
78+
public func channelInitializer(_ handler: @escaping @Sendable (Channel) -> EventLoopFuture<Void>) -> Self {
79+
self.channelInitializer = handler
80+
return self
81+
}
82+
83+
/// Specifies a `ChannelOption` to be applied to the `Channel`.
84+
///
85+
/// - parameters:
86+
/// - option: The option to be applied.
87+
/// - value: The value for the option.
88+
@inlinable
89+
public func channelOption<Option: ChannelOption>(_ option: Option, value: Option.Value) -> Self {
90+
self._channelOptions.append(key: option, value: value)
91+
return self
92+
}
93+
94+
/// Bind the `Channel` to `host`.
95+
/// All packets or errors matching the `ipProtocol` specified are passed to the resulting `Channel`.
96+
///
97+
/// - parameters:
98+
/// - host: The host to bind on.
99+
/// - ipProtocol: The IP protocol used in the IP protocol/nextHeader field.
100+
public func bind(host: String, ipProtocol: NIOIPProtocol) -> EventLoopFuture<Channel> {
101+
return bind0(ipProtocol: ipProtocol) {
102+
return try SocketAddress.makeAddressResolvingHost(host, port: 0)
103+
}
104+
}
105+
106+
private func bind0(ipProtocol: NIOIPProtocol, _ makeSocketAddress: () throws -> SocketAddress) -> EventLoopFuture<Channel> {
107+
let address: SocketAddress
108+
do {
109+
address = try makeSocketAddress()
110+
} catch {
111+
return group.next().makeFailedFuture(error)
112+
}
113+
precondition(address.port == nil || address.port == 0, "port must be 0 or not set")
114+
func makeChannel(_ eventLoop: SelectableEventLoop) throws -> DatagramChannel {
115+
return try DatagramChannel(eventLoop: eventLoop,
116+
protocolFamily: address.protocol,
117+
protocolSubtype: .init(ipProtocol),
118+
socketType: .raw)
119+
}
120+
return withNewChannel(makeChannel: makeChannel) { (eventLoop, channel) in
121+
channel.register().flatMap {
122+
channel.bind(to: address)
123+
}
124+
}
125+
}
126+
127+
/// Connect the `Channel` to `host`.
128+
///
129+
/// - parameters:
130+
/// - host: The host to connect to.
131+
/// - ipProtocol: The IP protocol used in the IP protocol/nextHeader field.
132+
public func connect(host: String, ipProtocol: NIOIPProtocol) -> EventLoopFuture<Channel> {
133+
return connect0(ipProtocol: ipProtocol) {
134+
return try SocketAddress.makeAddressResolvingHost(host, port: 0)
135+
}
136+
}
137+
138+
private func connect0(ipProtocol: NIOIPProtocol, _ makeSocketAddress: () throws -> SocketAddress) -> EventLoopFuture<Channel> {
139+
let address: SocketAddress
140+
do {
141+
address = try makeSocketAddress()
142+
} catch {
143+
return group.next().makeFailedFuture(error)
144+
}
145+
func makeChannel(_ eventLoop: SelectableEventLoop) throws -> DatagramChannel {
146+
return try DatagramChannel(eventLoop: eventLoop,
147+
protocolFamily: address.protocol,
148+
protocolSubtype: .init(ipProtocol),
149+
socketType: .raw)
150+
}
151+
return withNewChannel(makeChannel: makeChannel) { (eventLoop, channel) in
152+
channel.register().flatMap {
153+
channel.connect(to: address)
154+
}
155+
}
156+
}
157+
158+
private func withNewChannel(makeChannel: (_ eventLoop: SelectableEventLoop) throws -> DatagramChannel, _ bringup: @escaping (EventLoop, DatagramChannel) -> EventLoopFuture<Void>) -> EventLoopFuture<Channel> {
159+
let eventLoop = self.group.next()
160+
let channelInitializer = self.channelInitializer ?? { _ in eventLoop.makeSucceededFuture(()) }
161+
let channelOptions = self._channelOptions
162+
163+
let channel: DatagramChannel
164+
do {
165+
channel = try makeChannel(eventLoop as! SelectableEventLoop)
166+
} catch {
167+
return eventLoop.makeFailedFuture(error)
168+
}
169+
170+
func setupChannel() -> EventLoopFuture<Channel> {
171+
eventLoop.assertInEventLoop()
172+
return channelOptions.applyAllChannelOptions(to: channel).flatMap {
173+
channelInitializer(channel)
174+
}.flatMap {
175+
eventLoop.assertInEventLoop()
176+
return bringup(eventLoop, channel)
177+
}.map {
178+
channel
179+
}.flatMapError { error in
180+
eventLoop.makeFailedFuture(error)
181+
}
182+
}
183+
184+
if eventLoop.inEventLoop {
185+
return setupChannel()
186+
} else {
187+
return eventLoop.flatSubmit {
188+
setupChannel()
189+
}
190+
}
191+
}
192+
}
193+
194+
#if swift(>=5.6)
195+
@available(*, unavailable)
196+
extension NIORawSocketBootstrap: Sendable {}
197+
#endif

Sources/NIOPosix/Socket.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ typealias IOVector = iovec
3232
/// - parameters:
3333
/// - protocolFamily: The protocol family to use (usually `AF_INET6` or `AF_INET`).
3434
/// - type: The type of the socket to create.
35-
/// - protocolSubtype: The subtype of the protocol, corresponding to the `protocol`
35+
/// - protocolSubtype: The subtype of the protocol, corresponding to the `protocolSubtype`
3636
/// argument to the socket syscall. Defaults to 0.
3737
/// - setNonBlocking: Set non-blocking mode on the socket.
3838
/// - throws: An `IOError` if creation of the socket failed.

Tests/LinuxMain.swift

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,7 @@ class LinuxMainRunner {
120120
testCase(PendingDatagramWritesManagerTests.allTests),
121121
testCase(PipeChannelTest.allTests),
122122
testCase(PriorityQueueTest.allTests),
123+
testCase(RawSocketBootstrapTests.allTests),
123124
testCase(SALChannelTest.allTests),
124125
testCase(SALEventLoopTests.allTests),
125126
testCase(SNIHandlerTest.allTests),

Tests/NIOPosixTests/DatagramChannelTests.swift

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import NIOCore
1717
@testable import NIOPosix
1818
import XCTest
1919

20-
private extension Channel {
20+
extension Channel {
2121
func waitForDatagrams(count: Int) throws -> [AddressedEnvelope<ByteBuffer>] {
2222
return try self.pipeline.context(name: "ByteReadRecorder").flatMap { context in
2323
if let future = (context.handler as? DatagramReadRecorder<ByteBuffer>)?.notifyForDatagrams(count) {
@@ -47,7 +47,7 @@ private extension Channel {
4747
/// A class that records datagrams received and forwards them on.
4848
///
4949
/// Used extensively in tests to validate messaging expectations.
50-
private class DatagramReadRecorder<DataType>: ChannelInboundHandler {
50+
final class DatagramReadRecorder<DataType>: ChannelInboundHandler {
5151
typealias InboundIn = AddressedEnvelope<DataType>
5252
typealias InboundOut = AddressedEnvelope<DataType>
5353

0 commit comments

Comments
 (0)