Enforce Java 11 response body timeouts

java11/src/main/java/feign/http2client/Http2Client.java

@@ -32,12 +32,14 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpClient.Redirect;
 import java.net.http.HttpClient.Version;
+import java.net.http.HttpHeaders;
 import java.net.http.HttpRequest;
 import java.net.http.HttpRequest.BodyPublisher;
 import java.net.http.HttpRequest.BodyPublishers;
 import java.net.http.HttpRequest.Builder;
 import java.net.http.HttpResponse;
 import java.net.http.HttpResponse.BodyHandlers;
+import java.net.http.HttpTimeoutException;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collection;
@@ -52,13 +54,27 @@
 import java.util.TreeSet;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.zip.GZIPInputStream;
 import java.util.zip.InflaterInputStream;
+import javax.net.ssl.SSLSession;
 
 public class Http2Client implements Client, AsyncClient<Object> {
 
+  // Shared by all clients so response-body deadlines do not create one thread per client.
+  private static final ScheduledExecutorService BODY_READ_TIMEOUT_EXECUTOR =
+      Executors.newSingleThreadScheduledExecutor(
+          runnable -> {
+            Thread thread = new Thread(runnable, "feign-http2client-body-timeout");
+            thread.setDaemon(true);
+            return thread;
+          });
+
   private final HttpClient client;
 
   private final Map<Integer, SoftReference<HttpClient>> clients = new ConcurrentHashMap<>();
@@ -109,7 +125,7 @@ public Response execute(Request request, Options options) throws IOException {
       throw new IOException(e);
     }
 
-    return toFeignResponse(request, httpResponse);
+    return toFeignResponse(request, httpResponse, options);
   }
 
   @Override
