From 769caeee2cbae9e170031a467f52a38518182044 Mon Sep 17 00:00:00 2001
From: ShiKhu <alex.kaminsky.11@gmail.com>
Date: Thu, 6 May 2021 21:44:56 +0200
Subject: [PATCH] feat (backend-assets): ASSETS_SIZE_LIMIT env var

---
 backend/Dockerfile                       |  1 +
 backend/pkg/env/vars.go                  |  4 ++++
 backend/services/assets/cacher/cacher.go | 13 +++++++------
 backend/services/assets/main.go          |  1 +
 4 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/backend/Dockerfile b/backend/Dockerfile
index 085e15544..af8d4ff71 100644
--- a/backend/Dockerfile
+++ b/backend/Dockerfile
@@ -39,6 +39,7 @@ ENV TZ=UTC \
     AWS_REGION_WEB=eu-central-1 \
     AWS_REGION_IOS=eu-west-1 \
     AWS_REGION_ASSETS=eu-central-1 \
+    ASSETS_SIZE_LIMIT=6291456
 
 
 ARG SERVICE_NAME
diff --git a/backend/pkg/env/vars.go b/backend/pkg/env/vars.go
index e6541d002..64e586cc3 100644
--- a/backend/pkg/env/vars.go
+++ b/backend/pkg/env/vars.go
@@ -36,6 +36,10 @@ func Uint64(key string) uint64 {
 	return n
 }
 
+func Int(key string) int {
+	return int(Uint64(key))
+}
+
 func Bool(key string) bool {
 	v := String(key)
 	if v != "true" && v != "false" {
diff --git a/backend/services/assets/cacher/cacher.go b/backend/services/assets/cacher/cacher.go
index d9f39aa77..85a8a9d61 100644
--- a/backend/services/assets/cacher/cacher.go
+++ b/backend/services/assets/cacher/cacher.go
@@ -17,7 +17,6 @@ import (
 	"openreplay/backend/pkg/storage"
 )
 
-const BODY_LIMIT = 6 * (1 << 20) // 6 Mb
 const MAX_CACHE_DEPTH = 5
 
 type cacher struct {
@@ -26,9 +25,10 @@ type cacher struct {
 	httpClient *http.Client 	// Docs: "Clients are safe for concurrent use by multiple goroutines."
 	rewriter *assets.Rewriter // Read only
 	Errors chan error
+	sizeLimit int
 }
 
-func NewCacher(region string, bucket string, origin string) *cacher {
+func NewCacher(region string, bucket string, origin string, sizeLimit int) *cacher {
 	rewriter := assets.NewRewriter(origin)
 	return &cacher{
 		timeoutMap: newTimeoutMap(),
@@ -39,8 +39,9 @@ func NewCacher(region string, bucket string, origin string) *cacher {
         TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
       },
 		},
-		rewriter: rewriter,
-		Errors: make(chan error),
+		rewriter:   rewriter,
+		Errors:     make(chan error),
+		sizeLimit:  sizeLimit,
 	}
 }
 
@@ -72,12 +73,12 @@ func (c *cacher) cacheURL(requestURL string, sessionID uint64, depth byte, conte
 		c.Errors <- errors.Wrap(fmt.Errorf("Status code is %v, ", res.StatusCode), context)
 		return
 	}
-	data, err := ioutil.ReadAll(io.LimitReader(res.Body, BODY_LIMIT+1))
+	data, err := ioutil.ReadAll(io.LimitReader(res.Body, int64(c.sizeLimit+1)))
 	if err != nil {
 		c.Errors <- errors.Wrap(err, context)
 		return
 	}
-	if len(data) > BODY_LIMIT {
+	if len(data) > c.sizeLimit {
 		c.Errors <- errors.Wrap(errors.New("Maximum size exceeded"), context)
 		return
 	}
diff --git a/backend/services/assets/main.go b/backend/services/assets/main.go
index e62c277de..0193be7bb 100644
--- a/backend/services/assets/main.go
+++ b/backend/services/assets/main.go
@@ -26,6 +26,7 @@ func main() {
 		env.String("AWS_REGION"),
 		env.String("S3_BUCKET_ASSETS"),
 		env.String("ASSETS_ORIGIN"),
+		env.Int("ASSETS_SIZE_LIMIT"),
 	)
 
 	consumer := queue.NewMessageConsumer(