384 lines
16 KiB
Swift
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
|