Read and write null terminated strings (#1990)

Sources/NIOCore/ByteBuffer-aux.swift

@@ -90,6 +90,18 @@ extension ByteBuffer {
         self._moveWriterIndex(forwardBy: written)
         return written
     }
+
+    /// Write `string` into this `ByteBuffer` null terminated using UTF-8 encoding, moving the writer index forward appropriately.
+    ///
+    /// - parameters:
+    ///     - string: The string to write.
+    /// - returns: The number of bytes written.
+    @discardableResult
+    public mutating func writeNullTerminatedString(_ string: String) -> Int {
+        let written = self.setNullTerminatedString(string, at: self.writerIndex)
+        self._moveWriterIndex(forwardBy: written)
+        return written
+    }
 
     @inline(never)
     @usableFromInline
@@ -125,6 +137,18 @@ extension ByteBuffer {
             return self._setStringSlowpath(string, at: index)
         }
     }
+
+    /// Write `string` null terminated into this `ByteBuffer` at `index` using UTF-8 encoding. Does not move the writer index.
+    ///
+    /// - parameters:
+    ///     - string: The string to write.
+    ///     - index: The index for the first serialized byte.
+    /// - returns: The number of bytes written.
+    public mutating func setNullTerminatedString(_ string: String, at index: Int) -> Int {
+        let length = self.setString(string, at: index)
+        self.setInteger(UInt8(0), at: index &+ length)
+        return length &+ 1
+    }
 
     /// Get the string at `index` from this `ByteBuffer` decoding using the UTF-8 encoding. Does not move the reader index.
     /// The selected bytes must be readable or else `nil` will be returned.
@@ -143,6 +167,30 @@ extension ByteBuffer {
             return String(decoding: UnsafeRawBufferPointer(fastRebase: pointer[range]), as: Unicode.UTF8.self)
         }
     }
+
+    /// Get the string at `index` from this `ByteBuffer` decoding using the UTF-8 encoding. Does not move the reader index.
+    /// The selected bytes must be readable or else `nil` will be returned.
+    ///
+    /// - parameters:
+    ///     - index: The starting index into `ByteBuffer` containing the null terminated string of interest.
+    /// - returns: A `String` value deserialized from this `ByteBuffer` or `nil` if there isn't a complete null-terminated string,
+    ///            including null-terminator, in the readable bytes after `index` in the buffer
+    public func getNullTerminatedString(at index: Int) -> String? {
+        guard let stringLength = self.getNullTerminatedStringLength(at: index) else {
+            return nil
+        }
+        return self.getString(at: index, length: stringLength)
+    }
+
+    private func getNullTerminatedStringLength(at index: Int) -> Int? {
+        guard self.readerIndex <= index && index < self.writerIndex else {
+            return nil
+        }
+        guard let endIndex = self.readableBytesView[index...].firstIndex(of: 0) else {
+            return nil
+        }
+        return endIndex &- index
+    }
 
     /// Read `length` bytes off this `ByteBuffer`, decoding it as `String` using the UTF-8 encoding. Move the reader index forward by `length`.
     ///
@@ -156,6 +204,20 @@ extension ByteBuffer {
         self._moveReaderIndex(forwardBy: length)
         return result
     }
+
+    /// Read a null terminated string off this `ByteBuffer`, decoding it as `String` using the UTF-8 encoding. Move the reader index
+    /// forward by the string's length and its null terminator.
+    ///
+    /// - returns: A `String` value deserialized from this `ByteBuffer` or `nil` if there isn't a complete null-terminated string,
+    ///            including null-terminator, in the readable bytes of the buffer
+    public mutating func readNullTerminatedString() -> String? {
+        guard let stringLength = self.getNullTerminatedStringLength(at: self.readerIndex) else {
+            return nil
+        }
+        let result = self.readString(length: stringLength)
+        self.moveReaderIndex(forwardBy: 1) // move forward by null terminator
+        return result
+    }
 
     // MARK: Substring APIs
     /// Write `substring` into this `ByteBuffer` using UTF-8 encoding, moving the writer index forward appropriately.

Tests/NIOPosixTests/ByteBufferTest+XCTest.swift

@@ -37,6 +37,9 @@ extension ByteBufferTest {
                 ("testReadWrite", testReadWrite),
                 ("testStaticStringReadTests", testStaticStringReadTests),
                 ("testString", testString),
+                ("testNullTerminatedString", testNullTerminatedString),
+                ("testReadNullTerminatedStringWithoutNullTermination", testReadNullTerminatedStringWithoutNullTermination),
+                ("testGetNullTerminatedStringOutOfRangeTests", testGetNullTerminatedStringOutOfRangeTests),
                 ("testWriteSubstring", testWriteSubstring),
                 ("testSetSubstring", testSetSubstring),
                 ("testSliceEasy", testSliceEasy),

Tests/NIOPosixTests/ByteBufferTest.swift

@@ -204,6 +204,48 @@ class ByteBufferTest: XCTestCase {
         XCTAssertEqual("Hello", string)
     }
 
+    func testNullTerminatedString() {
+        let writtenHello = buf.writeNullTerminatedString("Hello")
+        XCTAssertEqual(writtenHello, 6)
+        XCTAssertEqual(buf.readableBytes, 6)
+
+        let writtenEmpty = buf.writeNullTerminatedString("")
+        XCTAssertEqual(writtenEmpty, 1)
+        XCTAssertEqual(buf.readableBytes, 7)
+
+        let writtenFoo = buf.writeNullTerminatedString("foo")
+        XCTAssertEqual(writtenFoo, 4)
+        XCTAssertEqual(buf.readableBytes, 11)
+
+        XCTAssertEqual(buf.getNullTerminatedString(at: 0), "Hello")
+        XCTAssertEqual(buf.getNullTerminatedString(at: 6), "")
+        XCTAssertEqual(buf.getNullTerminatedString(at: 7), "foo")
+
+        XCTAssertEqual(buf.readNullTerminatedString(), "Hello")
+        XCTAssertEqual(buf.readerIndex, 6)
+
+        XCTAssertEqual(buf.readNullTerminatedString(), "")
+        XCTAssertEqual(buf.readerIndex, 7)
+
+        XCTAssertEqual(buf.readNullTerminatedString(), "foo")
+        XCTAssertEqual(buf.readerIndex, 11)
+    }
+
+    func testReadNullTerminatedStringWithoutNullTermination() {
+        buf.writeString("Hello")
+        XCTAssertNil(buf.readNullTerminatedString())
+    }
+
+    func testGetNullTerminatedStringOutOfRangeTests() {
+        buf.writeNullTerminatedString("Hello")
+        XCTAssertNil(buf.getNullTerminatedString(at: 100))
+        buf.moveReaderIndex(forwardBy: 6)
+        XCTAssertNil(buf.readNullTerminatedString())
+        XCTAssertNil(buf.getNullTerminatedString(at: 0))
+        buf.writeInteger(UInt8(0))
+        XCTAssertEqual(buf.readNullTerminatedString(), "")
+    }
+
     func testWriteSubstring() {
         var text = "Hello"
         let written = buf.writeSubstring(text[...])