Merge pull request #499 from fkautz/pr_out_bandwidth_quotas_now_work_on_data_returned_from_server

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
 }

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 {

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
-		return true
+}
-	}
+
-	return false
+func (q *quotaMap) WillExceedQuota(ip uint32, size int64) (result bool) {
+	return q.GetQuotaUsed(ip)+size > q.limit
 }
 
 func (q *quotaMap) clean() {