From afd3f1618198dec831f2ea6ae843b61eb2bbf0e8 Mon Sep 17 00:00:00 2001
From: Fornax <wim@fornaxian.tech>
Date: Thu, 25 Apr 2024 14:50:42 +0200
Subject: [PATCH] Use websocket proto for getting file stats and rate limits.
 Remove view_token

---
 svelte/src/file_viewer/DownloadLimitStore.js  |  65 ----------
 svelte/src/file_viewer/FileStats.svelte       |  96 +++------------
 svelte/src/file_viewer/FileViewer.svelte      |   5 +-
 svelte/src/file_viewer/StatsSocket.js         | 115 ++++++++++++++++++
 svelte/src/file_viewer/TransferLimit.svelte   |  14 +--
 .../file_viewer/viewers/BandwidthUsage.svelte |  12 +-
 .../file_viewer/viewers/FilePreview.svelte    |   4 +-
 .../src/file_viewer/viewers/RateLimit.svelte  |  10 +-
 .../src/file_viewer/viewers/SpeedLimit.svelte |   6 +-
 webcontroller/file_viewer.go                  |  11 --
 10 files changed, 159 insertions(+), 179 deletions(-)
 delete mode 100644 svelte/src/file_viewer/DownloadLimitStore.js
 create mode 100644 svelte/src/file_viewer/StatsSocket.js

diff --git a/svelte/src/file_viewer/DownloadLimitStore.js b/svelte/src/file_viewer/DownloadLimitStore.js
deleted file mode 100644
index 79615b6..0000000
--- a/svelte/src/file_viewer/DownloadLimitStore.js
+++ /dev/null
@@ -1,65 +0,0 @@
-import { readable } from 'svelte/store';
-
-let limits = {
-	download_limit: 0,
-	download_limit_used: 0,
-	transfer_limit: 0,
-	transfer_limit_used: 0,
-	ipv6_bonus: 0,
-	loaded: false,
-}
-
-export const download_limits = readable(
-	limits,
-	function start(set) {
-		set_func = set;
-		loop();
-
-		return function stop() {
-			clearTimeout(timeout_ref);
-		};
-	}
-);
-
-let set_func;
-let timeout_ref = 0;
-let timeout_ms = 10000;
-
-const loop = async () => {
-	let new_limits;
-	try {
-		let resp = await fetch(window.api_endpoint + "/misc/rate_limits")
-		if (resp.status >= 400) {
-			throw new Error(await resp.text())
-		}
-		new_limits = await resp.json()
-		new_limits.loaded = true
-	} catch (err) {
-		console.error("Failed to get rate limits: " + err)
-		timeout_ref = setTimeout(loop, 30000)
-		return
-	}
-
-	// If the usage has not changed we increase the timeout with one second. If
-	// it did change we halve the timeout
-	if (
-		new_limits.transfer_limit_used === limits.transfer_limit_used &&
-		new_limits.transfer_limit === limits.transfer_limit
-	) {
-		timeout_ms += 2000
-		if (timeout_ms > 60000) {
-			timeout_ms = 60000
-		}
-	} else {
-		timeout_ms /= 2
-		if (timeout_ms < 5000) {
-			timeout_ms = 5000
-		}
-
-		limits = new_limits
-		set_func(new_limits)
-	}
-
-	console.debug("Sleep", timeout_ms / 1000, "seconds")
-	timeout_ref = setTimeout(loop, timeout_ms)
-}
diff --git a/svelte/src/file_viewer/FileStats.svelte b/svelte/src/file_viewer/FileStats.svelte
index 95ad03e..9c49795 100644
--- a/svelte/src/file_viewer/FileStats.svelte
+++ b/svelte/src/file_viewer/FileStats.svelte
@@ -1,6 +1,6 @@
 <script>
-import { onMount } from "svelte";
 import { formatDataVolume, formatThousands } from "../util/Formatting.svelte"
