001package com.box.sdk;
002
003import com.github.luben.zstd.ZstdInputStream;
004import java.io.IOException;
005import java.io.InputStream;
006import okhttp3.Interceptor;
007import okhttp3.MediaType;
008import okhttp3.Request;
009import okhttp3.Response;
010import okhttp3.ResponseBody;
011import okio.BufferedSource;
012import okio.Okio;
013import okio.Source;
014import org.jetbrains.annotations.NotNull;
015
016/**
017 * Interceptor that adds zstd compression support to API requests. This interceptor adds zstd to the
018 * Accept-Encoding header and handles decompression of zstd responses.
019 */
020public class ZstdInterceptor implements Interceptor {
021  @NotNull
022  @Override
023  public Response intercept(Chain chain) throws IOException {
024    Request request = chain.request();
025
026    // Add zstd to the Accept-Encoding header
027    String acceptEncoding;
028    String acceptEncodingHeader = request.header("Accept-Encoding");
029    if (acceptEncodingHeader == null || acceptEncodingHeader.isEmpty()) {
030      acceptEncoding = "zstd";
031    } else {
032      acceptEncoding = acceptEncodingHeader + ", zstd";
033    }
034
035    Request compressedRequest =
036        request
037            .newBuilder()
038            .removeHeader("Accept-Encoding")
039            .addHeader("Accept-Encoding", acceptEncoding)
040            .build();
041
042    Response response = chain.proceed(compressedRequest);
043    String contentEncoding = response.header("Content-Encoding");
044
045    // Only handle zstd encoded responses, let OkHttp handle gzip and others
046    if (contentEncoding == null || !contentEncoding.equalsIgnoreCase("zstd")) {
047      return response;
048    }
049
050    ResponseBody originalBody = response.body();
051    if (originalBody == null) {
052      return response;
053    }
054
055    // Create a streaming response body
056    ResponseBody decompressedBody = createStreamingResponseBody(originalBody);
057
058    return response
059        .newBuilder()
060        .body(decompressedBody)
061        .addHeader("X-Content-Encoding", "zstd")
062        .removeHeader("Content-Encoding")
063        .removeHeader("Content-Length")
064        .build();
065  }
066
067  /** Wraps the original response body in a streaming Zstd decompressor. */
068  private ResponseBody createStreamingResponseBody(ResponseBody originalBody) {
069    return new ResponseBody() {
070      @Override
071      public MediaType contentType() {
072        return originalBody.contentType();
073      }
074
075      @Override
076      public long contentLength() {
077        return -1;
078      }
079
080      @Override
081      public BufferedSource source() {
082        InputStream decompressedStream;
083        try {
084          decompressedStream = new ZstdInputStream(originalBody.byteStream());
085        } catch (IOException e) {
086          throw new RuntimeException("Failed to create ZstdInputStream", e);
087        }
088
089        Source source = Okio.source(decompressedStream);
090        return Okio.buffer(source);
091      }
092    };
093  }
094}