Limited read and write to ByteBuffer.

Author: catap

CHANGELOG.md

@@ -5,6 +5,7 @@ All notable changes to this project will be documented in this file.
 ## [unreleased]
 - Introduced support direct allocated ByteBuffer.
 - Cross build for scala-2.10.7.
+- Limited read and write to ByteBuffer.
 
 ## [2.4.1] - 2021-02-19
 - Fixed regression on RFC4648 encoding introduced in 2.4.0.

shared/src/main/scala/ky/korins/blake3/Hasher.scala

@@ -49,17 +49,31 @@ trait Hasher {
    */
   def update(input: String): Hasher
 
+  /**
+   * Updates a hasher from specified InputStream, returns the same hasher
+   *
+   * It reads `input` until it returns `-1` or it reads len bytes
+   */
+  def update(input: InputStream, len: Int): Hasher
+
   /**
    * Updates a hasher from specified InputStream, returns the same hasher
    *
    * It reads `input` until it returns `-1`
    */
-  def update(input: InputStream): Hasher
+  def update(input: InputStream): Hasher =
+    update(input, Int.MaxValue)
+
+  /**
+   * Updates a hasher from specified ByteBuffer with no more than len bytes, returns the same hasher
+   */
+  def update(input: ByteBuffer, len: Int): Hasher
 
   /**
-   * Updates a hasher from specified ByteBuffer, returns the same hasher
+   * Updates a hasher from specified ByteBuffer until the end, returns the same hasher
    */
-  def update(input: ByteBuffer): Hasher
+  def update(input: ByteBuffer): Hasher =
+    update(input, input.remaining())
 
   /**
    * Calculate a hash into specified byte array
@@ -91,9 +105,15 @@ trait Hasher {
   def done(out: OutputStream, len: Int): Unit
 
   /**
-   * Calculate a hash into specified OutputStream with specified output length in bytes
+   * Calculate a hash into specified ByteBuffer with specified output length in bytes
+   */
+  def done(out: ByteBuffer, len: Int): Unit
+
+  /**
+   * Calculate a hash into specified ByteBuffer until the end of ByteBuffer
    */
-  def done(out: ByteBuffer): Unit
+  def done(out: ByteBuffer): Unit =
+    done(out, out.remaining())
 
   /**
    * Calculate a hash as single short

shared/src/main/scala/ky/korins/blake3/HasherImpl.scala

@@ -134,7 +134,7 @@ private[blake3] class HasherImpl(
     this
   }
 
-  def update(input: InputStream): Hasher = synchronized {
+  def update(input: InputStream, len: Int): Hasher = synchronized {
     val bytes = new Array[Byte](CHUNK_LEN)
 
     var consume = chunkState.len() match {
@@ -145,24 +145,28 @@ private[blake3] class HasherImpl(
     }
 
     var read = input.read(bytes, 0, consume)
-    while (read > 0) {
+    var remaining = len - read
+    while (remaining > 0 && read >= 0) {
       val len = finalizeWhenCompleted()
       chunkState.update(bytes, 0, read)
       consume = CHUNK_LEN - len
       read = input.read(bytes, 0, consume)
+      remaining -= read
     }
 
     this
   }
 
-  def update(input: ByteBuffer): Hasher = synchronized {
+  def update(input: ByteBuffer, len: Int): Hasher = synchronized {
     val bytes = new Array[Byte](CHUNK_LEN)
 
-    while (input.hasRemaining) {
-      val len = finalizeWhenCompleted()
-      val consume = Math.min(CHUNK_LEN - len, input.remaining())
+    var remaining = len
+    while (remaining > 0 && input.hasRemaining) {
+      val chunkLen = finalizeWhenCompleted()
+      val consume = Math.min(CHUNK_LEN - chunkLen, remaining)
       input.get(bytes, 0, consume)
       chunkState.update(bytes, 0, consume)
+      remaining -= consume
     }
 
     this
@@ -216,6 +220,6 @@ private[blake3] class HasherImpl(
   def done(out: OutputStream, len: Int): Unit =
     getOutput.rootBytes(out, len)
 
-  def done(out: ByteBuffer): Unit =
-    getOutput.rootBytes(out)
+  def done(out: ByteBuffer, len: Int): Unit =
+    getOutput.rootBytes(out, len)
 }

shared/src/main/scala/ky/korins/blake3/Output.scala

@@ -172,10 +172,9 @@ private[blake3] class Output(
     }
   }
 
-  def rootBytes(out: ByteBuffer): Unit = {
+  def rootBytes(out: ByteBuffer, len: Int): Unit = {
     var outputBlockCounter = 0
     var pos = 0
-    val len = out.remaining()
 
     val blockLenWords = BLOCK_LEN_WORDS
     val words = new Array[Int](blockLenWords)