feat: Support zstd encoding for downloads (#1287)

congminh1254 · web-flow · commit 0e3c4c07e65e · 2025-02-19T14:50:45.000+01:00
diff --git a/README.md b/README.md
@@ -207,6 +207,9 @@ public final class BoxDeveloperEditionAPIConnectionAsEnterpriseUser {
 
 The Box Java SDK is compatible with Java 8 and up.
 
+## Compression Support
+The SDK supports both gzip and zstd compression for API requests. Compression is handled automatically based on server capabilities.
+
 ## Building
 
 The SDK uses Gradle for its build system. SDK comes with Gradle wrapper. Running `./gradlew build` from the root
diff --git a/build.gradle b/build.gradle
@@ -54,6 +54,7 @@ dependencies {
     implementation "org.bouncycastle:bcprov-jdk18on:1.78.1"
     implementation "org.bouncycastle:bcpkix-jdk18on:1.78.1"
     implementation "com.squareup.okhttp3:okhttp:4.12.0"
+    implementation "com.github.luben:zstd-jni:1.5.5-5"
     testsCommonImplementation "junit:junit:4.13.2"
     testsCommonImplementation "org.hamcrest:hamcrest-library:2.2"
     testsCommonImplementation "org.mockito:mockito-core:4.8.0"
diff --git a/src/intTest/java/com/box/sdk/BoxFileIT.java b/src/intTest/java/com/box/sdk/BoxFileIT.java
@@ -1,7 +1,9 @@
 package com.box.sdk;
 
+import static com.box.sdk.BinaryBodyUtils.writeStream;
 import static com.box.sdk.BoxApiProvider.jwtApiForServiceAccount;
 import static com.box.sdk.BoxFile.ALL_VERSION_FIELDS;
+import static com.box.sdk.BoxFile.CONTENT_URL_TEMPLATE;
 import static com.box.sdk.BoxRetentionPolicyAssignment.createAssignmentToFolder;
 import static com.box.sdk.BoxSharedLink.Access.OPEN;
 import static com.box.sdk.CleanupTools.deleteFile;
@@ -225,6 +227,64 @@ public void uploadAndDownloadFileSucceeds() throws IOException {
 
     }
 
+    @Test
+    public void downloadFileUseZstdSucceeds() throws IOException {
+        BoxAPIConnection api = jwtApiForServiceAccount();
+        api.setUseZstdCompression(true);
+
+        String fileName = "smalltest.pdf";
+        URL fileURL = this.getClass().getResource("/sample-files/" + fileName);
+        String filePath = URLDecoder.decode(fileURL.getFile(), "utf-8");
+        byte[] fileContent = readAllBytes(filePath);
+        BoxFile file = null;
+        try {
+            file = uploadSampleFileToUniqueFolder(api, fileName);
+
+            ByteArrayOutputStream downloadStream = new ByteArrayOutputStream();
+            ProgressListener mockDownloadListener = mock(ProgressListener.class);
+
+            URL url = CONTENT_URL_TEMPLATE.build(api.getBaseURL(), file.getID());
+            BoxAPIRequest request = new BoxAPIRequest(api, url, "GET");
+            BoxAPIResponse response = request.send();
+            writeStream(response, downloadStream, mockDownloadListener);
+
+            byte[] downloadedFileContent = downloadStream.toByteArray();
+            assertThat(response.getHeaders().get("X-Content-Encoding").get(0), is(equalTo("zstd")));
+            assertThat(downloadedFileContent, is(equalTo(fileContent)));
+        } finally {
+            deleteFile(file);
+        }
+    }
+
+    @Test
+    public void uploadAndDownloadFileDisabledZstdSucceeds() throws IOException {
+        BoxAPIConnection api = jwtApiForServiceAccount();
+        api.setUseZstdCompression(false);
+
+        String fileName = "smalltest.pdf";
+        URL fileURL = this.getClass().getResource("/sample-files/" + fileName);
+        String filePath = URLDecoder.decode(fileURL.getFile(), "utf-8");
+        byte[] fileContent = readAllBytes(filePath);
+        BoxFile file = null;
+        try {
+            file = uploadSampleFileToUniqueFolder(api, fileName);
+
+            ByteArrayOutputStream downloadStream = new ByteArrayOutputStream();
+            ProgressListener mockDownloadListener = mock(ProgressListener.class);
+
+            URL url = CONTENT_URL_TEMPLATE.build(api.getBaseURL(), file.getID());
+            BoxAPIRequest request = new BoxAPIRequest(api, url, "GET");
+            BoxAPIResponse response = request.send();
+            writeStream(response, downloadStream, mockDownloadListener);
+
+            byte[] downloadedFileContent = downloadStream.toByteArray();
+            assertThat(response.getHeaders().get("X-Content-Encoding"), is(nullValue()));
+            assertThat(downloadedFileContent, is(equalTo(fileContent)));
+        } finally {
+            deleteFile(file);
+        }
+    }
+
     @Test
     public void downloadFileRangeSucceeds() throws IOException {
         BoxAPIConnection api = jwtApiForServiceAccount();
diff --git a/src/main/java/com/box/sdk/BoxAPIConnection.java b/src/main/java/com/box/sdk/BoxAPIConnection.java
@@ -122,6 +122,7 @@ public class BoxAPIConnection {
     private int maxRetryAttempts;
     private int connectTimeout;
     private int readTimeout;
+    private boolean useZstdCompression;
     private final List<BoxAPIConnectionListener> listeners;
     private RequestInterceptor interceptor;
     private final Map<String, String> customHeaders;
@@ -160,6 +161,7 @@ public BoxAPIConnection(String clientID, String clientSecret, String accessToken
         this.maxRetryAttempts = BoxGlobalSettings.getMaxRetryAttempts();
         this.connectTimeout = BoxGlobalSettings.getConnectTimeout();
         this.readTimeout = BoxGlobalSettings.getReadTimeout();
+        this.useZstdCompression = BoxGlobalSettings.getUseZstdCompression();
         this.refreshLock = new ReentrantReadWriteLock();
         this.userAgent = "Box Java SDK v" + SDK_VERSION + " (Java " + JAVA_VERSION + ")";
         this.listeners = new ArrayList<>();
@@ -237,6 +239,9 @@ private void buildHttpClients() {
             }
         }
         builder = modifyHttpClientBuilder(builder);
+        if (this.useZstdCompression) {
+            builder.addNetworkInterceptor(new ZstdInterceptor());
+        }
 
         this.httpClient = builder.build();
         this.noRedirectsHttpClient = new OkHttpClient.Builder(httpClient)
@@ -657,6 +662,23 @@ public void setConnectTimeout(int connectTimeout) {
         buildHttpClients();
     }
 
+    /*
+     * Gets if request use zstd encoding when possible
+     * @return true if request use zstd encoding when possible
+     */
+    public boolean getUseZstdCompression() {
+        return this.useZstdCompression;
+    }
+
+    /*
+     * Sets if request use zstd encoding when possible
+     * @param useZstdCompression true if request use zstd encoding when possible
+     */
+    public void setUseZstdCompression(boolean useZstdCompression) {
+        this.useZstdCompression = useZstdCompression;
+        buildHttpClients();
+    }
+
     /**
      * Gets the read timeout for this connection in milliseconds.
      *
diff --git a/src/main/java/com/box/sdk/BoxGlobalSettings.java b/src/main/java/com/box/sdk/BoxGlobalSettings.java
@@ -7,6 +7,7 @@ public final class BoxGlobalSettings {
     private static int connectTimeout = 0;
     private static int readTimeout = 0;
     private static int maxRetryAttempts = BoxAPIConnection.DEFAULT_MAX_RETRIES;
+    private static boolean useZstdCompression = true;
 
     private BoxGlobalSettings() {
     }
@@ -67,4 +68,20 @@ public static int getMaxRetryAttempts() {
     public static void setMaxRetryAttempts(int attempts) {
         BoxGlobalSettings.maxRetryAttempts = attempts;
     }
+
+    /*
+     * Returns the global settings for using Zstd compression.
+     * @return true if Zstd compression is enabled, false otherwise
+     */
+    public static boolean getUseZstdCompression() {
+        return useZstdCompression;
+    }
+
+    /*
+     * Sets the global settings for using Zstd compression.
+     * @param useZstdCompression true to enable Zstd compression, false otherwise
+     */
+    public static void setUseZstdCompression(boolean useZstdCompression) {
+        BoxGlobalSettings.useZstdCompression = useZstdCompression;
+    }
 }
diff --git a/src/main/java/com/box/sdk/ZstdInterceptor.java b/src/main/java/com/box/sdk/ZstdInterceptor.java
@@ -0,0 +1,93 @@
+package com.box.sdk;
+
+import com.github.luben.zstd.ZstdInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import okhttp3.Interceptor;
+import okhttp3.MediaType;
+import okhttp3.Request;
+import okhttp3.Response;
+import okhttp3.ResponseBody;
+import okio.BufferedSource;
+import okio.Okio;
+import okio.Source;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * Interceptor that adds zstd compression support to API requests.
+ * This interceptor adds zstd to the Accept-Encoding header and handles decompression of zstd responses.
+ */
+public class ZstdInterceptor implements Interceptor {
+    @NotNull
+    @Override
+    public Response intercept(Chain chain) throws IOException {
+        Request request = chain.request();
+
+        // Add zstd to the Accept-Encoding header
+        String acceptEncoding;
+        String acceptEncodingHeader = request.header("Accept-Encoding");
+        if (acceptEncodingHeader == null || acceptEncodingHeader.isEmpty()) {
+            acceptEncoding = "zstd";
+        } else {
+            acceptEncoding = acceptEncodingHeader + ", zstd";
+        }
+
+        Request compressedRequest = request.newBuilder()
+            .removeHeader("Accept-Encoding")
+            .addHeader("Accept-Encoding", acceptEncoding)
+            .build();
+
+        Response response = chain.proceed(compressedRequest);
+        String contentEncoding = response.header("Content-Encoding");
+
+        // Only handle zstd encoded responses, let OkHttp handle gzip and others
+        if (contentEncoding == null || !contentEncoding.equalsIgnoreCase("zstd")) {
+            return response;
+        }
+
+        ResponseBody originalBody = response.body();
+        if (originalBody == null) {
+            return response;
+        }
+
+        // Create a streaming response body
+        ResponseBody decompressedBody = createStreamingResponseBody(originalBody);
+
+        return response.newBuilder()
+                .body(decompressedBody)
+                .addHeader("X-Content-Encoding", "zstd")
+                .removeHeader("Content-Encoding")
+                .removeHeader("Content-Length")
+                .build();
+    }
+
+    /**
+     * Wraps the original response body in a streaming Zstd decompressor.
+     */
+    private ResponseBody createStreamingResponseBody(ResponseBody originalBody) {
+        return new ResponseBody() {
+            @Override
+            public MediaType contentType() {
+                return originalBody.contentType();
+            }
+
+            @Override
+            public long contentLength() {
+                return -1;
+            }
+
+            @Override
+            public BufferedSource source() {
+                InputStream decompressedStream;
+                try {
+                    decompressedStream = new ZstdInputStream(originalBody.byteStream());
+                } catch (IOException e) {
+                    throw new RuntimeException("Failed to create ZstdInputStream", e);
+                }
+
+                Source source = Okio.source(decompressedStream);
+                return Okio.buffer(source);
+            }
+        };
+    }
+}
diff --git a/src/test/java/com/box/sdk/BoxAPIRequestTest.java b/src/test/java/com/box/sdk/BoxAPIRequestTest.java
@@ -31,6 +31,7 @@
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.util.zip.GZIPOutputStream;
+import com.github.luben.zstd.ZstdOutputStream;
 import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.Mockito;
@@ -225,7 +226,7 @@ public void handlesGZIPResponse() {
             aResponse()
                 .withStatus(200)
                 .withHeader("content-type", APPLICATION_JSON)
-                .withHeader("content-encoding", "GZIP")
+                .withHeader("content-encoding", "gzip")
                 .withBody(gzipped(jsonString))
         ));
 
@@ -338,14 +339,43 @@ public void willNotAddAuthenticationHeaderWhenDisabled() {
     private byte[] gzipped(String str) {
         try {
             ByteArrayOutputStream obj = new ByteArrayOutputStream();
-            GZIPOutputStream gzip = new GZIPOutputStream(obj);
-            gzip.write(str.getBytes(UTF_8));
-            gzip.close();
+            try (GZIPOutputStream gzip = new GZIPOutputStream(obj)) {
+                byte[] bytes = str.getBytes(UTF_8);
+                gzip.write(bytes, 0, bytes.length);
+                gzip.flush();
+                gzip.finish();
+            }
             return obj.toByteArray();
         } catch (IOException e) {
             throw new RuntimeException(e);
         }
+    }
+
+    @Test
+    public void handlesZstdResponse() throws IOException {
+        BoxAPIConnection api = createConnectionWith(boxMockUrl().toString());
+        BoxAPIRequest request = new BoxAPIRequest(api, boxMockUrl(), "GET");
+        String jsonString = "{\"foo\":\"bar\"}";
 
+        stubFor(get(urlEqualTo("/")).willReturn(
+            aResponse()
+                .withStatus(200)
+                .withHeader("content-type", APPLICATION_JSON)
+                .withHeader("content-encoding", "zstd")
+                .withBody(zstdCompressed(jsonString))
+        ));
+
+        try (BoxAPIResponse response = request.send()) {
+            assertThat(response.bodyToString(), is(jsonString));
+        }
+    }
+
+    private byte[] zstdCompressed(String str) throws IOException {
+        ByteArrayOutputStream obj = new ByteArrayOutputStream();
+        try (ZstdOutputStream zstd = new ZstdOutputStream(obj)) {
+            zstd.write(str.getBytes(UTF_8));
+        }
+        return obj.toByteArray();
     }
 
     private URL boxMockUrl() {
