Add visitorData option to config

lavalink-devs · Aug 9, 2024 · 4aa02b3 · 4aa02b3
1 parent d10c18f
commit 4aa02b3
Show file tree

Hide file tree

Showing 6 changed files with 75 additions and 23 deletions.
diff --git a/README.md b/README.md
@@ -203,9 +203,10 @@ Specifying the token is as simple as doing:
 ```yaml
 plugins:
   youtube:
-    poToken: "paste your token here"
+    pot:
+      token: "paste your po_token here"
+      visitorData: "paste your visitor_data here"
 ```
-You do **not** need the `visitor_data` that is also produced by the script.
 
 > [!NOTE]
 > A `poToken` is not a silver bullet, and currently it only applies to requests made via the `WEB` client.

diff --git a/common/src/main/java/dev/lavalink/youtube/clients/ClientConfig.java b/common/src/main/java/dev/lavalink/youtube/clients/ClientConfig.java
@@ -73,9 +73,31 @@ public ClientConfig withUserAgent(@NotNull String userAgent) {
         return this;
     }
 
-    public ClientConfig withVisitorData(@NotNull String visitorData) {
+    public ClientConfig withVisitorData(@Nullable String visitorData) {
         this.visitorData = visitorData;
-        withClientField("visitorData", visitorData);
+
+        if (visitorData != null) {
+            withClientField("visitorData", visitorData);
+        } else {
+            Map<String, Object> context = (Map<String, Object>) root.get("context");
+
+            if (context != null) {
+                Map<String, Object> client = (Map<String, Object>) context.get("client");
+
+                if (client != null) {
+                    client.remove("visitorData");
+
+                    if (client.isEmpty()) {
+                        context.remove("client");
+                    }
+                }
+
+                if (context.isEmpty()) {
+                    root.remove("context");
+                }
+            }
+        }
+
         return this;
     }
 

diff --git a/common/src/main/java/dev/lavalink/youtube/clients/Web.java b/common/src/main/java/dev/lavalink/youtube/clients/Web.java
@@ -51,16 +51,18 @@ public Web(@NotNull ClientOptions options) {
         this.options = options;
     }
 
-    public static void setPoToken(String poToken) {
+    public static void setPoTokenAndVisitorData(String poToken, String visitorData) {
         Web.poToken = poToken;
 
-        if (poToken == null) {
+        if (poToken == null || visitorData == null) {
             BASE_CONFIG.getRoot().remove("serviceIntegrityDimensions");
+            BASE_CONFIG.withVisitorData(null);
             return;
         }
 
         Map<String, Object> sid = BASE_CONFIG.putOnceAndJoin(BASE_CONFIG.getRoot(), "serviceIntegrityDimensions");
         sid.put("poToken", poToken);
+        BASE_CONFIG.withVisitorData(visitorData);
     }
 
     protected void fetchClientConfig(@NotNull HttpInterface httpInterface) {
@@ -109,12 +111,11 @@ protected void fetchClientConfig(@NotNull HttpInterface httpInterface) {
                     BASE_CONFIG.withClientField("clientVersion", clientVersion);
                 }
 
-                String visitorData = client.get("visitorData").text();
-
-                if (visitorData != null && !visitorData.isEmpty() && BASE_CONFIG.getVisitorData() == null) {
-                    // don't overwrite if visitorData was already set.
-                    BASE_CONFIG.withVisitorData(visitorData);
-                }
+//                String visitorData = client.get("visitorData").text();
+//
+//                if (visitorData != null && !visitorData.isEmpty()) {
+//                    BASE_CONFIG.withVisitorData(visitorData);
+//                }
             }
         } catch (IOException e) {
             throw ExceptionTools.toRuntimeException(e);

diff --git a/plugin/src/main/java/dev/lavalink/youtube/plugin/Pot.java b/plugin/src/main/java/dev/lavalink/youtube/plugin/Pot.java
@@ -0,0 +1,22 @@
+package dev.lavalink.youtube.plugin;
+
+public class Pot {
+    private String token;
+    private String visitorData;
+
+    public String getToken() {
+        return token != null && !token.isEmpty() ? token : null;
+    }
+
+    public String getVisitorData() {
+        return visitorData != null && !visitorData.isEmpty() ? visitorData : null;
+    }
+
+    public void setPoToken(String token) {
+        this.token = token;
+    }
+
+    public void setVisitorData(String visitorData) {
+        this.visitorData = visitorData;
+    }
+}
diff --git a/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubeConfig.java b/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubeConfig.java
@@ -14,7 +14,7 @@ public class YoutubeConfig {
     private boolean allowSearch = true;
     private boolean allowDirectVideoIds = true;
     private boolean allowDirectPlaylistIds = true;
-    private String poToken;
+    private Pot pot = null;
     private String[] clients;
     private Map<String, ClientOptions> clientOptions = new HashMap<>();
 
@@ -34,8 +34,8 @@ public boolean getAllowDirectPlaylistIds() {
         return allowDirectPlaylistIds;
     }
 
-    public String getPoToken() {
-        return poToken;
+    public Pot getPot() {
+        return pot;
     }
 
     public String[] getClients() {
@@ -62,8 +62,8 @@ public void setAllowDirectPlaylistIds(boolean allowDirectPlaylistIds) {
         this.allowDirectPlaylistIds = allowDirectPlaylistIds;
     }
 
-    public void setPoToken(String poToken) {
-        this.poToken = poToken;
+    public void setPot(Pot pot) {
+        this.pot = pot;
     }
 
     public void setClients(String[] clients) {

diff --git a/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubePluginLoader.java b/plugin/src/main/java/dev/lavalink/youtube/plugin/YoutubePluginLoader.java
@@ -8,7 +8,6 @@
 import com.sedmelluq.lava.extensions.youtuberotator.tools.ip.Ipv6Block;
 import dev.arbjerg.lavalink.api.AudioPlayerManagerConfiguration;
 import dev.lavalink.youtube.YoutubeAudioSourceManager;
-import dev.lavalink.youtube.clients.ClientConfig;
 import dev.lavalink.youtube.clients.ClientOptions;
 import dev.lavalink.youtube.clients.Web;
 import dev.lavalink.youtube.clients.skeleton.Client;
@@ -163,11 +162,18 @@ public AudioPlayerManager configure(AudioPlayerManager audioPlayerManager) {
                 clients = clientProvider.getDefaultClients();
             } else {
                 clients = youtubeConfig.getClients();
-                String poToken = youtubeConfig.getPoToken();
-
-                if (poToken != null && !poToken.isEmpty()) {
-                    log.debug("Setting poToken for WEB client to {}", poToken);
-                    Web.setPoToken(poToken);
+                Pot pot = youtubeConfig.getPot();
+
+                if (pot != null) {
+                    String token = pot.getToken();
+                    String visitorData = pot.getVisitorData();
+
+                    if (token != null && visitorData != null) {
+                        log.debug("Applying poToken and visitorData to WEB client (token: {}, vd: {})", token, visitorData);
+                        Web.setPoTokenAndVisitorData(token, visitorData);
+                    } else if (token != null || visitorData != null) {
+                        log.warn("Both pot.token and pot.visitorData must be specified and valid for pot to apply.");
+                    }
                 }
             }