Limit the file storage by an optional quota

1 files changed, 78 insertions, 0 deletions

diff --git a/xmpp/component.go b/xmpp/component.go
index 961761b..1749100 100644
--- a/xmpp/component.go
+++ b/xmpp/component.go
@@ -2,6 +2,8 @@ package xmpp
 
 import (
 	"github.com/pkg/errors"
+	"regexp"
+	"strconv"
 	"sync"
 	"time"
 
@@ -20,6 +22,19 @@ var sessions map[string]*telegram.Client
 var db *persistence.SessionsYamlDB
 var sessionLock sync.Mutex
 
+const (
+	B uint64 = 1
+	KB       = B << 10
+	MB       = KB << 10
+	GB       = MB << 10
+	TB       = GB << 10
+	PB       = TB << 10
+	EB       = PB << 10
+
+	maxUint64 uint64 = (1 << 64) - 1
+)
+var sizeRegex = regexp.MustCompile("\\A([0-9]+) ?([KMGTPE]?B?)\\z")
+
 // NewComponent starts a new component and wraps it in
 // a stream manager that you should start yourself
 func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.StreamManager, *xmpp.Component, error) {
@@ -32,6 +47,13 @@ func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.Strea
 
 	tgConf = tc
 
+	if tc.Content.Quota != "" {
+		gateway.StorageQuota, err = parseSize(tc.Content.Quota)
+		if err != nil {
+			log.Warnf("Error parsing the storage quota: %v; the cleaner is disabled", err)
+		}
+	}
+
 	options := xmpp.ComponentOptions{
 		TransportConfiguration: xmpp.TransportConfiguration{
 			Address: conf.Host + ":" + conf.Port,
@@ -80,10 +102,22 @@ func heartbeat(component *xmpp.Component) {
 	}
 	sessionLock.Unlock()
 
+	quotaLowThreshold := gateway.StorageQuota / 10 * 9
+
 	log.Info("Starting heartbeat queue")
 
 	// status updater thread
 	for {
+		gateway.StorageLock.Lock()
+		if quotaLowThreshold > 0 && tgConf.Content.Path != "" {
+			gateway.MeasureStorageSize(tgConf.Content.Path)
+
+			if gateway.CachedStorageSize > quotaLowThreshold {
+				gateway.CleanOldFiles(tgConf.Content.Path, quotaLowThreshold)
+			}
+		}
+		gateway.StorageLock.Unlock()
+
 		time.Sleep(60e9)
 		now := time.Now().Unix()
 
@@ -201,3 +235,47 @@ func Close(component *xmpp.Component) {
 	// close stream
 	component.Disconnect()
 }
+
+// based on https://github.com/c2h5oh/datasize/blob/master/datasize.go
+func parseSize(sSize string) (uint64, error) {
+	sizeParts := sizeRegex.FindStringSubmatch(sSize)
+
+	if len(sizeParts) > 2 {
+		numPart, err := strconv.ParseInt(sizeParts[1], 10, 64)
+		if err != nil {
+			return 0, err
+		}
+
+		var divisor uint64
+		val := uint64(numPart)
+
+		if len(sizeParts[2]) > 0 {
+			switch sizeParts[2][0] {
+			case 'B':
+				divisor = 1
+			case 'K':
+				divisor = KB
+			case 'M':
+				divisor = MB
+			case 'G':
+				divisor = GB
+			case 'T':
+				divisor = TB
+			case 'P':
+				divisor = PB
+			case 'E':
+				divisor = EB
+			}
+		}
+
+		if divisor == 0 {
+			return 0, &strconv.NumError{"Wrong suffix", sSize, strconv.ErrSyntax}
+		}
+		if val > maxUint64/divisor {
+			return 0, &strconv.NumError{"Overflow", sSize, strconv.ErrRange}
+		}
+		return val * divisor, nil
+	}
+
+	return 0, &strconv.NumError{"Not enough parts", sSize, strconv.ErrSyntax}
+}