cleanups, usability, etc.

lavalink-devs · Jul 11, 2024 · de3b7dc · de3b7dc
1 parent f05dd6a
commit de3b7dc
Show file tree

Hide file tree

Showing 5 changed files with 92 additions and 5 deletions.
diff --git a/common/src/main/java/dev/lavalink/youtube/http/YoutubeOauth2Handler.java b/common/src/main/java/dev/lavalink/youtube/http/YoutubeOauth2Handler.java
@@ -4,6 +4,7 @@
 import com.grack.nanojson.JsonParser;
 import com.grack.nanojson.JsonParserException;
 import com.grack.nanojson.JsonWriter;
+import com.sedmelluq.discord.lavaplayer.tools.DataFormatTools;
 import com.sedmelluq.discord.lavaplayer.tools.ExceptionTools;
 import com.sedmelluq.discord.lavaplayer.tools.io.HttpClientTools;
 import com.sedmelluq.discord.lavaplayer.tools.io.HttpInterface;
@@ -14,7 +15,6 @@
 import org.apache.http.client.protocol.HttpClientContext;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
-import org.apache.http.util.EntityUtils;
 import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -52,7 +52,9 @@ public void setRefreshToken(@Nullable String refreshToken) {
         this.refreshToken = refreshToken;
         this.tokenExpires = System.currentTimeMillis(); // to trigger an access token refresh
 
-        if (refreshToken == null) {
+        // TODO: need to check what error is returned for invalid refresh tokens and fall back to
+        //       initialization if invalid.
+        if (DataFormatTools.isNullOrEmpty(refreshToken)) {
             initializeAccessToken();
         }
     }
@@ -173,8 +175,8 @@ private void pollForToken(String deviceCode) {
     }
 
     private void refreshAccessToken() {
-        if (refreshToken == null) {
-            throw new IllegalStateException("Cannot refresh access token without a refresh token!");
+        if (DataFormatTools.isNullOrEmpty(refreshToken)) {
+            throw new IllegalStateException("Cannot fetch access token without a refresh token!");
         }
 
         // @formatter:off
@@ -214,7 +216,7 @@ private void refreshAccessToken() {
     }
 
     public void applyToken(HttpUriRequest request) {
-        if (!enabled || refreshToken == null) {
+        if (!enabled || DataFormatTools.isNullOrEmpty(refreshToken)) {
             return;
         }
 

diff --git a/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubeConfig.java b/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubeConfig.java
@@ -17,6 +17,9 @@ public class YoutubeConfig {
     private String[] clients;
     private Map<String, ClientOptions> clientOptions = new HashMap<>();
 
+    private boolean useOauth2 = false;
+    private String oauth2RefreshToken = null;
+
     public boolean getEnabled() {
         return this.enabled;
     }
@@ -41,6 +44,14 @@ public Map<String, ClientOptions> getClientOptions() {
         return this.clientOptions;
     }
 
+    public boolean getUseOauth2() {
+        return this.useOauth2;
+    }
+
+    public String getOauth2RefreshToken() {
+        return this.oauth2RefreshToken;
+    }
+
     public void setEnabled(boolean enabled) {
         this.enabled = enabled;
     }
@@ -64,4 +75,12 @@ public void setClients(String[] clients) {
     public void setClientOptions(Map<String, ClientOptions> clientOptions) {
         this.clientOptions = clientOptions;
     }
+
+    public void setUseOauth2(boolean useOauth2) {
+        this.useOauth2 = useOauth2;
+    }
+
+    public void setOauth2RefreshToken(String oauth2RefreshToken) {
+        this.oauth2RefreshToken = oauth2RefreshToken;
+    }
 }
diff --git a/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubeOauthConfig.java b/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubeOauthConfig.java
@@ -0,0 +1,9 @@
+package dev.lavalink.youtube.plugin;
+
+public class YoutubeOauthConfig {
+    private String refreshToken;
+
+    public String getRefreshToken() {
+        return refreshToken;
+    }
+}
diff --git a/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubePluginLoader.java b/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubePluginLoader.java
@@ -184,6 +184,12 @@ public AudioPlayerManager configure(AudioPlayerManager audioPlayerManager) {
             source.setPlaylistPageCount(playlistLoadLimit);
         }
 
+        if (youtubeConfig != null) {
+            if (youtubeConfig.getUseOauth2()) {
+                source.useOauth2(youtubeConfig.getOauth2RefreshToken());
+            }
+        }
+
         audioPlayerManager.registerSourceManager(source);
         return audioPlayerManager;
     }

diff --git a/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubeRestHandler.java b/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubeRestHandler.java
@@ -0,0 +1,51 @@
+package dev.lavalink.youtube.plugin;
+
+import com.sedmelluq.discord.lavaplayer.player.AudioPlayerManager;
+import dev.lavalink.youtube.YoutubeAudioSourceManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.http.HttpStatus;
+import org.springframework.stereotype.Service;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestBody;
+import org.springframework.web.bind.annotation.ResponseStatus;
+import org.springframework.web.server.ResponseStatusException;
+
+import java.util.Collections;
+import java.util.Map;
+
+@Service
+public class YoutubeRestHandler {
+    private static final Logger log = LoggerFactory.getLogger(YoutubeRestHandler.class);
+
+    private final AudioPlayerManager playerManager;
+
+    public YoutubeRestHandler(AudioPlayerManager playerManager) {
+        this.playerManager = playerManager;
+    }
+
+    @GetMapping("/v4/youtube")
+    public Map<String, String> getYoutubeConfig() {
+        YoutubeAudioSourceManager source = playerManager.source(YoutubeAudioSourceManager.class);
+
+        if (source == null) {
+            throw new ResponseStatusException(HttpStatus.BAD_REQUEST, "The YouTube source manager is not registered.");
+        }
+
+        return Collections.singletonMap("refreshToken", source.getOauth2RefreshToken());
+    }
+
+    @PostMapping("/v4/youtube")
+    @ResponseStatus(HttpStatus.NO_CONTENT)
+    public void updateRefreshToken(@RequestBody YoutubeOauthConfig config) {
+        YoutubeAudioSourceManager source = playerManager.source(YoutubeAudioSourceManager.class);
+
+        if (source == null) {
+            throw new ResponseStatusException(HttpStatus.BAD_REQUEST, "The YouTube source manager is not registered.");
+        }
+
+        source.useOauth2(config.getRefreshToken());
+        log.debug("Updated YouTube OAuth2 refresh token to {}", config.getRefreshToken());
+    }
+}