fix(xml): close DTD resolver resources

CHANGES.txt

@@ -1,4 +1,5 @@
 Current (7.13.0)
+Fixed: GITHUB-3316: close DTD resolver resources and bound connection waits (w3lld1)
 Fixed: GITHUB-3138: Clarified that class-level @Ignore applies to subclasses, not nested classes (ShinyHero666)
 New: Added OpenRewrite to the build with a hand-picked recipe list (see rewrite.yml), and applied it to the main sources (Julien Herr)
 Fixed: Remove leftover dead JUnit code: the deprecated unused ConversionUtils and orphaned JUnit test samples, following the removal of JUnit execution support in 7.10.0 (Julien Herr)

testng-core/src/main/java/org/testng/xml/TestNGContentHandler.java

@@ -2,6 +2,7 @@
 
 import static org.testng.internal.Utils.isStringBlank;
 
+import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
@@ -16,7 +17,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.Stack;
 import org.testng.ITestObjectFactory;
@@ -43,6 +43,8 @@
  */
 // TODO move to internal
 public class TestNGContentHandler extends DefaultHandler {
+  private static final int DTD_CONNECTION_TIMEOUT_MILLIS = 10_000;
+
   private XmlSuite m_currentSuite = null;
   private XmlTest m_currentTest = null;
   private XmlDefine m_currentDefine = null;
@@ -71,27 +73,49 @@ public class TestNGContentHandler extends DefaultHandler {
                   "Failed to read [%s] from CLASSPATH. " + "Attempting to read from [%s].",
                   url.getPath(), systemId);
           Logger.getLogger(getClass()).warn(msg);
-          URLConnection urlConnection = url.openConnection();
-          if (urlConnection instanceof HttpURLConnection) {
-            HttpURLConnection conn = (HttpURLConnection) urlConnection;
-
-            int status = conn.getResponseCode();
-            if (status == HttpURLConnection.HTTP_MOVED_TEMP
-                || status == HttpURLConnection.HTTP_MOVED_PERM
-                || status == HttpURLConnection.HTTP_SEE_OTHER) {
-
-              String newUrl = conn.getHeaderField("Location");
-              conn = (HttpURLConnection) new URL(newUrl).openConnection();
-            }
-            stream = conn.getInputStream();
-          } else {
-            stream = urlConnection.getInputStream();
-          }
+          return readUrlAsInputSource(url);
+        }
+        try (InputStream input = stream) {
+          return new InputSource(new ByteArrayInputStream(input.readAllBytes()));
         }
-        return new InputSource(
-            Objects.requireNonNull(stream, "Failed to load DTD from " + systemId));
       };
 
