From 8f168ad2e933fcbf3598911ef71db4c3dba10947 Mon Sep 17 00:00:00 2001
From: Lennard Berger <lennard.berger@student.uni-tuebingen.de>
Date: Sun, 30 Sep 2018 22:29:18 +0200
Subject: [PATCH] [brotli] Add brotli support backed by org.brotli.dec

---
 browsermob-core/pom.xml                       |  7 ++++
 .../filters/ServerResponseCaptureFilter.java  |  8 ++++
 .../bmp/util/BrowserMobHttpUtil.java          | 40 ++++++++++++++++++-
 3 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/browsermob-core/pom.xml b/browsermob-core/pom.xml
index 2071434ee..8e5cdab71 100644
--- a/browsermob-core/pom.xml
+++ b/browsermob-core/pom.xml
@@ -139,6 +139,13 @@
             <version>${project.version}</version>
         </dependency>
 
+        <!-- use google brotli for brotli decoding-->
+        <dependency>
+            <groupId>org.brotli</groupId>
+            <artifactId>dec</artifactId>
+            <version>0.1.2</version>
+        </dependency>
+
         <!-- Netty will use javassist to improve performance if it is present on the classpath, but it is not required -->
         <dependency>
             <groupId>org.javassist</groupId>
diff --git a/browsermob-core/src/main/java/net/lightbody/bmp/filters/ServerResponseCaptureFilter.java b/browsermob-core/src/main/java/net/lightbody/bmp/filters/ServerResponseCaptureFilter.java
index d69ad8f76..1bce61ad7 100644
--- a/browsermob-core/src/main/java/net/lightbody/bmp/filters/ServerResponseCaptureFilter.java
+++ b/browsermob-core/src/main/java/net/lightbody/bmp/filters/ServerResponseCaptureFilter.java
@@ -25,6 +25,7 @@
  */
 public class ServerResponseCaptureFilter extends HttpFiltersAdapter {
     private static final Logger log = LoggerFactory.getLogger(ServerResponseCaptureFilter.class);
+    private static final String BROTLI_COMPRESSION = "br";
 
     /**
      * Populated by serverToProxyResponse() when processing the HttpResponse object
@@ -133,6 +134,13 @@ protected void decompressContents() {
             } catch (RuntimeException e) {
                 log.warn("Failed to decompress response with encoding type " + contentEncoding + " when decoding request from " + originalRequest.getUri(), e);
             }
+        } else if(contentEncoding.equals(BROTLI_COMPRESSION)) {
+            try {
+                fullResponseContents = BrowserMobHttpUtil.decompressBrotliContents(getRawResponseContents());
+                decompressionSuccessful = true;
+            } catch (RuntimeException e) {
+                log.warn("Failed to decompress response with encoding type " + contentEncoding + " when decoding request from " + originalRequest.getUri(), e);
+            }
         } else {
             log.warn("Cannot decode unsupported content encoding type {}", contentEncoding);
         }
diff --git a/browsermob-core/src/main/java/net/lightbody/bmp/util/BrowserMobHttpUtil.java b/browsermob-core/src/main/java/net/lightbody/bmp/util/BrowserMobHttpUtil.java
index 98172810f..264b3a2ee 100644
--- a/browsermob-core/src/main/java/net/lightbody/bmp/util/BrowserMobHttpUtil.java
+++ b/browsermob-core/src/main/java/net/lightbody/bmp/util/BrowserMobHttpUtil.java
@@ -3,6 +3,7 @@
 import com.google.common.io.BaseEncoding;
 import com.google.common.net.HostAndPort;
 import com.google.common.net.MediaType;
+import org.brotli.dec.BrotliInputStream;
 import io.netty.buffer.ByteBuf;
 import io.netty.handler.codec.http.HttpHeaders;
 import io.netty.handler.codec.http.HttpRequest;
@@ -14,6 +15,7 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
+import java.io.InputStream;
 import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -78,7 +80,7 @@ public static long getHeaderSize(HttpHeaders headers) {
     /**
      * Decompresses the gzipped byte stream.
      *
-     * @param fullMessage gzipped byte stream to decomress
+     * @param fullMessage gzipped byte stream to decompress
      * @return decompressed bytes
      * @throws DecompressionException thrown if the fullMessage cannot be read or decompressed for any reason
      */
@@ -111,6 +113,42 @@ public static byte[] decompressContents(byte[] fullMessage) throws Decompression
         return fullMessage;
     }
 
+    /**
+     * Decompresses the brotli byze stream
+     *
+     * @param fullMessage brotli byte stream to decompress
+     * @return decompressed bytes
+    * @throws DecompressionException thrown if the fullMessage cannot be read or decompressed for any reason
+     */
+    public static byte[] decompressBrotliContents(byte[] fullMessage) throws DecompressionException {
+        InputStream brotliReader = null;
+        ByteArrayOutputStream uncompressed;
+        try {
+            brotliReader = new BrotliInputStream(new ByteArrayInputStream(fullMessage));
+
+            uncompressed = new ByteArrayOutputStream(fullMessage.length);
+
+            byte[] decompressBuffer = new byte[DECOMPRESS_BUFFER_SIZE];
+            int bytesRead;
+            while ((bytesRead = brotliReader.read(decompressBuffer)) > -1) {
+                uncompressed.write(decompressBuffer, 0, bytesRead);
+            }
+
+            fullMessage = uncompressed.toByteArray();
+        } catch (IOException e) {
+            throw new DecompressionException("Unable to decompress response", e);
+        } finally {
+            try {
+                if (brotliReader != null) {
+                    brotliReader.close();
+                }
+            } catch (IOException e) {
+                log.warn("Unable to close brotli stream", e);
+            }
+        }
+        return fullMessage;
+    }
+
     /**
      * Returns true if the content type string indicates textual content. Currently these are any Content-Types that start with one of the
      * following: