stream: harden streaming redirectable urls

TAAPArthur · Aug 13, 2023 · 2d6b603 · 2d6b603
1 parent 44217e7
commit 2d6b603
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git a/amt/media_reader.py b/amt/media_reader.py
@@ -188,16 +188,16 @@ def search_for_media(self, name, media_type=None, exact=False, server_id=None, s
             return False
         return media_data
 
-    def get_server_for_url(self, url, streamable=False, server_id=None):
-        for server in self.get_servers():
-            func = server.can_stream_url if streamable else server.can_add_media_from_url
-            if server_id in (None, server.id):
-                if func(url):
-                    return url, server
-                elif server.domain and server.domain in url:
-                    new_url = server.get_redirect_url(url)
-                    if func(new_url):
-                        return new_url, server
+    def get_server_for_url(self, url, streamable=False, server_id=None, disallow_redirects=False):
+        servers = list(filter(lambda server: server_id in (None, server.id), self.get_servers()))
+        for server in servers:
+            if (server.can_stream_url if streamable else server.can_add_media_from_url)(url):
+                return url, server
+        if not disallow_redirects:
+            for server in filter(lambda server: server.domain and server.domain in url, servers):
+                new_url = server.get_redirect_url(url)
+                if new_url != url:
+                    return self.get_server_for_url(new_url, streamable=streamable, server_id=server_id, disallow_redirects=True)
         return url, None
 
     def get_media_from_url(self, url, server_id=None):

diff --git a/amt/tests/test.py b/amt/tests/test.py
@@ -2322,7 +2322,7 @@ class ServerStreamTest(RealBaseUnitTestClass):
     ]
 
     def test_a_verify_valid_stream_urls(self):
-        for streamable, url_list in [(True, self.streamable_urls), (False, self.addable_urls)]:
+        for streamable, url_list in [(True, self.streamable_urls), (False, self.streamable_urls), (False, self.addable_urls)]:
             for url_data in url_list:
                 sample_url = url_data[0]
                 for url in {sample_url, sample_url[:-1] if sample_url[-1] == "/" else sample_url, sample_url.replace("https:", "http:"), sample_url.replace("://", "://www.")}: