ResponseTeeingHTTPClient.java

/*
 * Copyright (C) 2023 B3Partners B.V.
 *
 * SPDX-License-Identifier: MIT
 */
package org.tailormap.api.geotools;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.function.BiConsumer;
import org.geotools.http.HTTPClient;
import org.geotools.http.HTTPResponse;

/**
 * Wrapper for a GeoTools HTTPClient that allows access to the response headers and body after the original response has
 * been consumed and disposed of by other code. The response body for the latest request is kept in memory, a consumer
 * can be specified to get responses of earlier requests. Response headers of an unwrapped HTTPClient can only be
 * retrieved when the response isn't disposed, this wrapper allows response headers to be cached for retrieval even
 * after disposal.
 */
public class ResponseTeeingHTTPClient implements HTTPClient {

  public class TeeHTTPResponseWrapper implements HTTPResponse {
    private final HTTPResponse wrapped;

    private boolean disposed = false;
    private final ByteArrayOutputStream copy = new ByteArrayOutputStream();

    private String contentType;
    private final Map<String, String> cachedResponseHeaders = new HashMap<>();

    public TeeHTTPResponseWrapper(HTTPResponse wrapped) {
      this.wrapped = wrapped;
    }

    public byte[] getCopy() {
      return copy.toByteArray();
    }

    // <editor-fold desc="methods directly delegated to wrapped object">
    @Override
    public void dispose() {
      disposed = true;
      wrapped.dispose();
    }

    public boolean isDisposed() {
      return disposed;
    }

    @Override
    public String getContentType() {
      if (contentType == null) {
        contentType = wrapped.getContentType();
      }
      return contentType;
    }

    @Override
    public String getResponseHeader(String header) {
      if (cachedResponseHeaders.containsKey(header)) {
        return cachedResponseHeaders.get(header);
      }
      // When disposed, the wrapped client will probably throw a NPE when trying to get a response
      // header
      if (isDisposed()) {
        return null;
      }
      return wrapped.getResponseHeader(header);
    }

    @Override
    public InputStream getResponseStream() throws IOException {
      // Cache response headers now, when the internal connection is still available
      for (String header : responseHeadersToCache) {
        cachedResponseHeaders.put(header, wrapped.getResponseHeader(header));
      }
      return new org.apache.commons.io.input.TeeInputStream(wrapped.getResponseStream(), copy);
    }

    @Override
    public String getResponseCharset() {
      return wrapped.getResponseCharset();
    }
    // </editor-fold>
  }

  private TeeHTTPResponseWrapper responseWrapper;

  private final HTTPClient wrapped;

  private final BiConsumer<URL, TeeHTTPResponseWrapper> requestConsumer;

  private final Set<String> responseHeadersToCache;

  /**
   * Wrap a GeoTools HTTPClient allowing access to the latest response body after it has been consumed by handing the
   * client to a GeoTools module. Note that getting the response headers after the original response has been consumed
   * will return null without explicitly calling the other constructor with the responseHeadersToCache parameter.
   *
   * @param wrapped The GeoTools HTTPClient to wrap
   */
  public ResponseTeeingHTTPClient(HTTPClient wrapped) {
    this(wrapped, null, null);
  }

  /**
   * Wrap a GeoTools HTTPClient allowing access to the all responses after they have been consumed by handing the
   * client to a GeoTools module by passing a consumer for each request URL and wrapped response. The
   * responseHeadersToCache parameter allows access to response headers after the original wrapped response has been
   * disposed (only cached when getResponseStream() is requested).
   *
   * @param wrapped The GeoTools HTTPClient to wrap
   * @param requestConsumer Consumer for each request so not only the latest response can be accessed, may be null
   * @param responseHeadersToCache Which response headers to cache, may be null
   */
  public ResponseTeeingHTTPClient(
      HTTPClient wrapped,
      BiConsumer<URL, TeeHTTPResponseWrapper> requestConsumer,
      Set<String> responseHeadersToCache) {
    this.wrapped = wrapped;
    this.requestConsumer = requestConsumer == null ? (url, response) -> {} : requestConsumer;
    this.responseHeadersToCache = responseHeadersToCache == null ? Collections.emptySet() : responseHeadersToCache;
  }

  @Override
  public HTTPResponse get(URL url, Map<String, String> headers) throws IOException {
    this.responseWrapper = new TeeHTTPResponseWrapper(wrapped.get(url, headers));
    requestConsumer.accept(url, this.responseWrapper);
    return responseWrapper;
  }

  @Override
  public HTTPResponse post(URL url, InputStream inputStream, String s) throws IOException {
    this.responseWrapper = new TeeHTTPResponseWrapper(wrapped.post(url, inputStream, s));
    requestConsumer.accept(url, this.responseWrapper);
    return responseWrapper;
  }

  @Override
  public HTTPResponse get(URL url) throws IOException {
    this.responseWrapper = new TeeHTTPResponseWrapper(wrapped.get(url));
    requestConsumer.accept(url, this.responseWrapper);
    return responseWrapper;
  }

  public HTTPResponse getLatestResponse() {
    return responseWrapper;
  }

  public byte[] getLatestResponseCopy() {
    return responseWrapper == null ? null : responseWrapper.getCopy();
  }

  // <editor-fold desc="methods directly delegated to wrapped object">
  @Override
  public String getUser() {
    return wrapped.getUser();
  }

  @Override
  public void setUser(String user) {
    wrapped.setUser(user);
  }

  @Override
  public String getPassword() {
    return wrapped.getPassword();
  }

  @Override
  public void setPassword(String password) {
    wrapped.setPassword(password);
  }

  @Override
  public Map<String, String> getExtraParams() {
    return wrapped.getExtraParams();
  }

  @Override
  public void setExtraParams(Map<String, String> extraParams) {
    wrapped.setExtraParams(extraParams);
  }

  @Override
  public int getConnectTimeout() {
    return wrapped.getConnectTimeout();
  }

  @Override
  public void setConnectTimeout(int connectTimeout) {
    wrapped.setConnectTimeout(connectTimeout);
  }

  @Override
  public int getReadTimeout() {
    return wrapped.getReadTimeout();
  }

  @Override
  public void setReadTimeout(int readTimeout) {
    wrapped.setReadTimeout(readTimeout);
  }

  @Override
  public void setTryGzip(boolean tryGZIP) {
    wrapped.setTryGzip(tryGZIP);
  }

  @Override
  public boolean isTryGzip() {
    return wrapped.isTryGzip();
  }
  // </editor-fold>
}