Files
wehub-resource-sync 680845cb1c
Build containerization / Verify commit signatures (push) Has been skipped
Build containerization / containerization (push) Successful in 0s
Linux build / Linux compile check (push) Has been cancelled
Linux build / Determine Swift version (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:25:30 +08:00

384 lines
16 KiB
Swift

//===----------------------------------------------------------------------===//
// Copyright © 2026 Apple Inc. and the Containerization project authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//===----------------------------------------------------------------------===//
import ContainerizationError
import Foundation
import Logging
import NIOCore
import NIOPosix
#if os(macOS)
/// A minimal NBD server for integration testing.
///
/// Serves a file-backed block device using the NBD newstyle handshake protocol.
/// Supports both TCP and Unix domain socket transports.
final class NBDServer: Sendable {
private let channel: Channel
private let socketPath: String?
private let group: EventLoopGroup
let url: String
init(filePath: String, socketPath: String, logger: Logger? = nil) throws {
self.socketPath = socketPath
self.group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
try? FileManager.default.removeItem(atPath: socketPath)
self.channel = try Self.bootstrap(group: self.group, filePath: filePath, logger: logger)
.bind(unixDomainSocketPath: socketPath)
.wait()
self.url = "nbd+unix:///?socket=\(socketPath)"
}
init(filePath: String, port: Int, logger: Logger? = nil) throws {
self.socketPath = nil
self.group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
self.channel = try Self.bootstrap(group: self.group, filePath: filePath, logger: logger)
.bind(host: "127.0.0.1", port: port)
.wait()
guard let boundPort = channel.localAddress?.port, boundPort > 0 else {
throw ContainerizationError(.internalError, message: "NBD server failed to bind to a port")
}
self.url = "nbd://127.0.0.1:\(boundPort)"
}
func stop() {
try? channel.close().wait()
try? group.syncShutdownGracefully()
if let socketPath {
try? FileManager.default.removeItem(atPath: socketPath)
}
}
private static func bootstrap(group: EventLoopGroup, filePath: String, logger: Logger?) -> ServerBootstrap {
ServerBootstrap(group: group)
.serverChannelOption(.socketOption(.so_reuseaddr), value: 1)
.childChannelInitializer { channel in
channel.eventLoop.makeCompletedFuture {
try channel.pipeline.syncOperations.addHandler(
NBDConnectionHandler(filePath: filePath, logger: logger)
)
}
}
}
}
private final class NBDConnectionHandler: ChannelInboundHandler {
typealias InboundIn = ByteBuffer
typealias OutboundOut = ByteBuffer
// Protocol constants
static let magic: UInt64 = 0x4e42_444d_4147_4943
static let ihaveopt: UInt64 = 0x4948_4156_454f_5054
static let replyMagic: UInt64 = 0x3_e889_0455_65a9
static let requestMagic: UInt32 = 0x2560_9513
static let simpleReplyMagic: UInt32 = 0x6744_6698
static let optExportName: UInt32 = 1
static let optAbort: UInt32 = 2
static let optInfo: UInt32 = 6
static let optGo: UInt32 = 7
static let cmdRead: UInt16 = 0
static let cmdWrite: UInt16 = 1
static let cmdDisc: UInt16 = 2
static let cmdFlush: UInt16 = 3
static let flagFixedNewstyle: UInt16 = 0x1
static let flagNoZeroes: UInt16 = 0x2
static let clientFlagFixedNewstyle: UInt32 = 0x1
static let clientFlagNoZeroes: UInt32 = 0x2
static let transmitHasFlags: UInt16 = 0x1
static let transmitSendFlush: UInt16 = 0x4
static let transmitSendFUA: UInt16 = 0x8
static let repACK: UInt32 = 1
static let repInfo: UInt32 = 3
static let repErrUnsup: UInt32 = 0x8000_0001
static let infoExport: UInt16 = 0
static let infoBlockSize: UInt16 = 3
// NBD error codes
static let errOK: UInt32 = 0
static let errIO: UInt32 = 5
static let errNotsup: UInt32 = 95
private let fileFD: Int32
private let fileSize: UInt64
private let logger: Logger?
private var buffer: ByteBuffer = ByteBuffer()
private var state: ConnectionState = .handshake
private enum ConnectionState {
case handshake
case options(noZeroes: Bool)
case transmission
}
init(filePath: String, logger: Logger?) {
self.fileFD = open(filePath, O_RDWR)
self.logger = logger
guard fileFD >= 0 else {
self.fileSize = 0
logger?.error("NBD server: failed to open \(filePath), errno=\(errno)")
return
}
var st = stat()
if fstat(self.fileFD, &st) == 0 {
self.fileSize = UInt64(st.st_size)
} else {
self.fileSize = 0
}
}
func channelActive(context: ChannelHandlerContext) {
guard fileFD >= 0 else {
context.close(promise: nil)
return
}
// Send initial handshake.
var buf = context.channel.allocator.buffer(capacity: 18)
buf.writeInteger(Self.magic)
buf.writeInteger(Self.ihaveopt)
buf.writeInteger(Self.flagFixedNewstyle | Self.flagNoZeroes)
context.writeAndFlush(wrapOutboundOut(buf), promise: nil)
}
func channelInactive(context: ChannelHandlerContext) {
if fileFD >= 0 {
close(fileFD)
}
}
func channelRead(context: ChannelHandlerContext, data: NIOAny) {
var incoming = unwrapInboundIn(data)
buffer.writeBuffer(&incoming)
processBuffer(context: context)
}
private func processBuffer(context: ChannelHandlerContext) {
while true {
switch state {
case .handshake:
guard buffer.readableBytes >= 4,
let clientFlags = buffer.readInteger(as: UInt32.self)
else {
return
}
guard clientFlags & Self.clientFlagFixedNewstyle != 0 else {
context.close(promise: nil)
return
}
let noZeroes = clientFlags & Self.clientFlagNoZeroes != 0
state = .options(noZeroes: noZeroes)
case .options(let noZeroes):
guard buffer.readableBytes >= 16 else {
return
}
// Peek at the header without consuming.
let readerIndex = buffer.readerIndex
guard let magic = buffer.getInteger(at: readerIndex, as: UInt64.self),
let optType = buffer.getInteger(at: readerIndex + 8, as: UInt32.self),
let dataLen = buffer.getInteger(at: readerIndex + 12, as: UInt32.self)
else {
context.close(promise: nil)
return
}
// Wait until we have the full option data.
guard buffer.readableBytes >= 16 + Int(dataLen) else {
return
}
// Consume the header.
buffer.moveReaderIndex(forwardBy: 16)
guard magic == Self.ihaveopt else {
context.close(promise: nil)
return
}
let transmitFlags = Self.transmitHasFlags | Self.transmitSendFlush | Self.transmitSendFUA
switch optType {
case Self.optExportName:
if dataLen > 0 {
buffer.moveReaderIndex(forwardBy: Int(dataLen))
}
var reply = context.channel.allocator.buffer(capacity: 10)
reply.writeInteger(fileSize)
reply.writeInteger(transmitFlags)
if !noZeroes {
reply.writeRepeatingByte(0, count: 124)
}
context.writeAndFlush(wrapOutboundOut(reply), promise: nil)
state = .transmission
case Self.optInfo, Self.optGo:
// Parse InfoRequest to check for block size request.
var requestedBlockSize = false
if dataLen >= 6 {
let optDataStart = buffer.readerIndex
let nameLen = Int(buffer.getInteger(at: optDataStart, as: UInt32.self) ?? 0)
let infoOffset = optDataStart + 4 + nameLen
if infoOffset + 2 <= optDataStart + Int(dataLen) {
let numReqs = Int(buffer.getInteger(at: infoOffset, as: UInt16.self) ?? 0)
for i in 0..<numReqs {
let reqOffset = infoOffset + 2 + i * 2
if reqOffset + 2 <= optDataStart + Int(dataLen) {
let infoType = buffer.getInteger(at: reqOffset, as: UInt16.self) ?? 0
if infoType == Self.infoBlockSize {
requestedBlockSize = true
}
}
}
}
}
if dataLen > 0 {
buffer.moveReaderIndex(forwardBy: Int(dataLen))
}
// Send NBD_INFO_EXPORT reply.
var exportInfo = context.channel.allocator.buffer(capacity: 32)
writeOptReply(&exportInfo, optType: optType, replyType: Self.repInfo, dataLen: 12)
exportInfo.writeInteger(Self.infoExport)
exportInfo.writeInteger(fileSize)
exportInfo.writeInteger(transmitFlags)
// Send NBD_INFO_BLOCK_SIZE if requested.
if requestedBlockSize {
writeOptReply(&exportInfo, optType: optType, replyType: Self.repInfo, dataLen: 14)
exportInfo.writeInteger(Self.infoBlockSize)
exportInfo.writeInteger(UInt32(1)) // minimum
exportInfo.writeInteger(UInt32(4096)) // preferred
exportInfo.writeInteger(UInt32(4096 * 32)) // maximum
}
writeOptReply(&exportInfo, optType: optType, replyType: Self.repACK, dataLen: 0)
context.writeAndFlush(wrapOutboundOut(exportInfo), promise: nil)
if optType == Self.optGo {
state = .transmission
}
case Self.optAbort:
if dataLen > 0 {
buffer.moveReaderIndex(forwardBy: Int(dataLen))
}
context.close(promise: nil)
return
default:
if dataLen > 0 {
buffer.moveReaderIndex(forwardBy: Int(dataLen))
}
var reply = context.channel.allocator.buffer(capacity: 20)
writeOptReply(&reply, optType: optType, replyType: Self.repErrUnsup, dataLen: 0)
context.writeAndFlush(wrapOutboundOut(reply), promise: nil)
}
case .transmission:
// Request header: 4 magic + 2 flags + 2 type + 8 cookie + 8 offset + 4 length = 28
guard buffer.readableBytes >= 28 else {
return
}
let readerIndex = buffer.readerIndex
guard let magic = buffer.getInteger(at: readerIndex, as: UInt32.self),
let cmdType = buffer.getInteger(at: readerIndex + 6, as: UInt16.self),
let cookie = buffer.getInteger(at: readerIndex + 8, as: UInt64.self),
let offset = buffer.getInteger(at: readerIndex + 16, as: UInt64.self),
let length = buffer.getInteger(at: readerIndex + 24, as: UInt32.self)
else {
context.close(promise: nil)
return
}
guard magic == Self.requestMagic else {
context.close(promise: nil)
return
}
switch cmdType {
case Self.cmdWrite:
// Need the full write payload before processing.
guard buffer.readableBytes >= 28 + Int(length) else {
return
}
buffer.moveReaderIndex(forwardBy: 28)
var writeData = [UInt8](repeating: 0, count: Int(length))
buffer.readWithUnsafeReadableBytes { ptr in
writeData.withUnsafeMutableBytes { dst in
guard let dstBase = dst.baseAddress, let srcBase = ptr.baseAddress else {
return
}
_ = memcpy(dstBase, srcBase, Int(length))
}
return Int(length)
}
let n = pwrite(fileFD, &writeData, Int(length), off_t(offset))
var reply = context.channel.allocator.buffer(capacity: 16)
writeSimpleReply(&reply, cookie: cookie, error: n < 0 ? Self.errIO : Self.errOK)
context.writeAndFlush(wrapOutboundOut(reply), promise: nil)
case Self.cmdRead:
buffer.moveReaderIndex(forwardBy: 28)
var readBuf = [UInt8](repeating: 0, count: Int(length))
let n = pread(fileFD, &readBuf, Int(length), off_t(offset))
var reply = context.channel.allocator.buffer(capacity: 16 + Int(length))
writeSimpleReply(&reply, cookie: cookie, error: n < 0 ? Self.errIO : Self.errOK)
if n >= 0 {
reply.writeBytes(readBuf[0..<Int(length)])
}
context.writeAndFlush(wrapOutboundOut(reply), promise: nil)
case Self.cmdDisc:
buffer.moveReaderIndex(forwardBy: 28)
context.close(promise: nil)
return
case Self.cmdFlush:
buffer.moveReaderIndex(forwardBy: 28)
fsync(fileFD)
var reply = context.channel.allocator.buffer(capacity: 16)
writeSimpleReply(&reply, cookie: cookie, error: Self.errOK)
context.writeAndFlush(wrapOutboundOut(reply), promise: nil)
default:
buffer.moveReaderIndex(forwardBy: 28)
var reply = context.channel.allocator.buffer(capacity: 16)
writeSimpleReply(&reply, cookie: cookie, error: Self.errNotsup)
context.writeAndFlush(wrapOutboundOut(reply), promise: nil)
}
}
}
}
private func writeOptReply(_ buf: inout ByteBuffer, optType: UInt32, replyType: UInt32, dataLen: UInt32) {
buf.writeInteger(Self.replyMagic)
buf.writeInteger(optType)
buf.writeInteger(replyType)
buf.writeInteger(dataLen)
}
private func writeSimpleReply(_ buf: inout ByteBuffer, cookie: UInt64, error: UInt32) {
buf.writeInteger(Self.simpleReplyMagic)
buf.writeInteger(error)
buf.writeInteger(cookie)
}
}
#endif