From 43b9d5f1825186cec75965e300d73d07bb54693c Mon Sep 17 00:00:00 2001
From: LamGC <lam827@lamgc.net>
Date: Thu, 11 Jan 2024 00:46:00 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E7=99=BB=E5=BD=95=E5=90=8E=E5=85=88?=
 =?UTF-8?q?=E4=BD=BF=E7=94=A8=E7=94=A8=E6=88=B7=E8=87=AA=E5=B7=B1=E7=9A=84?=
 =?UTF-8?q?=E9=A9=B1=E5=8A=A8=E5=99=A8.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/main/kotlin/OneDriveTransferCenter.kt         |  7 ++++---
 src/main/kotlin/OneDriveTransferSettingManager.kt | 10 +++++++++-
 src/main/kotlin/Services.kt                       | 12 +++++++-----
 3 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/src/main/kotlin/OneDriveTransferCenter.kt b/src/main/kotlin/OneDriveTransferCenter.kt
index c5588c3..c9f0ee0 100644
--- a/src/main/kotlin/OneDriveTransferCenter.kt
+++ b/src/main/kotlin/OneDriveTransferCenter.kt
@@ -16,6 +16,7 @@ import org.telegram.telegrambots.meta.api.methods.GetFile
 import org.telegram.telegrambots.meta.api.methods.send.SendMessage
 import org.telegram.telegrambots.meta.api.methods.updatingmessages.EditMessageText
 import org.telegram.telegrambots.meta.api.objects.Document
+import org.telegram.telegrambots.meta.exceptions.TelegramApiException
 import java.io.File
 import java.io.InputStream
 import java.net.SocketTimeoutException
@@ -57,7 +58,7 @@ object DefaultOneDriveTransferCallback : OneDriveTransferCallback {
             .text("""
                 OneDrive 中转任务执行中
                 文件名： ${progress.currentTask.document.fileName}
-                进度：${progress.progress.get() * 100}%
+                进度：${String.format("%.3f", progress.progress.get() * 100)}%
             """.trimIndent())
             .build())
     }
@@ -158,8 +159,8 @@ class OneDriveTransferTaskExecutor(
                         .build()
                 )
                 break
-            } catch (e: Exception) {
-                if (e !is SocketTimeoutException) {
+            } catch (e: TelegramApiException) {
+                if (e.cause !is SocketTimeoutException) {
                     throw e
                 }
                 if (++retryCount > maxRetryCount) {
diff --git a/src/main/kotlin/OneDriveTransferSettingManager.kt b/src/main/kotlin/OneDriveTransferSettingManager.kt
index 402342b..8afbb83 100644
--- a/src/main/kotlin/OneDriveTransferSettingManager.kt
+++ b/src/main/kotlin/OneDriveTransferSettingManager.kt
@@ -55,10 +55,18 @@ class OneDriveTransferSettingManager(private val authClient: ConfidentialClientA
                 userName = result.account().username()
             }
                 ?: OneDriveTransferSetting.new {
+                    val drive = try {
+                        OneDriveTransferService.createGraphClient(authClient, result.account())
+                            .me().drive()
+                            .buildRequest().get()
+                    } catch (e: Exception) {
+                        null
+                    }
+
                     telegramUserId = userId
                     accountId = result.account().homeAccountId()
                     userName = result.account().username()
-                    driveId = ""
+                    driveId = drive?.id ?: ""
                     storagePath = "Telegram Files/"
                 }
         }
diff --git a/src/main/kotlin/Services.kt b/src/main/kotlin/Services.kt
index 8b4e340..23c26b9 100644
--- a/src/main/kotlin/Services.kt
+++ b/src/main/kotlin/Services.kt
@@ -51,11 +51,7 @@ class OneDriveTransferService(
             authClient.accounts.get().firstOrNull { account ->
                 account.homeAccountId() == it.accountId
             }
-        }?.let {
-            GraphServiceClient.builder()
-                .httpClient(HttpClients.createDefault(MsalAuthorizationProvider(authClient, it)))
-                .buildClient()
-        } ?: throw OneDriveNotLoginException()
+        }?.let { Companion.createGraphClient(authClient, it) } ?: throw OneDriveNotLoginException()
         THREAD_CURRENT_GRAPH_CLIENT.set(ClientCache(userId, serviceClient))
         return serviceClient
     }
@@ -127,6 +123,12 @@ class OneDriveTransferService(
 
     companion object {
         private val THREAD_CURRENT_GRAPH_CLIENT = ThreadLocal<ClientCache>()
+
+        fun createGraphClient(authClient: ConfidentialClientApplication, iAccount: IAccount): GraphServiceClient<Request> {
+            return GraphServiceClient.builder()
+                .httpClient(HttpClients.createDefault(MsalAuthorizationProvider(authClient, iAccount)))
+                .buildClient()
+        }
     }
 }