[FEATURE] Downloader fix

LaunchServer/src/main/java/pro/gravit/launchserver/auth/core/interfaces/provider/AuthSupportUserBan.java

@@ -3,11 +3,13 @@
 import pro.gravit.launchserver.auth.core.User;
 import pro.gravit.launchserver.auth.core.interfaces.user.UserSupportBanInfo;
 
-public interface AuthSupportUserBan {
+import java.time.LocalDateTime;
-    void banUser(User user, String reason);
 
-    default void banUser(User user) {
+public interface AuthSupportUserBan {
-        banUser(user, null);
+    UserSupportBanInfo.UserBanInfo banUser(User user, String reason, String moderator, LocalDateTime startTime, LocalDateTime endTime);
+
+    default UserSupportBanInfo.UserBanInfo banUser(User user) {
+        return banUser(user, null, null, LocalDateTime.now(), null);
     }
 
     void unbanUser(User user);

Launcher/src/main/java/pro/gravit/launcher/BasicLauncherEventHandler.java

@@ -0,0 +1,28 @@
+package pro.gravit.launcher;
+
+import pro.gravit.launcher.events.ExtendedTokenRequestEvent;
+import pro.gravit.launcher.events.request.SecurityReportRequestEvent;
+import pro.gravit.launcher.request.Request;
+import pro.gravit.launcher.request.WebSocketEvent;
+import pro.gravit.launcher.request.websockets.ClientWebSocketService;
+
+public class BasicLauncherEventHandler implements ClientWebSocketService.EventHandler {
+
+    @Override
+    public <T extends WebSocketEvent> boolean eventHandle(T event) {
+        if (event instanceof SecurityReportRequestEvent) {
+            SecurityReportRequestEvent event1 = (SecurityReportRequestEvent) event;
+            if (event1.action == SecurityReportRequestEvent.ReportAction.CRASH) {
+                LauncherEngine.exitLauncher(80);
+            }
+        }
+        if (event instanceof ExtendedTokenRequestEvent) {
+            ExtendedTokenRequestEvent event1 = (ExtendedTokenRequestEvent) event;
+            String token = event1.getExtendedToken();
+            if (token != null) {
+                Request.addExtendedToken(event1.getExtendedTokenName(), token);
+            }
+        }
+        return false;
+    }
+}

Launcher/src/main/java/pro/gravit/launcher/LauncherEngine.java

@@ -192,6 +192,7 @@ public void start(String... args) throws Throwable {
                     throw new RequestException("Connection failed", e);
                 }
             };
+            Request.service.registerEventHandler(new BasicLauncherEventHandler());
         }
         Objects.requireNonNull(args, "args");
         if (started.getAndSet(true))

Launcher/src/main/java/pro/gravit/launcher/client/ClientLauncherEntryPoint.java

@@ -1,9 +1,6 @@
 package pro.gravit.launcher.client;
 
-import pro.gravit.launcher.Launcher;
+import pro.gravit.launcher.*;
-import pro.gravit.launcher.LauncherAgent;
-import pro.gravit.launcher.LauncherConfig;
-import pro.gravit.launcher.LauncherEngine;
 import pro.gravit.launcher.api.AuthService;
 import pro.gravit.launcher.api.ClientService;
 import pro.gravit.launcher.client.events.client.*;
@@ -104,6 +101,7 @@ public static void main(String[] args) throws Throwable {
             LogHelper.info("Using Sessions");
             Request.setSession(params.session);
         }
+        Request.service.registerEventHandler(new BasicLauncherEventHandler());
         checkJVMBitsAndVersion(params.profile.getMinJavaVersion(), params.profile.getRecommendJavaVersion(), params.profile.getMaxJavaVersion(), params.profile.isWarnMissJavaVersion());
         LauncherEngine.modulesManager.invokeEvent(new ClientProcessInitPhase(engine, params));
 

LauncherAPI/src/main/java/pro/gravit/launcher/events/ExtendedTokenRequestEvent.java

@@ -0,0 +1,7 @@
+package pro.gravit.launcher.events;
+
+public interface ExtendedTokenRequestEvent {
+    String getExtendedTokenName();
+
+    String getExtendedToken();
+}

LauncherAPI/src/main/java/pro/gravit/launcher/request/Request.java

@@ -14,6 +14,7 @@
 
 import java.util.*;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.BiConsumer;
 
 public abstract class Request<R extends WebSocketEvent> implements WebSocketRequest {
     public static StdWebSocketService service;
@@ -23,6 +24,7 @@ public abstract class Request<R extends WebSocketEvent> implements WebSocketRequ
     private static String authId;
     private static long tokenExpiredTime;
     private static List<ExtendedTokenCallback> extendedTokenCallbacks = new ArrayList<>(4);
+    private static List<BiConsumer<String, AuthRequestEvent.OAuthRequestEvent>> oauthChangeHandlers = new ArrayList<>(4);
     @LauncherNetworkAPI
     public final UUID requestUUID = UUID.randomUUID();
     private transient final AtomicBoolean started = new AtomicBoolean(false);
@@ -43,6 +45,9 @@ public static void setOAuth(String authId, AuthRequestEvent.OAuthRequestEvent ev
         } else {
             tokenExpiredTime = 0;
         }
+        for (BiConsumer<String, AuthRequestEvent.OAuthRequestEvent> handler : oauthChangeHandlers) {
+            handler.accept(authId, event);
+        }
     }
 
     public static AuthRequestEvent.OAuthRequestEvent getOAuth() {
@@ -154,6 +159,18 @@ public void addExtendedTokenCallback(ExtendedTokenCallback cb) {
         extendedTokenCallbacks.add(cb);
     }
 
+    public void removeExtendedTokenCallback(ExtendedTokenCallback cb) {
+        extendedTokenCallbacks.remove(cb);
+    }
+
+    public void addOAuthChangeHandler(BiConsumer<String, AuthRequestEvent.OAuthRequestEvent> eventHandler) {
+        oauthChangeHandlers.add(eventHandler);
+    }
+
+    public void removeOAuthChangeHandler(BiConsumer<String, AuthRequestEvent.OAuthRequestEvent> eventHandler) {
+        oauthChangeHandlers.remove(eventHandler);
+    }
+
     public R request() throws Exception {
         if (!started.compareAndSet(false, true))
             throw new IllegalStateException("Request already started");

LauncherCore/src/main/java/pro/gravit/launcher/AsyncDownloader.java

@@ -32,6 +32,7 @@ public class AsyncDownloader {
     private static volatile SSLSocketFactory sslSocketFactory;
     private static volatile SSLContext sslContext;
     public final Callback callback;
+    public volatile boolean isClosed;
 
     public AsyncDownloader(Callback callback) {
         this.callback = callback;
@@ -170,6 +171,7 @@ public void transfer(InputStream input, Path file, long size) throws IOException
             // Download with digest update
             byte[] bytes = IOHelper.newBuffer();
             while (downloaded < size) {
+                if (isClosed) throw new IOException("Download interrupted");
                 int remaining = (int) Math.min(size - downloaded, bytes.length);
                 int length = input.read(bytes, 0, remaining);
                 if (length < 0)

LauncherCore/src/main/java/pro/gravit/utils/Downloader.java

@@ -13,7 +13,27 @@ public interface DownloadCallback {
         void apply(long fullDiff);
     }
 
-    public CompletableFuture<Void> downloadList(List<AsyncDownloader.SizedFile> files, String baseURL, Path targetDir, DownloadCallback callback, ExecutorService executor, int threads) throws Exception {
+    private final CompletableFuture<Void> future;
+    private final AsyncDownloader asyncDownloader;
+
+    private Downloader(CompletableFuture<Void> future, AsyncDownloader downloader) {
+        this.future = future;
+        this.asyncDownloader = downloader;
+    }
+
+    public CompletableFuture<Void> getFuture() {
+        return future;
+    }
+
+    public void cancel() {
+        this.asyncDownloader.isClosed = true;
+    }
+
+    public boolean isCanceled() {
+        return this.asyncDownloader.isClosed;
+    }
+
+    public static Downloader downloadList(List<AsyncDownloader.SizedFile> files, String baseURL, Path targetDir, DownloadCallback callback, ExecutorService executor, int threads) throws Exception {
         final boolean closeExecutor;
         if (executor == null) {
             executor = Executors.newWorkStealingPool(4);
@@ -30,10 +50,10 @@ public CompletableFuture<Void> downloadList(List<AsyncDownloader.SizedFile> file
         CompletableFuture<Void> future = CompletableFuture.allOf(asyncDownloader.runDownloadList(list, baseURL, targetDir, executor));
 
         ExecutorService finalExecutor = executor;
-        return future.thenAccept(e -> {
+        return new Downloader(future.thenAccept(e -> {
             if (closeExecutor) {
                 finalExecutor.shutdownNow();
             }
-        });
+        }), asyncDownloader);
     }
 }

LauncherCore/src/main/java11/pro/gravit/utils/Downloader.java

@@ -11,6 +11,7 @@
 import java.nio.ByteBuffer;
 import java.nio.file.Path;
 import java.util.Collections;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Queue;
 import java.util.concurrent.*;
@@ -29,12 +30,55 @@ public interface DownloadCallback {
         void onComplete(Path path);
     }
 
-    public static CompletableFuture<Void> downloadList(List<AsyncDownloader.SizedFile> files, String baseURL, Path targetDir, DownloadCallback callback, ExecutorService executor, int threads) throws Exception {
+    protected final HttpClient client;
+    protected final ExecutorService executor;
+    protected CompletableFuture<Void> future;
+    protected final LinkedList<DownloadTask> tasks = new LinkedList<>();
+
+    protected Downloader(HttpClient client, ExecutorService executor) {
+        this.client = client;
+        this.executor = executor;
+    }
+
+    public void cancel() {
+        for (DownloadTask task : tasks) {
+            if (!task.isCompleted()) {
+                task.cancel();
+            }
+        }
+        tasks.clear();
+        executor.shutdownNow();
+    }
+
+    public boolean isCanceled() {
+        return executor.isTerminated();
+    }
+
+    public CompletableFuture<Void> getFuture() {
+        return future;
+    }
+
+    public static Downloader downloadList(List<AsyncDownloader.SizedFile> files, String baseURL, Path targetDir, DownloadCallback callback, ExecutorService executor, int threads) throws Exception {
         boolean closeExecutor = false;
         if (executor == null) {
             executor = Executors.newWorkStealingPool(Math.min(3, threads));
             closeExecutor = true;
         }
+        Downloader downloader = newDownloader(executor);
+        downloader.future = downloader.downloadFiles(files, baseURL, targetDir, callback, executor, threads);
+        if (closeExecutor) {
+            ExecutorService finalExecutor = executor;
+            downloader.future = downloader.future.thenAccept(e -> {
+                finalExecutor.shutdownNow();
+            });
+        }
+        return downloader;
+    }
+
+    public static Downloader newDownloader(ExecutorService executor) {
+        if (executor == null) {
+            throw new NullPointerException();
+        }
         HttpClient.Builder builder = HttpClient.newBuilder()
                 .version(isNoHttp2 ? HttpClient.Version.HTTP_1_1 : HttpClient.Version.HTTP_2)
                 .followRedirects(HttpClient.Redirect.NORMAL)
@@ -46,21 +90,15 @@ public static CompletableFuture<Void> downloadList(List<AsyncDownloader.SizedFil
                 throw new SecurityException(e);
             }
         }
-        CompletableFuture<Void> future = downloadList(builder.build(), files, baseURL, targetDir, callback, executor, threads);
+        HttpClient client = builder.build();
-        if (closeExecutor) {
+        return new Downloader(client, executor);
-            ExecutorService finalExecutor = executor;
-            future = future.thenAccept(e -> {
-                finalExecutor.shutdownNow();
-            });
-        }
-        return future;
     }
 
     private static class ConsumerObject {
         Consumer<HttpResponse<Path>> next = null;
     }
 
-    public static CompletableFuture<Void> downloadList(HttpClient client, List<AsyncDownloader.SizedFile> files, String baseURL, Path targetDir, DownloadCallback callback, ExecutorService executor, int threads) throws Exception {
+    public CompletableFuture<Void> downloadFiles(List<AsyncDownloader.SizedFile> files, String baseURL, Path targetDir, DownloadCallback callback, ExecutorService executor, int threads) throws Exception {
         // URI scheme
         URI baseUri = new URI(baseURL);
         Collections.shuffle(files);
@@ -79,7 +117,8 @@ public static CompletableFuture<Void> downloadList(HttpClient client, List<Async
                 return;
             }
             try {
-                sendAsync(client, file, baseUri, targetDir, callback).thenAccept(consumerObject.next);
+                DownloadTask task = sendAsync(file, baseUri, targetDir, callback);
+                task.completableFuture.thenAccept(consumerObject.next);
             } catch (Exception exception) {
                 future.completeExceptionally(exception);
             }
@@ -91,11 +130,39 @@ public static CompletableFuture<Void> downloadList(HttpClient client, List<Async
         return future;
     }
 
-    private static CompletableFuture<HttpResponse<Path>> sendAsync(HttpClient client, AsyncDownloader.SizedFile file, URI baseUri, Path targetDir, DownloadCallback callback) throws Exception {
+    public static class DownloadTask {
-        return client.sendAsync(makeHttpRequest(baseUri, file.urlPath), makeBodyHandler(targetDir.resolve(file.filePath), callback));
+        public final ProgressTrackingBodyHandler<Path> bodyHandler;
+        public final CompletableFuture<HttpResponse<Path>> completableFuture;
+
+        public DownloadTask(ProgressTrackingBodyHandler<Path> bodyHandler, CompletableFuture<HttpResponse<Path>> completableFuture) {
+            this.bodyHandler = bodyHandler;
+            this.completableFuture = completableFuture;
         }
 
-    private static HttpRequest makeHttpRequest(URI baseUri, String filePath) throws URISyntaxException {
+        public boolean isCompleted() {
+            return completableFuture.isDone() | completableFuture.isCompletedExceptionally();
+        }
+
+        public void cancel() {
+            bodyHandler.cancel();
+        }
+    }
+
+    protected DownloadTask sendAsync(AsyncDownloader.SizedFile file, URI baseUri, Path targetDir, DownloadCallback callback) throws Exception {
+        ProgressTrackingBodyHandler<Path> bodyHandler = makeBodyHandler(targetDir.resolve(file.filePath), callback);
+        CompletableFuture<HttpResponse<Path>> future = client.sendAsync(makeHttpRequest(baseUri, file.urlPath), bodyHandler);
+        var ref = new Object() {
+            DownloadTask task = null;
+        };
+        ref.task = new DownloadTask(bodyHandler, future.thenApply((e) -> {
+            tasks.remove(ref.task);
+            return e;
+        }));
+        tasks.add(ref.task);
+        return ref.task;
+    }
+
+    protected HttpRequest makeHttpRequest(URI baseUri, String filePath) throws URISyntaxException {
         String scheme = baseUri.getScheme();
         String host = baseUri.getHost();
         int port = baseUri.getPort();
@@ -109,13 +176,14 @@ private static HttpRequest makeHttpRequest(URI baseUri, String filePath) throws
                 .build();
     }
 
-    private static HttpResponse.BodyHandler<Path> makeBodyHandler(Path file, DownloadCallback callback) {
+    protected ProgressTrackingBodyHandler<Path> makeBodyHandler(Path file, DownloadCallback callback) {
         return new ProgressTrackingBodyHandler<>(HttpResponse.BodyHandlers.ofFile(file), callback);
     }
 
     public static class ProgressTrackingBodyHandler<T> implements HttpResponse.BodyHandler<T> {
         private final HttpResponse.BodyHandler<T> delegate;
         private final DownloadCallback callback;
+        private ProgressTrackingBodySubscriber subscriber;
 
         public ProgressTrackingBodyHandler(HttpResponse.BodyHandler<T> delegate, DownloadCallback callback) {
             this.delegate = delegate;
@@ -124,11 +192,19 @@ public ProgressTrackingBodyHandler(HttpResponse.BodyHandler<T> delegate, Downloa
 
         @Override
         public HttpResponse.BodySubscriber<T> apply(HttpResponse.ResponseInfo responseInfo) {
-            return delegate.apply(responseInfo);
+            subscriber = new ProgressTrackingBodySubscriber(delegate.apply(responseInfo));
+            return subscriber;
+        }
+
+        public void cancel() {
+            if (subscriber != null) {
+                subscriber.cancel();
+            }
         }
 
         private class ProgressTrackingBodySubscriber implements HttpResponse.BodySubscriber<T> {
             private final HttpResponse.BodySubscriber<T> delegate;
+            private Flow.Subscription subscription;
 
             public ProgressTrackingBodySubscriber(HttpResponse.BodySubscriber<T> delegate) {
                 this.delegate = delegate;
@@ -141,6 +217,7 @@ public CompletionStage<T> getBody() {
 
             @Override
             public void onSubscribe(Flow.Subscription subscription) {
+                this.subscription = subscription;
                 delegate.onSubscribe(subscription);
             }
 
@@ -163,6 +240,12 @@ public void onError(Throwable throwable) {
             public void onComplete() {
                 delegate.onComplete();
             }
+
+            public void cancel() {
+                if (subscription != null) {
+                    subscription.cancel();
+                }
+            }
         }
     }
 }