Allow clients to transform playback URIs as needed.

lavalink-devs · Aug 9, 2024 · 8aad32d · 8aad32d
1 parent 7fef414
commit 8aad32d
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 10 deletions.
diff --git a/common/src/main/java/dev/lavalink/youtube/clients/Web.java b/common/src/main/java/dev/lavalink/youtube/clients/Web.java
@@ -9,13 +9,16 @@
 import dev.lavalink.youtube.clients.skeleton.StreamingNonMusicClient;
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpGet;
+import org.apache.http.client.utils.URIBuilder;
 import org.apache.http.util.EntityUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -132,6 +135,25 @@ public ClientConfig getBaseClientConfig(@NotNull HttpInterface httpInterface) {
         return BASE_CONFIG.copy();
     }
 
+    @Override
+    @NotNull
+    public URI transformPlaybackUri(@NotNull URI originalUri, @NotNull URI resolvedPlaybackUri) {
+        if (poToken == null) {
+            return resolvedPlaybackUri;
+        }
+
+        log.debug("Applying 'pot' parameter on playback URI: {}", resolvedPlaybackUri);
+        URIBuilder builder = new URIBuilder(resolvedPlaybackUri);
+        builder.addParameter("pot", poToken);
+
+        try {
+            return builder.build();
+        } catch (URISyntaxException e) {
+            log.debug("Failed to apply 'pot' parameter.", e);
+            return resolvedPlaybackUri;
+        }
+    }
+
     @Override
     @NotNull
     protected List<AudioTrack> extractSearchResults(@NotNull YoutubeAudioSourceManager source,

diff --git a/common/src/main/java/dev/lavalink/youtube/clients/skeleton/Client.java b/common/src/main/java/dev/lavalink/youtube/clients/skeleton/Client.java
@@ -14,6 +14,7 @@
 import org.jetbrains.annotations.Nullable;
 
 import java.io.IOException;
+import java.net.URI;
 import java.util.List;
 import java.util.stream.Collectors;
 
@@ -133,6 +134,21 @@ default AudioTrack findSelectedTrack(@NotNull List<AudioTrack> tracks,
         return null;
     }
 
+    /**
+     * Transforms a given playback URL as necessary. For example, you can add query parameters
+     * or resolve any challenge parameters that might be needed.
+     * @param originalUri The original stream URI. This will be completely unmodified, and is
+     *                    provided as it has been received from YouTube.
+     * @param resolvedPlaybackUri The playback URI. This will have already been
+     *                            transformed by the SignatureCipherManager.
+     * @return The new playback URI.
+     */
+    @NotNull
+    default URI transformPlaybackUri(@NotNull URI originalUri,
+                                     @NotNull URI resolvedPlaybackUri) {
+        return resolvedPlaybackUri;
+    }
+
     /**
      * Builds an audio track with the given parameters.
      * Hint: You can use {@link YoutubeAudioSourceManager#buildAudioTrack(AudioTrackInfo)} to

diff --git a/common/src/main/java/dev/lavalink/youtube/track/YoutubeAudioTrack.java b/common/src/main/java/dev/lavalink/youtube/track/YoutubeAudioTrack.java
@@ -15,11 +15,9 @@
 import dev.lavalink.youtube.UrlTools;
 import dev.lavalink.youtube.UrlTools.UrlInfo;
 import dev.lavalink.youtube.YoutubeAudioSourceManager;
-import dev.lavalink.youtube.clients.Web;
 import dev.lavalink.youtube.clients.skeleton.Client;
 import dev.lavalink.youtube.track.format.StreamFormat;
 import dev.lavalink.youtube.track.format.TrackFormats;
-import org.apache.http.client.utils.URIBuilder;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
@@ -190,17 +188,12 @@ private FormatWithUrl loadBestFormatWithUrl(@NotNull HttpInterface httpInterface
 
     StreamFormat format = formats.getBestFormat();
 
-    URI signedUrl = sourceManager.getCipherManager()
+    URI resolvedUrl = sourceManager.getCipherManager()
         .resolveFormatUrl(httpInterface, formats.getPlayerScriptUrl(), format);
 
-    if (client.getIdentifier().equals(Web.BASE_CONFIG.getName()) && Web.poToken != null) {
-      log.debug("Format origin is {}, setting 'pot' parameter.", client.getIdentifier());
-      URIBuilder builder = new URIBuilder(signedUrl);
-      builder.addParameter("pot", Web.poToken);
-      signedUrl = builder.build();
-    }
+    resolvedUrl = client.transformPlaybackUri(format.getUrl(), resolvedUrl);
 
-    return new FormatWithUrl(format, signedUrl);
+    return new FormatWithUrl(format, resolvedUrl);
   }
 
   @Override