Make use of connection patcher in more places

3 files changed, 17 insertions, 11 deletions

diff --git a/src/main/java/moe/nea/libautoupdate/MavenSource.java b/src/main/java/moe/nea/libautoupdate/MavenSource.java
index 102e8c7..d1e7d5d 100644
--- a/src/main/java/moe/nea/libautoupdate/MavenSource.java
+++ b/src/main/java/moe/nea/libautoupdate/MavenSource.java
@@ -44,7 +44,7 @@ public class MavenSource implements UpdateSource {
         dbf.setFeature(XMLConstants.FEATURE_SECURE_PROCESSING, true);
         var db = dbf.newDocumentBuilder();
         return CompletableFuture.supplyAsync(() -> {
-            try (val is = new URL(getMavenMetadataUrl()).openStream()) {
+            try (val is = UpdateUtils.openUrlConnection(new URL(getMavenMetadataUrl()))) {
                 var document = db.parse(new InputSource(is));
                 var metadata = (Element) document.getDocumentElement();
                 var versioning = (Element) metadata.getElementsByTagName("versioning").item(0);
diff --git a/src/main/java/moe/nea/libautoupdate/PotentialUpdate.java b/src/main/java/moe/nea/libautoupdate/PotentialUpdate.java
index 2825c0d..37a37de 100644
--- a/src/main/java/moe/nea/libautoupdate/PotentialUpdate.java
+++ b/src/main/java/moe/nea/libautoupdate/PotentialUpdate.java
@@ -78,7 +78,7 @@ public class PotentialUpdate {
      * Download the updated jar into the storage directory.
      */
     public void downloadUpdate() throws IOException {
-        try (val from = update.getDownloadAsURL().openStream();
+        try (val from = UpdateUtils.openUrlConnection(update.getDownloadAsURL());
              val to = new FileOutputStream(getUpdateJarStorage())) {
             UpdateUtils.connect(from, to);
         }
diff --git a/src/main/java/moe/nea/libautoupdate/UpdateUtils.java b/src/main/java/moe/nea/libautoupdate/UpdateUtils.java
index 13bc4eb..862e6df 100644
--- a/src/main/java/moe/nea/libautoupdate/UpdateUtils.java
+++ b/src/main/java/moe/nea/libautoupdate/UpdateUtils.java
@@ -7,7 +7,10 @@ import lombok.var;
 import java.io.*;
 import java.lang.reflect.Type;
 import java.math.BigInteger;
-import java.net.*;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.net.URL;
+import java.net.URLConnection;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
@@ -70,6 +73,13 @@ public class UpdateUtils {
         }
     }
 
+    public static InputStream openUrlConnection(URL url) throws IOException {
+        val conn = url.openConnection();
+        if (connectionPatcher != null)
+            connectionPatcher.accept(conn);
+        return conn.getInputStream();
+    }
+
     public static void deleteDirectory(Path path) throws IOException {
         if (!Files.exists(path)) return;
         Files.walkFileTree(path, new SimpleFileVisitor<Path>() {
@@ -89,24 +99,20 @@ public class UpdateUtils {
 
     private static Consumer<URLConnection> connectionPatcher = null;
 
+    /**
+     * Insert a connection patcher, which can modify connections before they are read from.
+     */
     public static void patchConnection(Consumer<URLConnection> connectionPatcher) {
         UpdateUtils.connectionPatcher = connectionPatcher;
     }
 
     public static <T> CompletableFuture<T> httpGet(String url, Gson gson, Type clazz) {
         return CompletableFuture.supplyAsync(() -> {
-            URLConnection conn = null;
             try {
-                conn = new URL(url).openConnection();
-                if (connectionPatcher != null)
-                    connectionPatcher.accept(conn);
-                try (val is = conn.getInputStream()) {
+                try (val is = openUrlConnection(new URL(url))) {
                     return gson.fromJson(new InputStreamReader(is, StandardCharsets.UTF_8), clazz);
                 }
             } catch (IOException e) {
-                if (conn instanceof HttpURLConnection) {
-                    ((HttpURLConnection) conn).disconnect();
-                }
                 throw new CompletionException(e);
             }
         });