From 55c77dd34301fa0f53bc4ddfd773a0ad1d1563ae Mon Sep 17 00:00:00 2001
From: Gravita <12893402+gravit0@users.noreply.github.com>
Date: Sun, 13 Aug 2023 18:07:19 +0700
Subject: [PATCH] [FIX] Auto refresh

---
 .../pro/gravit/launcher/request/Request.java  | 58 +++++++++++--------
 1 file changed, 33 insertions(+), 25 deletions(-)

diff --git a/LauncherAPI/src/main/java/pro/gravit/launcher/request/Request.java b/LauncherAPI/src/main/java/pro/gravit/launcher/request/Request.java
index b8df49b0..b29d76b8 100644
--- a/LauncherAPI/src/main/java/pro/gravit/launcher/request/Request.java
+++ b/LauncherAPI/src/main/java/pro/gravit/launcher/request/Request.java
@@ -12,6 +12,7 @@
 import pro.gravit.utils.helper.LogHelper;
 
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -22,11 +23,11 @@ public abstract class Request<R extends WebSocketEvent> implements WebSocketRequ
     private static final List<ExtendedTokenCallback> extendedTokenCallbacks = new ArrayList<>(4);
     private static final List<BiConsumer<String, AuthRequestEvent.OAuthRequestEvent>> oauthChangeHandlers = new ArrayList<>(4);
 
-    private static RequestService requestService;
-    private static AuthRequestEvent.OAuthRequestEvent oauth;
-    private static Map<String, ExtendedToken> extendedTokens;
-    private static String authId;
-    private static long tokenExpiredTime;
+    private static volatile RequestService requestService;
+    private static volatile AuthRequestEvent.OAuthRequestEvent oauth;
+    private static volatile Map<String, ExtendedToken> extendedTokens;
+    private static volatile String authId;
+    private static volatile long tokenExpiredTime;
     private static ScheduledExecutorService executorService;
     @LauncherNetworkAPI
     public final UUID requestUUID = UUID.randomUUID();
@@ -104,14 +105,14 @@ public static void clearExtendedTokens() {
 
     public static void addExtendedToken(String name, ExtendedToken token) {
         if (extendedTokens == null) {
-            extendedTokens = new HashMap<>();
+            extendedTokens = new ConcurrentHashMap<>();
         }
         extendedTokens.put(name, token);
     }
 
     public static void addAllExtendedToken(Map<String, ExtendedToken> map) {
         if (extendedTokens == null) {
-            extendedTokens = new HashMap<>();
+            extendedTokens = new ConcurrentHashMap<>();
         }
         extendedTokens.putAll(map);
     }
@@ -150,14 +151,17 @@ public static RequestRestoreReport restore() throws Exception {
         return restore(false, false);
     }
 
-    private static Map<String, String> getExpiredExtendedTokens() {
-        Map<String, String> map = new HashMap<>();
+    private synchronized static Map<String, String> getExpiredExtendedTokens() {
+        Set<String> set = new HashSet<>();
         for(Map.Entry<String, ExtendedToken> e : extendedTokens.entrySet()) {
             if(e.getValue().expire != 0 && e.getValue().expire < System.currentTimeMillis()) {
-                map.put(e.getKey(), e.getValue().token);
+                set.add(e.getKey());
             }
         }
-        return map;
+        if(set.isEmpty()) {
+            return new HashMap<>();
+        }
+        return makeNewTokens(set);
     }
 
     public static synchronized RequestRestoreReport restore(boolean needUserInfo, boolean refreshOnly) throws Exception {
@@ -180,19 +184,8 @@ public static synchronized RequestRestoreReport restore(boolean needUserInfo, bo
         RestoreRequestEvent event = request.request();
         List<String> invalidTokens = null;
         if (event.invalidTokens != null && event.invalidTokens.size() > 0) {
-            boolean needRequest = false;
-            Map<String, String> tokens = new HashMap<>();
-            for (ExtendedTokenCallback cb : extendedTokenCallbacks) {
-                for (String tokenName : event.invalidTokens) {
-                    ExtendedToken newToken = cb.tryGetNewToken(tokenName);
-                    if (newToken != null) {
-                        needRequest = true;
-                        tokens.put(tokenName, newToken.token);
-                        addExtendedToken(tokenName, newToken);
-                    }
-                }
-            }
-            if (needRequest) {
+            Map<String, String> tokens = makeNewTokens(event.invalidTokens);
+            if (!tokens.isEmpty()) {
                 request = new RestoreRequest(authId, null, tokens, false);
                 event = request.request();
                 if (event.invalidTokens != null && event.invalidTokens.size() > 0) {
@@ -204,6 +197,20 @@ public static synchronized RequestRestoreReport restore(boolean needUserInfo, bo
         return new RequestRestoreReport(refreshed, invalidTokens, event.userInfo);
     }
 
+    private synchronized static Map<String, String> makeNewTokens(Collection<String> keys) {
+        Map<String, String> tokens = new HashMap<>();
+        for (ExtendedTokenCallback cb : extendedTokenCallbacks) {
+            for (String tokenName : keys) {
+                ExtendedToken newToken = cb.tryGetNewToken(tokenName);
+                if (newToken != null) {
+                    tokens.put(tokenName, newToken.token);
+                    addExtendedToken(tokenName, newToken);
+                }
+            }
+        }
+        return tokens;
+    }
+
     public static void requestError(String message) throws RequestException {
         throw new RequestException(message);
     }
@@ -272,7 +279,8 @@ public static class ExtendedToken {
 
         public ExtendedToken(String token, long expire) {
             this.token = token;
-            this.expire = expire;
+            long time = System.currentTimeMillis();
+            this.expire = expire < time/2 ? time+expire : expire;
         }
     }