From 7867ee89fa89bb57a81b13f9df315eab105a1908 Mon Sep 17 00:00:00 2001
From: "Frederick F. Kautz IV" <fkautz@alumni.cmu.edu>
Date: Sun, 26 Apr 2015 15:49:46 -0700
Subject: [PATCH] Bandwidth quotas now work on data returned from server

---
 pkg/api/api_router.go          |  3 ++-
 pkg/api/quota/bandwidth_cap.go | 29 ++++++++++++++++++++++++++---
 pkg/api/quota/quota_handler.go | 17 ++++++++++++-----
 3 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/pkg/api/api_router.go b/pkg/api/api_router.go
index 4a1a40108..bbc88b5fb 100644
--- a/pkg/api/api_router.go
+++ b/pkg/api/api_router.go
@@ -92,7 +92,8 @@ func HTTPHandler(domain string, driver drivers.Driver) http.Handler {
 	}
 
 	h := validateHandler(conf, ignoreResourcesHandler(mux))
-	h = quota.BandwidthCap(h, 250*1024*1024, time.Duration(30*time.Minute))
+	h = quota.BandwidthCap(h, 256*1024*1024, time.Duration(30*time.Minute))
+	h = quota.BandwidthCap(h, 1024*1024*1024, time.Duration(24*time.Hour))
 	h = quota.RequestLimit(h, 100, time.Duration(30*time.Minute))
 	return h
 }
diff --git a/pkg/api/quota/bandwidth_cap.go b/pkg/api/quota/bandwidth_cap.go
index 31f8e4073..46c88665d 100644
--- a/pkg/api/quota/bandwidth_cap.go
+++ b/pkg/api/quota/bandwidth_cap.go
@@ -18,6 +18,7 @@ package quota
 
 import (
 	"errors"
+	"github.com/minio-io/minio/pkg/iodine"
 	"io"
 	"net"
 	"net/http"
@@ -39,6 +40,11 @@ func (h *bandwidthQuotaHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 		quotas:     h.quotas,
 		ip:         longIP,
 	}
+	w = quotaWriter{
+		ResponseWriter: w,
+		quotas:         h.quotas,
+		ip:             longIP,
+	}
 	h.handler.ServeHTTP(w, req)
 }
 
@@ -63,15 +69,32 @@ type quotaReader struct {
 
 func (q quotaReader) Read(b []byte) (int, error) {
 	if q.quotas.IsQuotaMet(q.ip) {
-		return 0, errors.New("Quota Met")
+		return 0, iodine.New(errors.New("Quota Met"), nil)
 	}
 	n, err := q.ReadCloser.Read(b)
 	q.quotas.Add(q.ip, int64(n))
-	return n, err
+	return n, iodine.New(err, nil)
 }
 
 func (q quotaReader) Close() error {
-	return q.ReadCloser.Close()
+	return iodine.New(q.ReadCloser.Close(), nil)
+}
+
+type quotaWriter struct {
+	http.ResponseWriter
+	quotas *quotaMap
+	ip     uint32
+}
+
+func (q quotaWriter) Write(b []byte) (int, error) {
+	if q.quotas.IsQuotaMet(q.ip) {
+		return 0, iodine.New(errors.New("Quota Met"), nil)
+	}
+	q.quotas.Add(q.ip, int64(len(b)))
+	n, err := q.ResponseWriter.Write(b)
+	// remove from quota if a full write isn't performed
+	q.quotas.Add(q.ip, int64(n-len(b)))
+	return n, iodine.New(err, nil)
 }
 
 func segmentSize(duration time.Duration) time.Duration {
diff --git a/pkg/api/quota/quota_handler.go b/pkg/api/quota/quota_handler.go
index 48b3c425f..461d3d919 100644
--- a/pkg/api/quota/quota_handler.go
+++ b/pkg/api/quota/quota_handler.go
@@ -47,20 +47,27 @@ func (q *quotaMap) Add(ip uint32, size int64) {
 
 func (q *quotaMap) IsQuotaMet(ip uint32) bool {
 	q.clean()
+	if q.GetQuotaUsed(ip) >= q.limit {
+		return true
+	}
+	return false
+}
+
+func (q *quotaMap) GetQuotaUsed(ip uint32) (total int64) {
 	currentMinute := time.Now().UnixNano() / q.segmentSize.Nanoseconds()
 	if _, ok := q.data[currentMinute]; !ok {
 		q.data[currentMinute] = make(map[uint32]int64)
 	}
-	var total int64
 	for _, segment := range q.data {
 		if used, ok := segment[ip]; ok {
 			total += used
 		}
 	}
-	if total >= q.limit {
-		return true
-	}
-	return false
+	return
+}
+
+func (q *quotaMap) WillExceedQuota(ip uint32, size int64) (result bool) {
+	return q.GetQuotaUsed(ip)+size > q.limit
 }
 
 func (q *quotaMap) clean() {