From db5870be289f3fd64beb869c0a67f303f0ac0299 Mon Sep 17 00:00:00 2001
From: "Frederick F. Kautz IV" <fkautz@alumni.cmu.edu>
Date: Mon, 27 Apr 2015 18:30:21 -0700
Subject: [PATCH] Fix multiple response.WriteHeader calls

---
 pkg/api/quota/bandwidth_cap.go       | 22 +++++++++++++++++++---
 pkg/storage/drivers/memory/memory.go |  4 +++-
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/pkg/api/quota/bandwidth_cap.go b/pkg/api/quota/bandwidth_cap.go
index 7b47a48cd..b4e8f3230 100644
--- a/pkg/api/quota/bandwidth_cap.go
+++ b/pkg/api/quota/bandwidth_cap.go
@@ -24,6 +24,8 @@ import (
 	"time"
 
 	"github.com/minio-io/minio/pkg/iodine"
+	"github.com/minio-io/minio/pkg/utils/log"
+	"sync"
 )
 
 // bandwidthQuotaHandler
@@ -40,17 +42,20 @@ func (h *bandwidthQuotaHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 		writeErrorResponse(w, req, BandWidthInsufficientToProceed, req.URL.Path)
 		return
 	}
-	req.Body = &quotaReader{
+	qr := &quotaReader{
 		ReadCloser: req.Body,
 		quotas:     h.quotas,
 		ip:         longIP,
 		w:          w,
 		req:        req,
+		lock:       &sync.RWMutex{},
 	}
+	req.Body = qr
 	w = &quotaWriter{
 		ResponseWriter: w,
 		quotas:         h.quotas,
 		ip:             longIP,
+		quotaReader:    qr,
 	}
 	h.handler.ServeHTTP(w, req)
 }
@@ -75,13 +80,19 @@ type quotaReader struct {
 	w      http.ResponseWriter
 	req    *http.Request
 	err    bool
+	lock   *sync.RWMutex
 }
 
 func (q *quotaReader) Read(b []byte) (int, error) {
+	log.Println(q.quotas.GetQuotaUsed(q.ip))
+	log.Println(q.quotas.limit)
+	q.lock.Lock()
+	defer q.lock.Unlock()
 	if q.err {
 		return 0, iodine.New(errors.New("Quota Met"), nil)
 	}
-	if q.quotas.IsQuotaMet(q.ip) {
+	if q.err == false && q.quotas.IsQuotaMet(q.ip) {
+		defer q.lock.Unlock()
 		q.err = true
 		writeErrorResponse(q.w, q.req, BandWidthQuotaExceeded, q.req.URL.Path)
 		return 0, iodine.New(errors.New("Quota Met"), nil)
@@ -99,9 +110,12 @@ type quotaWriter struct {
 	ResponseWriter http.ResponseWriter
 	quotas         *quotaMap
 	ip             uint32
+	quotaReader    *quotaReader
 }
 
 func (q *quotaWriter) Write(b []byte) (int, error) {
+	q.quotaReader.lock.RLock()
+	defer q.quotaReader.lock.RUnlock()
 	if q.quotas.IsQuotaMet(q.ip) {
 		return 0, iodine.New(errors.New("Quota Met"), nil)
 	}
@@ -116,7 +130,9 @@ func (q *quotaWriter) Header() http.Header {
 }
 
 func (q *quotaWriter) WriteHeader(status int) {
-	if q.quotas.IsQuotaMet(q.ip) {
+	q.quotaReader.lock.RLock()
+	defer q.quotaReader.lock.RUnlock()
+	if q.quotas.IsQuotaMet(q.ip) || q.quotaReader.err {
 		return
 	}
 	q.ResponseWriter.WriteHeader(status)
diff --git a/pkg/storage/drivers/memory/memory.go b/pkg/storage/drivers/memory/memory.go
index a7c2ca99b..97f854174 100644
--- a/pkg/storage/drivers/memory/memory.go
+++ b/pkg/storage/drivers/memory/memory.go
@@ -222,7 +222,9 @@ func (memory *memoryDriver) CreateObject(bucket, key, contentType, expectedMD5Su
 			summer.Write(chunk.Data)
 			_, err := io.Copy(&bytesBuffer, bytes.NewBuffer(chunk.Data))
 			if err != nil {
-				return iodine.New(err, nil)
+				err := iodine.New(err, nil)
+				log.Println(err)
+				return err
 			}
 			if uint64(totalLength) > memory.maxSize {
 				return iodine.New(drivers.EntityTooLarge{