Handle nullable original in TeeOutputStream

zalando · Nov 3, 2024 · c9e56c6 · c9e56c6
1 parent 337cc5e
commit c9e56c6
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 12 deletions.
diff --git a/logbook-jaxrs/src/main/java/org/zalando/logbook/jaxrs/TeeOutputStream.java b/logbook-jaxrs/src/main/java/org/zalando/logbook/jaxrs/TeeOutputStream.java
@@ -1,5 +1,6 @@
 package org.zalando.logbook.jaxrs;
 
+import javax.annotation.Nullable;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
@@ -12,31 +13,35 @@ final class TeeOutputStream extends OutputStream {
     private final OutputStream original;
     private final ByteArrayOutputStream copy = new ByteArrayOutputStream();
 
-    TeeOutputStream(final OutputStream original) {
+    TeeOutputStream(@Nullable final OutputStream original) {
         this.original = original;
     }
 
     @Override
     public void write(final int b) throws IOException {
-        original.write(b);
+        if (original != null)
+            original.write(b);
         copy.write(b);
     }
 
     @Override
     public void write(final byte[] b, final int off, final int len) throws IOException {
-        original.write(b, off, len);
+        if (original != null)
+            original.write(b, off, len);
         copy.write(b, off, len);
     }
 
     @Override
     public void flush() throws IOException {
-        original.flush();
+        if (original != null)
+            original.flush();
         copy.flush();
     }
 
     @Override
     public void close() throws IOException {
-        original.close();
+        if (original != null)
+            original.close();
         copy.close();
     }
 

diff --git a/logbook-jaxrs/src/test/java/org/zalando/logbook/jaxrs/TeeOutputStreamTest.java b/logbook-jaxrs/src/test/java/org/zalando/logbook/jaxrs/TeeOutputStreamTest.java
@@ -10,28 +10,35 @@
 class TeeOutputStreamTest {
 
     private final ByteArrayOutputStream output = new ByteArrayOutputStream();
-    private final TeeOutputStream unit = new TeeOutputStream(output);
+    private final TeeOutputStream unit1 = new TeeOutputStream(output);
+    private final TeeOutputStream unit2 = new TeeOutputStream(null);
 
     @Test
     void shouldWriteByte() throws IOException {
-        unit.write(17);
+        unit1.write(17);
+        unit2.write(17);
 
-        assertArrayEquals(unit.toByteArray(), output.toByteArray());
+        assertArrayEquals(unit1.toByteArray(), output.toByteArray());
+        assertArrayEquals(unit2.toByteArray(), output.toByteArray());
     }
 
     @Test
     void shouldWriteBytesWithoutOffsets() throws IOException {
-        unit.write(new byte[]{17});
+        unit1.write(new byte[]{17});
+        unit2.write(new byte[]{17});
 
-        assertArrayEquals(unit.toByteArray(), output.toByteArray());
+        assertArrayEquals(unit1.toByteArray(), output.toByteArray());
+        assertArrayEquals(unit2.toByteArray(), output.toByteArray());
 
     }
 
     @Test
     void shouldWriteBytesWithOffsets() throws IOException {
-        unit.write(new byte[]{17}, 0, 1);
+        unit1.write(new byte[]{17}, 0, 1);
+        unit2.write(new byte[]{17}, 0, 1);
 
-        assertArrayEquals(unit.toByteArray(), output.toByteArray());
+        assertArrayEquals(unit1.toByteArray(), output.toByteArray());
+        assertArrayEquals(unit2.toByteArray(), output.toByteArray());
 
     }