From 17afd3f8c7a016d5103be949990efb695de865b5 Mon Sep 17 00:00:00 2001
From: Bohdan Horbeshko <bodqhrohro@gmail.com>
Date: Thu, 31 Mar 2022 21:42:12 -0400
Subject: Limit the file storage by an optional quota

---
 telegram/utils.go | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

(limited to 'telegram/utils.go')

diff --git a/telegram/utils.go b/telegram/utils.go
index f60f35f..b41423a 100644
--- a/telegram/utils.go
+++ b/telegram/utils.go
@@ -361,6 +361,9 @@ func (c *Client) formatFile(file *client.File, compact bool) string {
 		return ""
 	}
 
+	gateway.StorageLock.Lock()
+	defer gateway.StorageLock.Unlock()
+
 	var link string
 	var src string
 
@@ -372,6 +375,9 @@ func (c *Client) formatFile(file *client.File, compact bool) string {
 			return ""
 		}
 
+		size64:= uint64(file.Size)
+		c.prepareDiskSpace(size64)
+
 		basename := file.Remote.UniqueId + filepath.Ext(src)
 		dest := c.content.Path + "/" + basename // destination path
 		link = c.content.Link + "/" + basename // download link
@@ -387,6 +393,8 @@ func (c *Client) formatFile(file *client.File, compact bool) string {
 				return "<ERROR>"
 			}
 		}
+		gateway.CachedStorageSize += size64
+
 		// chown
 		if c.content.User != "" {
 			user, err := osUser.Lookup(c.content.User)
@@ -729,7 +737,12 @@ func (c *Client) messageToPrefix(message *client.Message, previewString string,
 }
 
 func (c *Client) ensureDownloadFile(file *client.File) *client.File {
+	gateway.StorageLock.Lock()
+	defer gateway.StorageLock.Unlock()
+
 	if file != nil {
+		c.prepareDiskSpace(uint64(file.Size))
+
 		newFile, err := c.DownloadFile(file.Id, 1, true)
 		if err == nil {
 			return newFile
@@ -952,3 +965,16 @@ func (c *Client) subscribeToID(id int64, chat *client.Chat) {
 		args...,
 	)
 }
+
+func (c *Client) prepareDiskSpace(size uint64) {
+	if gateway.StorageQuota > 0 && c.content.Path != "" {
+		var loweredQuota uint64
+		if gateway.StorageQuota >= size {
+			loweredQuota = gateway.StorageQuota - size
+		}
+		if gateway.CachedStorageSize >= loweredQuota {
+			log.Warn("Storage is rapidly clogged")
+			gateway.CleanOldFiles(c.content.Path, loweredQuota)
+		}
+	}
+}
-- 
cgit v1.2.3