From 9d1bd5974e8cedffd96bb59009852f416c7e7d5b Mon Sep 17 00:00:00 2001
From: Rossen Georgiev <rossen@rgp.io>
Date: Mon, 6 May 2019 17:49:45 +0100
Subject: [PATCH] remove workshop lookup + add chunk caching

---
 requirements.txt    |  1 +
 setup.py            |  1 +
 steam/client/cdn.py | 32 +++++++++++++++++---------------
 3 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 08a91c2..def0e70 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,3 +13,4 @@ mock==1.3.0
 PyYAML==5.1
 vcrpy==1.7.4
 sphinx==1.3.5
+cachetools>=3.0.0
diff --git a/setup.py b/setup.py
index 9978cee..af0c1a7 100644
--- a/setup.py
+++ b/setup.py
@@ -22,6 +22,7 @@ install_requires = [
     'gevent>=1.2.0',
     'protobuf>=3.0.0',
     'gevent-eventemitter>=2.0',
+    'cachetools>=3.0.0',
 ]
 
 install_extras = {
diff --git a/steam/client/cdn.py b/steam/client/cdn.py
index 16e3a44..e1af182 100644
--- a/steam/client/cdn.py
+++ b/steam/client/cdn.py
@@ -9,6 +9,7 @@ import logging
 import struct
 
 import vdf
+from cachetools import LRUCache
 from steam import webapi
 from steam.enums import EResult, EServerType
 from steam.util.web import make_requests_session
@@ -99,13 +100,13 @@ class ContentServer(object):
 class CDNClient(object):
     _LOG = logging.getLogger("CDNClient")
     servers = deque()
+    _chunk_cache = LRUCache(20)
 
     def __init__(self, client):
         self.steam = client
         self.web = make_requests_session()
         self.cdn_auth_tokens = {}
         self.depot_keys = {}
-        self.workshop_depots = {}
         self.manifests = {}
         self.app_depots = {}
 
@@ -193,12 +194,12 @@ class CDNClient(object):
             server = self.get_content_server(rotate=True)
 
     def get_chunk(self, app_id, depot_id, chunk_id, cdn_auth_token=None):
-        if cdn_auth_token is None:
-            cdn_auth_token = self.get_cdn_auth_token(depot_id)
+        if (depot_id, chunk_id) not in self._chunk_cache:
+            if cdn_auth_token is None:
+                cdn_auth_token = self.get_cdn_auth_token(depot_id)
 
-        resp = self.get('depot', '%s/chunk/%s' % (depot_id, chunk_id), cdn_auth_token)
+            resp = self.get('depot', '%s/chunk/%s' % (depot_id, chunk_id), cdn_auth_token)
 
-        if resp.ok:
             data = symmetric_decrypt(resp.content, self.get_depot_key(app_id, depot_id))
 
             if data[:2] == b'VZ':
@@ -216,11 +217,15 @@ class CDNClient(object):
                 if crc32(udata) != checksum:
                     raise ValueError("CRC checksum doesn't match for decompressed data")
 
-                return udata[:decompressed_size]
+                data = udata[:decompressed_size]
 
             else:
                 with ZipFile(BytesIO(data)) as zf:
-                    return zf.read(zf.filelist[0])
+                    data = zf.read(zf.filelist[0])
+
+            self._chunk_cache[(depot_id, chunk_id)] = data
+
+        return self._chunk_cache[(depot_id, chunk_id)]
 
     def get_manifest(self, app_id, depot_id, manifest_id, cdn_auth_token=None, decrypt=True):
         if (app_id, depot_id, manifest_id) not in self.manifests:
@@ -301,14 +306,7 @@ class CDNClient(object):
         elif not wf.hcontent_file:
             raise ValueError("Workshop file is not on steampipe")
 
-        app_id = wf.consumer_appid
-
-        ws_app_id = self.workshop_depots.get(app_id)
-
-        if ws_app_id is None:
-            ws_app_id = int(self.steam.get_product_info([app_id])['apps'][app_id]['depots'].get(
-                'workshopdepot', app_id))
-            self.workshop_depots[app_id] = ws_app_id
+        app_id = ws_app_id = wf.consumer_appid
 
         manifest = self.get_manifest(app_id, ws_app_id, wf.hcontent_file)
         manifest.name = wf.title
@@ -379,6 +377,10 @@ class CDNDepotFile(DepotFile):
             'is_directory=True' if self.is_directory else self.size,
             )
 
+    @property
+    def name(self):
+        return self.filename
+
     @property
     def seekable(self):
         return self.is_file