+  static byte[] readUrl(URL url, boolean followRedirect) throws IOException {
+    URLConnection connection = url.openConnection();
+    configureConnection(connection);
+    if (!(connection instanceof HttpURLConnection)) {
+      try (InputStream input = connection.getInputStream()) {
+        return input.readAllBytes();
+      }
+    }
+
+    HttpURLConnection httpConnection = (HttpURLConnection) connection;
+    httpConnection.setInstanceFollowRedirects(false);
+    try {
+      int status = httpConnection.getResponseCode();
+      if (followRedirect
+          && (status == HttpURLConnection.HTTP_MOVED_TEMP
+              || status == HttpURLConnection.HTTP_MOVED_PERM
+              || status == HttpURLConnection.HTTP_SEE_OTHER)) {
+        return readUrl(new URL(url, httpConnection.getHeaderField("Location")), false);
+      }
+      try (InputStream input = httpConnection.getInputStream()) {
+        return input.readAllBytes();
+      }
+    } finally {
+      httpConnection.disconnect();
+    }
+  }
+
+  static InputSource readUrlAsInputSource(URL url) throws IOException {
+    return new InputSource(new ByteArrayInputStream(readUrl(url, true)));
+  }
+
+  static void configureConnection(URLConnection connection) {
+    connection.setConnectTimeout(DTD_CONNECTION_TIMEOUT_MILLIS);
+    connection.setReadTimeout(DTD_CONNECTION_TIMEOUT_MILLIS);
+  }
+
   enum Location {
     SUITE,
     TEST,
@@ -130,9 +154,11 @@ public InputSource resolveEntity(String publicId, String systemId)
 
     if (skipConsideringSystemId(systemId)) {
       m_validate = true;
-      InputStream is = loadDtdUsingClassLoader();
-      if (is != null) {
-        return new InputSource(is);
+      InputStream stream = loadDtdUsingClassLoader();
+      if (stream != null) {
+        try (InputStream input = stream) {
+          return new InputSource(new ByteArrayInputStream(input.readAllBytes()));
+        }
       }
       // If the classpath loading of DTD fails, then we try to load it from "https" TestNG site.
       System.out.println(

testng-core/src/test/java/org/testng/xml/TestNGContentHandlerTestSupport.java

@@ -0,0 +1,22 @@
+package org.testng.xml;
+
+import java.io.IOException;
+import java.net.URL;
+import java.net.URLConnection;
+import org.xml.sax.InputSource;
+
+public final class TestNGContentHandlerTestSupport {
+  private TestNGContentHandlerTestSupport() {}
+
+  public static void configureConnection(URLConnection connection) {
+    TestNGContentHandler.configureConnection(connection);
+  }
+
+  public static byte[] readUrl(URL url) throws IOException {
+    return TestNGContentHandler.readUrl(url, true);
+  }
+
+  public static InputSource readUrlAsInputSource(URL url) throws IOException {
+    return TestNGContentHandler.readUrlAsInputSource(url);
+  }
+}

testng-core/src/test/java/test/xml/TestNGContentHandlerTest.java

@@ -2,12 +2,23 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import com.sun.net.httpserver.HttpServer;
+import java.io.ByteArrayInputStream;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.net.HttpURLConnection;
+import java.net.InetSocketAddress;
+import java.net.URL;
+import java.net.URLConnection;
+import java.net.URLStreamHandler;
+import java.nio.charset.StandardCharsets;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.testng.annotations.Test;
 import org.testng.xml.SuiteXmlParser;
 import org.testng.xml.TestNGContentHandler;
+import org.testng.xml.TestNGContentHandlerTestSupport;
 import org.testng.xml.XmlClass;
 import org.testng.xml.XmlInclude;
 import org.xml.sax.SAXException;
@@ -36,6 +47,95 @@ public void ensureAppropriateConnectionObjectsAreUsed() throws IOException, SAXE
     assertThat(xmlClass.getSupportClass()).isEqualTo(test.xml.issue2501.TestClassSample.class);
   }
 
+  @Test(description = "GITHUB-3316")
+  public void resolverConfiguresDtdConnectionTimeouts() throws Exception {
+    URLConnection connection =
+        new URLConnection(new URL("https://testng.org/testng-1.1.dtd")) {
+          @Override
+          public void connect() throws IOException {}
+        };
+
+    TestNGContentHandlerTestSupport.configureConnection(connection);
+
+    assertThat(connection.getConnectTimeout()).isEqualTo(10_000);
+    assertThat(connection.getReadTimeout()).isEqualTo(10_000);
+  }
+
+  @Test(description = "GITHUB-3316")
+  public void resolverBuffersAndClosesExternalDtd() throws Exception {
+    byte[] content = "<!ELEMENT suite EMPTY>".getBytes(StandardCharsets.UTF_8);
+    AtomicBoolean sourceClosed = new AtomicBoolean();
+    URL dtd =
+        new URL(
+            null,
+            "memory://testng.dtd",
+            new URLStreamHandler() {
+              @Override
+              protected URLConnection openConnection(URL url) {
+                return new URLConnection(url) {
+                  @Override
+                  public void connect() throws IOException {}
+
+                  @Override
+                  public ByteArrayInputStream getInputStream() {
+                    return new ByteArrayInputStream(content) {
+                      @Override
+                      public void close() throws IOException {
+                        sourceClosed.set(true);
+                        super.close();
+                      }
+                    };
+                  }
+                };
+              }
+            });
+
+    var resolved = TestNGContentHandlerTestSupport.readUrlAsInputSource(dtd);
+    assertThat(sourceClosed).isTrue();
+    try (var input = resolved.getByteStream()) {
+      assertThat(input.readAllBytes()).isEqualTo(content);
+    }
+  }
+
+  @Test(description = "GITHUB-3316")
+  public void resolverFollowsOnlyOneRedirect() throws Exception {
+    AtomicInteger secondRequests = new AtomicInteger();
+    AtomicInteger thirdRequests = new AtomicInteger();
+    HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+    server.createContext(
+        "/first",
+        exchange -> {
+          exchange.getResponseHeaders().add("Location", "/second");
+          exchange.sendResponseHeaders(HttpURLConnection.HTTP_MOVED_TEMP, -1);
+          exchange.close();
+        });
+    server.createContext(
+        "/second",
+        exchange -> {
+          secondRequests.incrementAndGet();
+          exchange.getResponseHeaders().add("Location", "/third");
+          exchange.sendResponseHeaders(HttpURLConnection.HTTP_MOVED_TEMP, -1);
+          exchange.close();
+        });
+    server.createContext(
+        "/third",
+        exchange -> {
+          thirdRequests.incrementAndGet();
+          exchange.sendResponseHeaders(HttpURLConnection.HTTP_OK, 0);
+          exchange.close();
+        });
+    server.start();
+
+    try {
+      URL dtd = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/first");
+      TestNGContentHandlerTestSupport.readUrl(dtd);
+      assertThat(secondRequests).hasValue(1);
+      assertThat(thirdRequests).hasValue(0);
+    } finally {
+      server.stop(0);
+    }
+  }
+
   public static class LocalTestClass {
     @Test
     public void helloWorld() {}