From cfef67da634996f09e5f2509e198cc73c88ce8b2 Mon Sep 17 00:00:00 2001
From: Benjamin Peterson <benjamin@engflow.com>
Date: Thu, 16 Mar 2023 07:44:55 -0700
Subject: [PATCH] Fix seeking of empty chunkers.

Empty chunkers are strange in that they emit one empty chunk and then end. This change makes them properly resetable. Previously, reset() on a empty chunker could result in a NullPointerException.

Also, it's important to call close() on the underlying data stream even if it's empty.

Closes #17797.

PiperOrigin-RevId: 517119206
Change-Id: Iff7908d6cd0633aa2a355ea89f8e647a9fefffcd
---
 .../devtools/build/lib/remote/Chunker.java    |  6 ++--
 .../build/lib/remote/ChunkerTest.java         | 29 +++++++++++++++++--
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/remote/Chunker.java b/src/main/java/com/google/devtools/build/lib/remote/Chunker.java
index ce7372a1c18185..3ef5d31d4175ca 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/Chunker.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/Chunker.java
@@ -159,14 +159,14 @@ public void reset() throws IOException {
   public void seek(long toOffset) throws IOException {
     // For compressed stream, we need to reinitialize the stream since the offset refers to the
     // uncompressed form.
-    if (initialized && toOffset >= offset && !compressed) {
+    if (initialized && size > 0 && toOffset >= offset && !compressed) {
       ByteStreams.skipFully(data, toOffset - offset);
       offset = toOffset;
     } else {
       reset();
       initialize(toOffset);
     }
-    if (data.finished()) {
+    if (size > 0 && data.finished()) {
       close();
     }
   }
@@ -216,7 +216,7 @@ public Chunk next() throws IOException {
     maybeInitialize();
 
     if (size == 0) {
-      data = null;
+      close();
       return emptyChunk;
     }
 
diff --git a/src/test/java/com/google/devtools/build/lib/remote/ChunkerTest.java b/src/test/java/com/google/devtools/build/lib/remote/ChunkerTest.java
index cd4fa6943034c3..7c693480df1482 100644
--- a/src/test/java/com/google/devtools/build/lib/remote/ChunkerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/remote/ChunkerTest.java
@@ -84,8 +84,17 @@ public void nextShouldThrowIfNoMoreData() throws IOException {
 
   @Test
   public void emptyData() throws Exception {
-    byte[] data = new byte[0];
-    Chunker chunker = Chunker.builder().setInput(data).build();
+    var inp =
+        new ByteArrayInputStream(new byte[0]) {
+          private boolean closed;
+
+          @Override
+          public void close() throws IOException {
+            closed = true;
+            super.close();
+          }
+        };
+    Chunker chunker = Chunker.builder().setInput(0, inp).build();
 
     assertThat(chunker.hasNext()).isTrue();
 
@@ -96,6 +105,7 @@ public void emptyData() throws Exception {
     assertThat(next.getOffset()).isEqualTo(0);
 
     assertThat(chunker.hasNext()).isFalse();
+    assertThat(inp.closed).isTrue();
 
     assertThrows(NoSuchElementException.class, () -> chunker.next());
   }
@@ -193,6 +203,21 @@ public void seekForwards() throws IOException {
     assertThat(chunker.hasNext()).isFalse();
   }
 
+  @Test
+  public void seekEmptyData() throws IOException {
+    var chunker = Chunker.builder().setInput(new byte[0]).build();
+    for (var i = 0; i < 2; i++) {
+      chunker.seek(0);
+      var next = chunker.next();
+      assertThat(next).isNotNull();
+      assertThat(next.getData()).isEmpty();
+      assertThat(next.getOffset()).isEqualTo(0);
+
+      assertThat(chunker.hasNext()).isFalse();
+      assertThrows(NoSuchElementException.class, chunker::next);
+    }
+  }
+
   @Test
   public void testSingleChunkCompressed() throws IOException {
     byte[] data = {72, 101, 108, 108, 111, 32, 87, 111, 114, 108, 100, 33};