@@ -125,7 +141,7 @@ public CompletableFuture<Response> execute(
     HttpClient clientForRequest = getOrCreateClient(options);
     CompletableFuture<HttpResponse<InputStream>> future =
         clientForRequest.sendAsync(httpRequest, HttpResponse.BodyHandlers.ofInputStream());
-    return future.thenApply(httpResponse -> toFeignResponse(request, httpResponse));
+    return future.thenApply(httpResponse -> toFeignResponse(request, httpResponse, options));
   }
 
   protected Response toFeignResponse(Request request, HttpResponse<InputStream> httpResponse) {
@@ -159,6 +175,194 @@ protected Response toFeignResponse(Request request, HttpResponse<InputStream> ht
         .build();
   }
 
+  protected Response toFeignResponse(
+      Request request, HttpResponse<InputStream> httpResponse, Options options) {
+    final InputStream body = httpResponse.body();
+    final InputStream timedBody = withReadTimeout(body, options);
+    return toFeignResponse(
+        request,
+        body == timedBody ? httpResponse : new TimeoutHttpResponse(httpResponse, timedBody));
+  }
+
+  private static InputStream withReadTimeout(InputStream body, Options options) {
+    if (body == null || options == null || options.readTimeout() <= 0) {
+      return body;
+    }
+    return new TimeoutInputStream(body, options.readTimeout(), options.readTimeoutUnit());
+  }
+
+  private static final class TimeoutInputStream extends InputStream {
+
+    private final InputStream delegate;
+    private final ScheduledFuture<?> timeoutFuture;
+    private volatile boolean timedOut;
+
+    private TimeoutInputStream(InputStream delegate, long timeout, TimeUnit timeoutUnit) {
+      this.delegate = delegate;
+      this.timeoutFuture =
+          BODY_READ_TIMEOUT_EXECUTOR.schedule(
+              () -> {
+                timedOut = true;
+                try {
+                  delegate.close();
+                } catch (IOException ignored) {
+                }
+              },
+              timeout,
+              timeoutUnit);
+    }
+
+    @Override
+    public int read() throws IOException {
+      try {
+        return afterRead(delegate.read());
+      } catch (IOException e) {
+        throw translateException(e);
+      }
+    }
+
+    @Override
+    public int read(byte[] b, int off, int len) throws IOException {
+      try {
+        return afterRead(delegate.read(b, off, len));
+      } catch (IOException e) {
+        throw translateException(e);
+      }
+    }
+
+    @Override
+    public long skip(long n) throws IOException {
+      try {
+        final long skipped = delegate.skip(n);
+        checkTimedOut();
+        return skipped;
+      } catch (IOException e) {
+        throw translateException(e);
+      }
+    }
+
+    @Override
+    public int available() throws IOException {
+      return delegate.available();
+    }
+
+    @Override
+    public synchronized void mark(int readLimit) {
+      delegate.mark(readLimit);
+    }
+
+    @Override
+    public synchronized void reset() throws IOException {
+      delegate.reset();
+    }
+
+    @Override
+    public boolean markSupported() {
+      return delegate.markSupported();
+    }
+
+    @Override
+    public void close() throws IOException {
+      timeoutFuture.cancel(false);
+      delegate.close();
+    }
+
+    private int afterRead(int result) throws HttpTimeoutException {
+      checkTimedOut();
+      if (result == -1) {
+        timeoutFuture.cancel(false);
+      }
+      return result;
+    }
+
+    private void checkTimedOut() throws HttpTimeoutException {
+      if (timedOut) {
+        throw timeoutException(null);
+      }
+    }
+
+    private IOException translateException(IOException exception) {
+      final HttpTimeoutException timeoutException = findTimeoutException(exception);
+      if (timeoutException != null) {
+        return timeoutException;
+      }
+      if (timedOut) {
+        return timeoutException(exception);
+      }
+      return exception;
+    }
+
+    private static HttpTimeoutException timeoutException(IOException cause) {
+      final HttpTimeoutException exception = new HttpTimeoutException("response timed out");
+      if (cause != null) {
+        exception.initCause(cause);
+      }
+      return exception;
+    }
+
+    private static HttpTimeoutException findTimeoutException(Throwable throwable) {
+      Throwable current = throwable;
+      while (current != null) {
+        if (current instanceof HttpTimeoutException) {
+          return (HttpTimeoutException) current;
+        }
+        current = current.getCause();
+      }
+      return null;
+    }
+  }
+
+  private static final class TimeoutHttpResponse implements HttpResponse<InputStream> {
+
+    private final HttpResponse<InputStream> delegate;
+    private final InputStream body;
+
+    private TimeoutHttpResponse(HttpResponse<InputStream> delegate, InputStream body) {
+      this.delegate = delegate;
+      this.body = body;
+    }
+
+    @Override
+    public int statusCode() {
+      return delegate.statusCode();
+    }
+
+    @Override
+    public HttpRequest request() {
+      return delegate.request();
+    }
+
+    @Override
+    public Optional<HttpResponse<InputStream>> previousResponse() {
+      return delegate.previousResponse();
+    }
+
+    @Override
+    public HttpHeaders headers() {
+      return delegate.headers();
+    }
+
+    @Override
+    public InputStream body() {
+      return body;
+    }
+
+    @Override
+    public Optional<SSLSession> sslSession() {
+      return delegate.sslSession();
+    }
+
+    @Override
+    public URI uri() {
+      return delegate.uri();
+    }
+
+    @Override
+    public Version version() {
+      return delegate.version();
+    }
+  }
+
   private HttpClient getOrCreateClient(Options options) {
     if (doesClientConfigurationDiffer(options)) {
       // create a new client from the existing one - but with connectTimeout and followRedirect

java11/src/test/java/feign/http2client/Http2ClientContentLengthTest.java

@@ -16,31 +16,41 @@
 package feign.http2client;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import feign.Request;
 import feign.Request.HttpMethod;
 import feign.Response;
 import feign.Util;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.net.URI;
 import java.net.http.HttpClient.Version;
 import java.net.http.HttpHeaders;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.net.http.HttpTimeoutException;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.zip.GZIPOutputStream;
 import javax.net.ssl.SSLSession;
 import org.junit.jupiter.api.Test;
 
 class Http2ClientContentLengthTest {
 
   private static HttpResponse<InputStream> responseWithContentLength(String contentLength) {
+    return responseWithContentLength(contentLength, new ByteArrayInputStream(new byte[0]));
+  }
+
+  private static HttpResponse<InputStream> responseWithContentLength(
+      String contentLength, InputStream body) {
     final HttpHeaders headers =
         HttpHeaders.of(Map.of("Content-Length", List.of(contentLength)), (name, value) -> true);
     return new HttpResponse<>() {
@@ -66,7 +76,7 @@ public HttpHeaders headers() {
 
       @Override
       public InputStream body() {
-        return new ByteArrayInputStream(new byte[0]);
+        return body;
       }
 
       @Override
@@ -185,4 +195,46 @@ void gzipDecodedBodyReportsUnknownLength() throws Exception {
     assertThat(Util.toString(response.body().asReader(StandardCharsets.UTF_8)))
         .isEqualTo("Compressed Data");
   }
+
+  @Test
+  void responseBodyTimeoutUsesOneStreamDeadline() throws Exception {
+    final CountDownLatch closed = new CountDownLatch(1);
+    final InputStream body =
+        new ByteArrayInputStream(new byte[] {1}) {
+          @Override
+          public void close() throws IOException {
+            closed.countDown();
+            super.close();
+          }
+        };
+    final Request.Options options =
+        new Request.Options(1, TimeUnit.SECONDS, 1, TimeUnit.MILLISECONDS, true);
+    final Response response =
+        new Http2Client().toFeignResponse(request(), responseWithContentLength("1", body), options);
+
+    assertThat(closed.await(5, TimeUnit.SECONDS)).isTrue();
+    assertThrows(HttpTimeoutException.class, response.body().asInputStream()::read);
+  }
+
+  @Test
+  void responseBodyTimeoutPreservesStreamOperations() throws Exception {
+    final Request.Options options =
+        new Request.Options(1, TimeUnit.SECONDS, 1, TimeUnit.MINUTES, true);
+    final Response response =
+        new Http2Client()
+            .toFeignResponse(
+                request(),
+                responseWithContentLength("3", new ByteArrayInputStream(new byte[] {1, 2, 3})),
+                options);
+    final InputStream body = response.body().asInputStream();
+
+    assertThat(body.markSupported()).isTrue();
+    body.mark(3);
+    assertThat(body.read()).isEqualTo(1);
+    assertThat(body.skip(1)).isEqualTo(1);
+    assertThat(body.read()).isEqualTo(3);
+    body.reset();
+    assertThat(body.read()).isEqualTo(1);
+    body.close();
+  }
 }