diff --git a/actor-tests/src/test/scala/org/apache/pekko/util/ByteStringSpec.scala b/actor-tests/src/test/scala/org/apache/pekko/util/ByteStringSpec.scala index e7d15c7541d..67009e9a1ae 100644 --- a/actor-tests/src/test/scala/org/apache/pekko/util/ByteStringSpec.scala +++ b/actor-tests/src/test/scala/org/apache/pekko/util/ByteStringSpec.scala @@ -1744,6 +1744,137 @@ class ByteStringSpec extends AnyWordSpec with Matchers with Checkers { } } + "ByteString.copyToArray with a source offset" must { + // Same content in every internal representation, so the new overload is exercised against + // the compacted, sliced, two-fragment and multi-fragment layouts. + def representations(bytes: Array[Byte]): List[(String, ByteString)] = { + val n = bytes.length + val whole = ByteString(bytes) + var result = List("compact" -> whole.compact, "whole" -> whole) + if (n >= 1) { + val padded = ByteString(Array[Byte](0x7F, 0x7F)) ++ whole ++ ByteString(Array[Byte](0x7F)) + result ::= "sliced" -> padded.drop(2).dropRight(1) + result ::= "per-byte rope" -> bytes.map(b => ByteString(Array(b))).reduce(_ ++ _) + } + if (n >= 2) result ::= "two fragments" -> (ByteString(bytes.take(1)) ++ ByteString(bytes.drop(1))) + if (n >= 4) { + val q = n / 4 + result ::= "four fragments" -> (ByteString(bytes.slice(0, q)) ++ ByteString(bytes.slice(q, 2 * q)) ++ + ByteString(bytes.slice(2 * q, 3 * q)) ++ ByteString(bytes.slice(3 * q, n))) + } + result + } + + def sample(n: Int): Array[Byte] = Array.tabulate[Byte](n)(i => ((i * 31 + 7) % 251).toByte) + + "copy the same bytes as slice(...).copyToArray, for every offset and length" in { + for (n <- List(0, 1, 2, 3, 8, 16, 17, 64)) { + val bytes = sample(n) + for { + (name, bs) <- representations(bytes) + srcOffset <- -1 to n + destOffset <- List(0, 1, 3) + len <- List(0, 1, 3, n, n + 5) + } withClue(s"size $n, $name, srcOffset $srcOffset, destOffset $destOffset, len $len: ") { + val destSize = n + 8 + val actual = Array.fill[Byte](destSize)(0x7F.toByte) + val expected = Array.fill[Byte](destSize)(0x7F.toByte) + val expectedCopied = + if (srcOffset < 0 || len <= 0 || srcOffset >= n) 0 + else math.max(0, math.min(math.min(len, n - srcOffset), destSize - destOffset)) + if (expectedCopied > 0) System.arraycopy(bytes, srcOffset, expected, destOffset, expectedCopied) + + bs.copyToArray(srcOffset, actual, destOffset, len) should ===(expectedCopied) + actual should ===(expected) + } + } + } + + "clamp out of range arguments instead of throwing" in { + val bs = ByteString(sample(8)) + val dest = new Array[Byte](8) + bs.copyToArray(-1, dest, 0, 4) should ===(0) + bs.copyToArray(0, dest, -1, 4) should ===(0) + bs.copyToArray(8, dest, 0, 4) should ===(0) + bs.copyToArray(100, dest, 0, 4) should ===(0) + bs.copyToArray(0, dest, 0, 0) should ===(0) + bs.copyToArray(0, dest, 0, -1) should ===(0) + bs.copyToArray(0, dest, 8, 4) should ===(0) + } + + "not write past the end of the destination" in { + val bs = ByteString(sample(64)) + val dest = Array.fill[Byte](10)(0x7F.toByte) + bs.copyToArray(0, dest, 6, 64) should ===(4) + dest.take(6) should ===(Array.fill[Byte](6)(0x7F.toByte)) + } + + "leave an empty ByteString as a no-op" in { + val dest = Array.fill[Byte](4)(0x7F.toByte) + ByteString.empty.copyToArray(0, dest, 0, 4) should ===(0) + dest should ===(Array.fill[Byte](4)(0x7F.toByte)) + } + } + + "ByteString.decodeString over a range" must { + def representations(bytes: Array[Byte]): List[(String, ByteString)] = { + val n = bytes.length + val whole = ByteString(bytes) + var result = List("compact" -> whole.compact, "whole" -> whole) + if (n >= 1) { + val padded = ByteString(Array[Byte](0x7F, 0x7F)) ++ whole ++ ByteString(Array[Byte](0x7F)) + result ::= "sliced" -> padded.drop(2).dropRight(1) + result ::= "per-byte rope" -> bytes.map(b => ByteString(Array(b))).reduce(_ ++ _) + } + if (n >= 2) result ::= "two fragments" -> (ByteString(bytes.take(1)) ++ ByteString(bytes.drop(1))) + result + } + + val texts = List("", "a", "hello world", "Content-Type: application/json", "héllo wörld", "\u00e9\u00e8\u00ea") + + "decode the same string as slice(...).decodeString" in { + for (text <- texts) { + val bytes = text.getBytes(StandardCharsets.UTF_8) + val n = bytes.length + for { + (name, bs) <- representations(bytes) + from <- -1 to n + until <- -1 to n + 1 + } withClue(s"[$text], $name, from $from, until $until: ") { + val start = math.max(0, from) + val end = math.min(n, until) + val expected = if (start >= end) "" else new String(bytes, start, end - start, StandardCharsets.UTF_8) + bs.decodeString(StandardCharsets.UTF_8, from, until) should ===(expected) + bs.decodeString(StandardCharsets.UTF_8, from, until) should ===(bs.slice(start, end).decodeString( + StandardCharsets.UTF_8)) + } + } + } + + "return the empty string for an empty or inverted range" in { + val bs = ByteString("hello".getBytes(StandardCharsets.UTF_8)) + bs.decodeString(StandardCharsets.UTF_8, 2, 2) should ===("") + bs.decodeString(StandardCharsets.UTF_8, 3, 1) should ===("") + bs.decodeString(StandardCharsets.UTF_8, 10, 20) should ===("") + bs.decodeString(StandardCharsets.UTF_8, -5, -1) should ===("") + ByteString.empty.decodeString(StandardCharsets.UTF_8, 0, 1) should ===("") + } + + "clamp a range that runs past the end" in { + val bs = ByteString("hello".getBytes(StandardCharsets.UTF_8)) + bs.decodeString(StandardCharsets.UTF_8, 3, 99) should ===("lo") + bs.decodeString(StandardCharsets.UTF_8, -3, 2) should ===("he") + } + + "decode a multi-byte character that spans two fragments" in { + // 'é' is two bytes in UTF-8; split the ByteString between them + val bytes = "aéb".getBytes(StandardCharsets.UTF_8) + val split = ByteString(bytes.take(2)) ++ ByteString(bytes.drop(2)) + split.decodeString(StandardCharsets.UTF_8, 0, bytes.length) should ===("aéb") + split.decodeString(StandardCharsets.UTF_8, 1, 3) should ===("é") + } + } + "A ByteString" must { "have correct size" when { "concatenating" in { check((a: ByteString, b: ByteString) => (a ++ b).size == a.size + b.size) } diff --git a/actor/src/main/scala/org/apache/pekko/util/ByteString.scala b/actor/src/main/scala/org/apache/pekko/util/ByteString.scala index af10c991a99..e6805977093 100644 --- a/actor/src/main/scala/org/apache/pekko/util/ByteString.scala +++ b/actor/src/main/scala/org/apache/pekko/util/ByteString.scala @@ -469,6 +469,15 @@ object ByteString { toCopy } + override def copyToArray(srcOffset: Int, dest: Array[Byte], destOffset: Int, len: Int): Int = { + val toCopy = copyableBytes(srcOffset, dest.length, destOffset, len) + if (toCopy > 0) System.arraycopy(bytes, srcOffset, dest, destOffset, toCopy) + toCopy + } + + private[pekko] override def decodeStringUnchecked(charset: Charset, offset: Int, len: Int): String = + new String(bytes, offset, len, charset) + override def toArrayUnsafe(): Array[Byte] = bytes override def asInputStream: InputStream = new UnsynchronizedByteArrayInputStream(bytes) @@ -856,6 +865,15 @@ object ByteString { toCopy } + override def copyToArray(srcOffset: Int, dest: Array[Byte], destOffset: Int, len: Int): Int = { + val toCopy = copyableBytes(srcOffset, dest.length, destOffset, len) + if (toCopy > 0) System.arraycopy(bytes, startIndex + srcOffset, dest, destOffset, toCopy) + toCopy + } + + private[pekko] override def decodeStringUnchecked(charset: Charset, offset: Int, len: Int): String = + new String(bytes, startIndex + offset, len, charset) + protected def writeReplace(): AnyRef = new SerializationProxy(this) override def toArrayUnsafe(): Array[Byte] = { @@ -1284,6 +1302,23 @@ object ByteString { totalToCopy } + override def copyToArray(srcOffset: Int, dest: Array[Byte], destOffset: Int, len: Int): Int = { + val toCopy = copyableBytes(srcOffset, dest.length, destOffset, len) + if (toCopy <= 0) 0 + else { + val firstLength = first.length + var copied = 0 + if (srcOffset < firstLength) { + copied = first.copyToArray(srcOffset, dest, destOffset, toCopy) + } + if (copied < toCopy) { + val secondOffset = math.max(0, srcOffset - firstLength) + second.copyToArray(secondOffset, dest, destOffset + copied, toCopy - copied) + } + toCopy + } + } + override def foreach[@specialized U](f: Byte => U): Unit = { first.foreach(f) second.foreach(f) @@ -1863,6 +1898,31 @@ object ByteString { ByteString1C(result) } + override def copyToArray(srcOffset: Int, dest: Array[Byte], destOffset: Int, len: Int): Int = { + val toCopy = copyableBytes(srcOffset, dest.length, destOffset, len) + if (toCopy <= 0) 0 + else { + // Locate the fragment holding srcOffset, then copy whole runs out of consecutive + // fragments. Indexing rather than iterating keeps this allocation free. + var fragIdx = 0 + var fragOffset = srcOffset + while (fragOffset >= bytestrings(fragIdx).length) { + fragOffset -= bytestrings(fragIdx).length + fragIdx += 1 + } + var copied = 0 + while (copied < toCopy) { + val frag = bytestrings(fragIdx) + val fromFrag = math.min(toCopy - copied, frag.length - fragOffset) + frag.copyToArray(fragOffset, dest, destOffset + copied, fromFrag) + copied += fromFrag + fragIdx += 1 + fragOffset = 0 + } + toCopy + } + } + override def copyToArray[B >: Byte](dest: Array[B], start: Int, len: Int): Int = { if (bytestrings.size == 1) bytestrings.head.copyToArray(dest, start, len) else { @@ -2402,6 +2462,45 @@ object ByteString { override def copyToArray[B >: Byte](xs: Array[B], start: Int, len: Int): Int = throw new UnsupportedOperationException("Method copyToArray is not implemented in ByteString") + /** + * Copies `len` bytes of this ByteString, starting at `srcOffset`, into `dest` starting at + * `destOffset`. + * + * `copyToArray(dest, destOffset, len)` can only copy from the start of this ByteString, so + * copying from an offset otherwise requires slicing first, which allocates. This overload + * copies directly out of the backing array(s). + * + * Out of range values are clamped: the number of bytes copied is the largest value that fits in + * both this ByteString from `srcOffset` and `dest` from `destOffset`, and is never negative. + * + * @param srcOffset the index in this ByteString to start copying from + * @param dest the array to copy into + * @param destOffset the index in `dest` to start copying to + * @param len the maximum number of bytes to copy + * @return the number of bytes actually copied + * @since 2.0.0 + */ + def copyToArray(srcOffset: Int, dest: Array[Byte], destOffset: Int, len: Int): Int = { + val toCopy = copyableBytes(srcOffset, dest.length, destOffset, len) + if (toCopy <= 0) 0 + else { + var i = 0 + while (i < toCopy) { + dest(destOffset + i) = byteAtUnchecked(srcOffset + i) + i += 1 + } + toCopy + } + } + + /** + * INTERNAL API: the number of bytes `copyToArray(srcOffset, dest, destOffset, len)` may copy, + * clamped to what is available in this ByteString and what fits in the destination. + */ + private[pekko] final def copyableBytes(srcOffset: Int, destLength: Int, destOffset: Int, len: Int): Int = + if (srcOffset < 0 || destOffset < 0 || len <= 0 || srcOffset >= length) 0 + else math.max(0, math.min(math.min(len, length - srcOffset), destLength - destOffset)) + /** * Unsafe API: Use only in situations you are completely confident that this is what * you need, and that you understand the implications documented below. @@ -2525,6 +2624,38 @@ object ByteString { */ def decodeString(charset: Charset): String + /** + * Decodes the bytes in `[from, until)` using the given charset. + * + * Equivalent to `slice(from, until).decodeString(charset)` but without allocating the + * intermediate ByteString, which matters for parsers that repeatedly decode small ranges of a + * larger buffer. + * + * The range is clamped to the bounds of this ByteString, so an empty or inverted range decodes + * to the empty string rather than throwing. + * + * @param charset the charset to decode with + * @param from the index to start decoding at, inclusive + * @param until the index to stop decoding at, exclusive + * @since 2.0.0 + */ + def decodeString(charset: Charset, from: Int, until: Int): String = { + val start = math.max(0, from) + val end = math.min(length, until) + if (start >= end) "" + else decodeStringUnchecked(charset, start, end - start) + } + + /** + * INTERNAL API: decode `len` bytes starting at `offset`, with the range already validated. + * Overridden by the array-backed layouts so no intermediate ByteString or array is allocated. + */ + private[pekko] def decodeStringUnchecked(charset: Charset, offset: Int, len: Int): String = { + val bytes = new Array[Byte](len) + copyToArray(offset, bytes, 0, len) + new String(bytes, charset) + } + /** * Returns a ByteString which is the binary representation of this ByteString * if this ByteString is Base64-encoded.