finish me

pkg/api/quota/bandwidth_cap.go

@@ -20,6 +20,7 @@ import (
 	"errors"
 	"github.com/minio-io/minio/pkg/iodine"
 	"io"
+	"log"
 	"net"
 	"net/http"
 	"time"
@@ -31,20 +32,31 @@ type bandwidthQuotaHandler struct {
 	quotas  *quotaMap
 }
 
+var bandwidthQuotaExceeded = ErrorResponse{
+	Code:      "BandwithQuotaExceeded",
+	Message:   "Bandwidth Quota Exceeded",
+	Resource:  "",
+	RequestID: "",
+	HostID:    "",
+}
+
 // ServeHTTP is an http.Handler ServeHTTP method
 func (h *bandwidthQuotaHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	host, _, _ := net.SplitHostPort(req.RemoteAddr)
 	longIP := longIP{net.ParseIP(host)}.IptoUint32()
-	req.Body = quotaReader{
+	req.Body = &quotaReader{
 		ReadCloser: req.Body,
 		quotas:     h.quotas,
 		ip:         longIP,
+		w:          w,
+		req:        req,
 	}
-	w = quotaWriter{
+	w = &quotaWriter{
 		ResponseWriter: w,
 		quotas:         h.quotas,
 		ip:             longIP,
 	}
+	log.Println("serving")
 	h.handler.ServeHTTP(w, req)
 }
 
@@ -65,10 +77,20 @@ type quotaReader struct {
 	io.ReadCloser
 	quotas *quotaMap
 	ip     uint32
+	w      http.ResponseWriter
+	req    *http.Request
+	err    bool
 }
 
-func (q quotaReader) Read(b []byte) (int, error) {
+func (q *quotaReader) Read(b []byte) (int, error) {
+	if q.err {
+		log.Println("Shortcut, quitting")
+		return 0, iodine.New(errors.New("Quota Met"), nil)
+	}
 	if q.quotas.IsQuotaMet(q.ip) {
+		q.err = true
+		log.Println("QUOTA!!!")
+		writeError(q.w, q.req, bandwidthQuotaExceeded, 429)
 		return 0, iodine.New(errors.New("Quota Met"), nil)
 	}
 	n, err := q.ReadCloser.Read(b)
@@ -76,17 +98,17 @@ func (q quotaReader) Read(b []byte) (int, error) {
 	return n, iodine.New(err, nil)
 }
 
-func (q quotaReader) Close() error {
+func (q *quotaReader) Close() error {
 	return iodine.New(q.ReadCloser.Close(), nil)
 }
 
 type quotaWriter struct {
-	http.ResponseWriter
-	quotas *quotaMap
-	ip     uint32
+	ResponseWriter http.ResponseWriter
+	quotas         *quotaMap
+	ip             uint32
 }
 
-func (q quotaWriter) Write(b []byte) (int, error) {
+func (q *quotaWriter) Write(b []byte) (int, error) {
 	if q.quotas.IsQuotaMet(q.ip) {
 		return 0, iodine.New(errors.New("Quota Met"), nil)
 	}
@@ -96,6 +118,16 @@ func (q quotaWriter) Write(b []byte) (int, error) {
 	q.quotas.Add(q.ip, int64(n-len(b)))
 	return n, iodine.New(err, nil)
 }
+func (q *quotaWriter) Header() http.Header {
+	return q.ResponseWriter.Header()
+}
+
+func (q *quotaWriter) WriteHeader(status int) {
+	if q.quotas.IsQuotaMet(q.ip) {
+		return
+	}
+	q.ResponseWriter.WriteHeader(status)
+}
 
 func segmentSize(duration time.Duration) time.Duration {
 	var segmentSize time.Duration