From d360d0fec35c09cfaef410411a73b2dac928c3cc Mon Sep 17 00:00:00 2001
From: Anis Elleuch <vadmeste@gmail.com>
Date: Thu, 28 May 2015 19:33:35 +0100
Subject: [PATCH] Forbid the upload of files bigger than the memory backend
 capacity

---
 pkg/storage/drivers/errors.go              |  9 +++++++++
 pkg/storage/drivers/memory/memory.go       |  5 ++++-
 pkg/storage/drivers/memory/memory_cache.go | 19 ++++++++++++-------
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/pkg/storage/drivers/errors.go b/pkg/storage/drivers/errors.go
index 935596d13..442dd1be6 100644
--- a/pkg/storage/drivers/errors.go
+++ b/pkg/storage/drivers/errors.go
@@ -18,6 +18,10 @@ package drivers
 
 import "fmt"
 
+// InternalError - generic internal error
+type InternalError struct {
+}
+
 // BackendError - generic disk backend error
 type BackendError struct {
 	Path string
@@ -127,6 +131,11 @@ func EmbedError(bucket, object string, err error) ImplementationError {
 	}
 }
 
+// Return string an error formatted as the given text
+func (e InternalError) Error() string {
+	return "Internal error occured"
+}
+
 // Return string an error formatted as the given text
 func (e ObjectNotFound) Error() string {
 	return "Object not Found: " + e.Bucket + "#" + e.Object
diff --git a/pkg/storage/drivers/memory/memory.go b/pkg/storage/drivers/memory/memory.go
index 9757fe340..cab3a8058 100644
--- a/pkg/storage/drivers/memory/memory.go
+++ b/pkg/storage/drivers/memory/memory.go
@@ -271,8 +271,11 @@ func (memory *memoryDriver) createObject(bucket, key, contentType, expectedMD5Su
 	totalLength := len(readBytes)
 
 	memory.lock.Lock()
-	memory.objects.Set(objectKey, readBytes)
+	ok := memory.objects.Set(objectKey, readBytes)
 	memory.lock.Unlock()
+	if !ok {
+		return "", iodine.New(drivers.InternalError{}, nil)
+	}
 	// setting up for de-allocation
 	readBytes = nil
 
diff --git a/pkg/storage/drivers/memory/memory_cache.go b/pkg/storage/drivers/memory/memory_cache.go
index e72ae9c46..a543880fd 100644
--- a/pkg/storage/drivers/memory/memory_cache.go
+++ b/pkg/storage/drivers/memory/memory_cache.go
@@ -115,23 +115,28 @@ func (r *Cache) Get(key string) (interface{}, bool) {
 }
 
 // Set will persist a value to the cache
-func (r *Cache) Set(key string, value interface{}) {
+func (r *Cache) Set(key string, value interface{}) bool {
 	r.Lock()
-	// remove random key if only we reach the maxSize threshold,
-	// if not assume infinite memory
+	defer r.Unlock()
+	valueLen := uint64(len(value.([]byte)))
 	if r.maxSize > 0 {
+		// check if the size of the object is not bigger than the
+		// capacity of the cache
+		if valueLen > r.maxSize {
+			return false
+		}
+		// remove random key if only we reach the maxSize threshold
 		for key := range r.items {
-			for (r.currentSize + uint64(len(value.([]byte)))) > r.maxSize {
+			for (r.currentSize + valueLen) > r.maxSize {
 				r.Delete(key)
 			}
 			break
 		}
 	}
 	r.items[key] = value
-	r.currentSize += uint64(len(value.([]byte)))
+	r.currentSize += valueLen
 	r.updatedAt[key] = time.Now()
-	r.Unlock()
-	return
+	return true
 }
 
 // Expire expires keys which have expired