+import { set_file, stats } from "./StatsSocket"
 
 export let file = {
 	id: "",
@@ -10,98 +10,42 @@ export let file = {
 	bandwidth_used: 0,
 	bandwidth_used_paid: 0,
 }
-export let view_token = ""
 
 let views = 0
 let downloads = 0
 let size = 0
-let socket = null
-let error_msg = ""
-
-$: update_stats(file.id)
-let update_stats = async id => {
-	if (id === "" || id == "demo") {
-		return
-	}
-
-	views = file.views
-	if (file.size === 0) {
-		downloads = file.downloads
-	} else {
-		downloads = Math.round((file.bandwidth_used + file.bandwidth_used_paid) / file.size)
-	}
+$: {
 	size = file.size
 
-	send_watch_command()
-}
+	if ($stats.file_stats_init) {
+		views = $stats.file_stats.views
 
-const send_watch_command = () => {
-	if (socket !== null && socket.readyState === WebSocket.OPEN) {
-		socket.send(
-			JSON.stringify(
-				{cmd: "watch_file", a1: file.id, a2: view_token}
-			)
-		)
-	}
-}
-
-const init_socket = () => {
-	if (socket !== null && socket.readyState !== WebSocket.CLOSED) {
-		return
-	}
-
-	console.log("initializing socket")
-	socket = new WebSocket(location.origin.replace(/^http/, 'ws') + "/api/file_stats")
-
-	socket.onopen = () => send_watch_command()
-	socket.onmessage = msg => {
-		let j = JSON.parse(msg.data)
-		console.debug("WS update", j)
-
-		error_msg = ""
-		views = j.views
 		if (file.size === 0) {
-			downloads = j.downloads
+			downloads = $stats.file_stats.downloads
 		} else {
-			downloads = Math.round((j.bandwidth + j.bandwidth_paid) / file.size)
+			downloads = Math.round(($stats.file_stats.bandwidth + $stats.file_stats.bandwidth_paid) / file.size)
 		}
-	}
-	socket.onerror = err => {
-		if (socket === null) {
-			return
-		}
-		console.error("WS error", err)
-		socket.close()
-		socket = null
-		error_msg = "failed to get stats, retrying..."
+	} else {
+		views = file.views
 
-		window.setTimeout(init_socket, 2000)
+		if (file.size === 0) {
+			downloads = file.downloads
+		} else {
+			downloads = Math.round((file.bandwidth_used + file.bandwidth_used_paid) / file.size)
+		}
 	}
 }
 
-onMount(() => {
-	init_socket()
-
-	return () => {
-		if (socket !== null) {
-			socket.close()
-			socket = null
-		}
-	}
-})
+$: set_file(file.id)
 </script>
 
 <div>
-	{#if error_msg !== ""}
-		{error_msg}
-	{:else}
-		<div class="label">Views</div>
-		<div class="stat">{formatThousands(views)}</div>
-		<div class="label">Downloads</div>
-		<div class="stat">{formatThousands(downloads)}</div>
-		<div class="label">Size</div>
-		<div class="stat">{formatDataVolume(size, 3)}</div>
-	{/if}
+	<div class="label">Views</div>
+	<div class="stat">{formatThousands(views)}</div>
+	<div class="label">Downloads</div>
+	<div class="stat">{formatThousands(downloads)}</div>
+	<div class="label">Size</div>
+	<div class="stat">{formatDataVolume(size, 3)}</div>
 </div>
 
 <style>
diff --git a/svelte/src/file_viewer/FileViewer.svelte b/svelte/src/file_viewer/FileViewer.svelte
index 1d207df..70f00c5 100644
--- a/svelte/src/file_viewer/FileViewer.svelte
+++ b/svelte/src/file_viewer/FileViewer.svelte
@@ -23,7 +23,6 @@ import CopyButton from "../layout/CopyButton.svelte";
 
 let loading = true
 let embedded = false
-let view_token = ""
 let ads_enabled = false
 
 let view = "" // file or gallery
@@ -88,8 +87,6 @@ onMount(() => {
 		toolbar_visible = false
 	}
 
-	view_token = viewer_data.view_token
-
 	if (viewer_data.type === "list") {
 		open_list(viewer_data.api_response)
 	} else {
@@ -418,7 +415,7 @@ const keyboard_event = evt => {
 	<div class="file_preview_row">
 		<div class="toolbar" class:toolbar_visible>
 			{#if view === "file"}
-				<FileStats file={file} view_token={view_token}/>
+				<FileStats file={file}/>
 			{:else if view === "gallery"}
 				<ListStats list={list}/>
 			{/if}
diff --git a/svelte/src/file_viewer/StatsSocket.js b/svelte/src/file_viewer/StatsSocket.js
new file mode 100644
index 0000000..e2b2a6e
--- /dev/null
+++ b/svelte/src/file_viewer/StatsSocket.js
@@ -0,0 +1,115 @@
+import { readable } from "svelte/store";
+
+let results = {
+	connected: false,
+	file_stats_init: false,
+	file_stats: {
+		views: 0,
+		downloads: 0,
+		bandwidth: 0,
+		bandwidth_paid: 0,
+	},
+	limits_init: false,
+	limits: {
+		download_limit: 0,
+		download_limit_used: 0,
+		transfer_limit: 0,
+		transfer_limit_used: 0,
+	},
+}
+
+export const stats = readable(
+	results,
+	(set) => {
+		start_sock(set)
+		return stop_sock
+	},
+);
+
+let socket = null
+const start_sock = (set_func) => {
+	if (socket !== null) {
+		return
+	}
+
+	console.log("initializing stats socket")
+	socket = new WebSocket(location.origin.replace(/^http/, 'ws') + "/api/file_stats")
+
+	socket.onopen = () => {
+		results.connected = true
+		set_func(results)
+
+		// Subscribe to the rate limit feed. This will also process any queued
+		// commands built up while the socket was down
+		send_cmd({ type: "limits" })
+	}
+	socket.onmessage = msg => {
+		let j = JSON.parse(msg.data)
+		console.debug("WS update", j)
+
+		if (j.type === "file_stats") {
+			results.file_stats = j.file_stats
+			results.file_stats_init = true
+			set_func(results)
+		} else if (j.type === "limits") {
+			results.limits = j.limits
+			results.limits_init = true
+			set_func(results)
+		} else {
+			console.error("Unknown ws message type", j.type, "data", msg.data)
+		}
+	}
+	socket.onerror = err => {
+		console.error("socket error", err)
+		stop_sock(set_func)
+		window.setTimeout(() => start_sock(set_func), 2000)
+	}
+	socket.onclose = () => {
+		stop_sock(set_func)
+		window.setTimeout(() => start_sock(set_func), 2000)
+	}
+}
+
+const stop_sock = (set_func) => {
+	if (socket === null) {
+		return
+	}
+
+	// Prevent error handlers from re-initializing the socket
+	socket.onerror = null
+	socket.onclose = null
+
+	// Close and delete the socket
+	socket.close()
+	socket = null
+
+	// Reset the state
+	results.connected = false
+	results.file_stats_init = false
+	results.limits_init = false
+	set_func(results)
+}
+
+export const set_file = file_id => {
+	send_cmd({
+		type: "file_stats",
+		data: { file_id: file_id },
+	})
+}
+
+let queued_commands = []
+const send_cmd = cmd => {
+	if (socket !== null && socket.readyState === WebSocket.OPEN) {
+
+		// First empty the queue
+		while (queued_commands.length !== 0) {
+			socket.send(JSON.stringify(queued_commands.shift()))
+		}
+
+		// Send the requested command
+		socket.send(JSON.stringify(cmd))
+	} else if (cmd !== null) {
+		queued_commands.push(cmd)
+		console.debug("Socket is closed, command", cmd, "added to queue")
+	}
+}
diff --git a/svelte/src/file_viewer/TransferLimit.svelte b/svelte/src/file_viewer/TransferLimit.svelte
index 39ea5fb..db815f5 100644
--- a/svelte/src/file_viewer/TransferLimit.svelte
+++ b/svelte/src/file_viewer/TransferLimit.svelte
@@ -1,28 +1,28 @@
 <script>
 import { formatDataVolume } from "../util/Formatting.svelte";
-import { download_limits } from "./DownloadLimitStore.js"
+import { stats } from "./StatsSocket.js"
 
 let percent = 0
 let title = ""
 $: {
-	if ($download_limits.transfer_limit === 0) {
+	if ($stats.limits.transfer_limit === 0) {
 		percent = 0 // Avoid division by 0
-	}else if ($download_limits.transfer_limit_used / $download_limits.transfer_limit > 1) {
+	} else if ($stats.limits.transfer_limit_used / $stats.limits.transfer_limit > 1) {
 		percent = 100
 	} else {
-		percent = ($download_limits.transfer_limit_used / $download_limits.transfer_limit) * 100
+		percent = ($stats.limits.transfer_limit_used / $stats.limits.transfer_limit) * 100
 	}
 
 	title = "Download limit used: " +
-		formatDataVolume($download_limits.transfer_limit_used, 3) +
+		formatDataVolume($stats.limits.transfer_limit_used, 3) +
 		" of " +
-		formatDataVolume($download_limits.transfer_limit, 3);
+		formatDataVolume($stats.limits.transfer_limit, 3);
 }
 </script>
 
 <!-- Always show the outer bar to prevent layout shift -->
 <div class="progress_bar_outer" title="{title}">
-	{#if $download_limits.loaded}
+	{#if $stats.limits_init}
 		<div class="progress_bar_text">
 			{title}
 		</div>
diff --git a/svelte/src/file_viewer/viewers/BandwidthUsage.svelte b/svelte/src/file_viewer/viewers/BandwidthUsage.svelte
index cd84596..769f7ab 100644
--- a/svelte/src/file_viewer/viewers/BandwidthUsage.svelte
+++ b/svelte/src/file_viewer/viewers/BandwidthUsage.svelte
@@ -2,16 +2,16 @@
 import { formatDataVolume } from "../../util/Formatting.svelte";
 import TextBlock from "./TextBlock.svelte";
 import ProgressBar from "../../util/ProgressBar.svelte";
-import { download_limits } from "../DownloadLimitStore";
+import { stats } from "../StatsSocket"
 
 export let file = {
 	size: 0,
 }
 
-$: transfer_left = $download_limits.transfer_limit - $download_limits.transfer_limit_used
+$: transfer_left = $stats.limits.transfer_limit - $stats.limits.transfer_limit_used
 </script>
 
-{#if $download_limits.loaded}
+{#if $stats.limits_init}
 	<TextBlock center={true}>
 		{#if file.size > transfer_left}
 			<div class="highlight_yellow">
@@ -24,8 +24,8 @@ $: transfer_left = $download_limits.transfer_limit - $download_limits.transfer_l
 		{/if}
 
 		<p>
-			You have used {formatDataVolume($download_limits.transfer_limit_used, 3)} of
-			your daily {formatDataVolume($download_limits.transfer_limit, 3)} transfer
+			You have used {formatDataVolume($stats.limits.transfer_limit_used, 3)} of
+			your daily {formatDataVolume($stats.limits.transfer_limit, 3)} transfer
 			limit. When the transfer limit is exceeded your download speed will
 			be reduced.
 		</p>
@@ -39,6 +39,6 @@ $: transfer_left = $download_limits.transfer_limit - $download_limits.transfer_l
 			</strong>
 		</p>
 
-		<ProgressBar total={$download_limits.transfer_limit} used={$download_limits.transfer_limit_used}></ProgressBar>
+		<ProgressBar total={$stats.limits.transfer_limit} used={$stats.limits.transfer_limit_used}></ProgressBar>
 	</TextBlock>
 {/if}
diff --git a/svelte/src/file_viewer/viewers/FilePreview.svelte b/svelte/src/file_viewer/viewers/FilePreview.svelte
index dd4d6ff..9d39606 100644
--- a/svelte/src/file_viewer/viewers/FilePreview.svelte
+++ b/svelte/src/file_viewer/viewers/FilePreview.svelte
@@ -12,7 +12,7 @@ import { file_type } from "../FileUtilities.svelte";
 import RateLimit from "./RateLimit.svelte";
 import Torrent from "./Torrent.svelte";
 import SpeedLimit from "./SpeedLimit.svelte";
-import { download_limits } from "../DownloadLimitStore";
+import { stats } from "../StatsSocket";
 import Zip from "./Zip.svelte";
 
 let viewer
@@ -60,7 +60,7 @@ export const toggle_playback = () => {
 	</div>
 {:else if viewer_type === "abuse"}
 	<Abuse bind:this={viewer} on:download></Abuse>
-{:else if !premium_download && $download_limits.transfer_limit_used > $download_limits.transfer_limit}
+{:else if !premium_download && $stats.limits.transfer_limit_used > $stats.limits.transfer_limit}
 	<SpeedLimit file={current_file} on:download></SpeedLimit>
 {:else if viewer_type === "rate_limit"}
 	<RateLimit bind:this={viewer} on:download></RateLimit>
diff --git a/svelte/src/file_viewer/viewers/RateLimit.svelte b/svelte/src/file_viewer/viewers/RateLimit.svelte
index 8a31612..e932f82 100644
--- a/svelte/src/file_viewer/viewers/RateLimit.svelte
+++ b/svelte/src/file_viewer/viewers/RateLimit.svelte
@@ -1,7 +1,7 @@
 <script>
 import { createEventDispatcher } from "svelte";
 import { formatDataVolume } from "../../util/Formatting.svelte";
-import { download_limits } from "../DownloadLimitStore";
+import { stats } from "../StatsSocket";
 import IconBlock from "./IconBlock.svelte";
 import TextBlock from "./TextBlock.svelte";
 let dispatch = createEventDispatcher()
@@ -36,15 +36,15 @@ let file = {
 		</h1>
 		<p>
 			You have reached your download limit for today. Without a pixeldrain
-			account you are limited to downloading {$download_limits.download_limit} files
-			or {formatDataVolume($download_limits.transfer_limit, 3)} per 48 hours. This limit
+			account you are limited to downloading {$stats.limits.download_limit} files
+			or {formatDataVolume($stats.limits.transfer_limit, 3)} per 48 hours. This limit
 			is counted per IP address, so if you're on a shared network it's
 			possible that others have also contributed to this limit.
 		</p>
 		<p>
 			In the last 24 hours you have downloaded
-			{$download_limits.download_limit_used} files and used
-			{formatDataVolume($download_limits.transfer_limit_used, 3)} bandwidth.
+			{$stats.limits.download_limit_used} files and used
+			{formatDataVolume($stats.limits.transfer_limit_used, 3)} bandwidth.
 		</p>
 	{/if}
 	<p>
diff --git a/svelte/src/file_viewer/viewers/SpeedLimit.svelte b/svelte/src/file_viewer/viewers/SpeedLimit.svelte
index 506c299..a0d7e18 100644
--- a/svelte/src/file_viewer/viewers/SpeedLimit.svelte
+++ b/svelte/src/file_viewer/viewers/SpeedLimit.svelte
@@ -1,7 +1,7 @@
 <script>
 import { createEventDispatcher } from "svelte";
 import { formatDataVolume, formatDuration } from "../../util/Formatting.svelte";
-import { download_limits } from "../DownloadLimitStore";
+import { stats } from "../StatsSocket";
 import IconBlock from "./IconBlock.svelte";
 import TextBlock from "./TextBlock.svelte";
 let dispatch = createEventDispatcher()
@@ -20,10 +20,10 @@ export let file = {
 	<p>
 		Pixeldrain's free tier is supported by my Patrons (be grateful). There's
 		only so much that you can do with the budget they provide.
-		{formatDataVolume($download_limits.transfer_limit, 3)} per day is about
+		{formatDataVolume($stats.limits.transfer_limit, 3)} per day is about
 		the most I can give away for free while keeping it fair for everyone,
 		and according to our records you have already downloaded
-		{formatDataVolume($download_limits.transfer_limit_used, 3)}.
+		{formatDataVolume($stats.limits.transfer_limit_used, 3)}.
 	</p>
 	<p>
 		It's not that I want to withold this file from you, it's just that I
diff --git a/webcontroller/file_viewer.go b/webcontroller/file_viewer.go
index 44a80a1..e2ddfe9 100644
--- a/webcontroller/file_viewer.go
+++ b/webcontroller/file_viewer.go
@@ -16,14 +16,6 @@ import (
 	blackfriday "github.com/russross/blackfriday/v2"
 )
 
-func (wc *WebController) viewTokenOrBust() (t string) {
-	var err error
-	if t, err = wc.api.GetMiscViewToken(); err != nil && !wc.config.ProxyAPIRequests {
-		log.Error("Could not get viewtoken: %s", err)
-	}
-	return t
-}
-
 func browserCompat(ua string) bool {
 	return strings.Contains(ua, "MSIE") || strings.Contains(ua, "Trident/7.0")
 }
@@ -32,7 +24,6 @@ type fileViewerData struct {
 	Type           string       `json:"type"` // file or list
 	APIResponse    interface{}  `json:"api_response"`
 	CaptchaKey     string       `json:"captcha_key"`
-	ViewToken      string       `json:"view_token"`
 	Embedded       bool         `json:"embedded"`
 	UserAdsEnabled bool         `json:"user_ads_enabled"`
 	ThemeURI       template.URL `json:"theme_uri"`
@@ -101,7 +92,6 @@ func (wc *WebController) serveFileViewer(w http.ResponseWriter, r *http.Request,
 
 	var vd = fileViewerData{
 		CaptchaKey:     wc.captchaKey(),
-		ViewToken:      wc.viewTokenOrBust(),
 		UserAdsEnabled: templateData.User.Subscription.ID == "",
 	}
 
@@ -176,7 +166,6 @@ func (wc *WebController) serveListViewer(w http.ResponseWriter, r *http.Request,
 	var vd = fileViewerData{
 		Type:           "list",
 		CaptchaKey:     wc.captchaSiteKey,
-		ViewToken:      wc.viewTokenOrBust(),
 		UserAdsEnabled: templateData.User.Subscription.ID == "",
 		APIResponse:    list,
 	}