Merge pull request #503 from fkautz/pr_out_bandwidth_quota_now_supports_100_continue

pkg/api/api_router.go

@@ -92,8 +92,8 @@ func HTTPHandler(domain string, driver drivers.Driver) http.Handler {
 	}
 
 	h := validateHandler(conf, ignoreResourcesHandler(mux))
-	h = quota.BandwidthCap(h, 1*1024*1024*1024, time.Duration(30*time.Minute))
-	h = quota.BandwidthCap(h, 1024*1024*1024, time.Duration(24*time.Hour))
+	h = quota.BandwidthCap(h, 25*1024*1024, time.Duration(30*time.Minute))
+	h = quota.BandwidthCap(h, 100*1024*1024, time.Duration(24*time.Hour))
 	h = quota.RequestLimit(h, 100, time.Duration(30*time.Minute))
 	h = quota.RequestLimit(h, 1000, time.Duration(24*time.Hour))
 	h = quota.ConnectionLimit(h, 5)

pkg/api/quota/bandwidth_cap.go

@@ -31,16 +31,38 @@ type bandwidthQuotaHandler struct {
 	quotas  *quotaMap
 }
 
+var bandwidthQuotaExceeded = ErrorResponse{
+	Code:      "BandwithQuotaExceeded",
+	Message:   "Bandwidth Quota Exceeded",
+	Resource:  "",
+	RequestID: "",
+	HostID:    "",
+}
+
+var bandwidthInsufficientToProceed = ErrorResponse{
+	Code:      "BandwidthQuotaWillBeExceeded",
+	Message:   "Bandwidth quota will be exceeded with this request",
+	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{
+	if h.quotas.WillExceedQuota(longIP, req.ContentLength) {
+		writeError(w, req, bandwidthInsufficientToProceed, 429)
+		return
+	}
+	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,
@@ -65,10 +87,18 @@ 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 {
+		return 0, iodine.New(errors.New("Quota Met"), nil)
+	}
 	if q.quotas.IsQuotaMet(q.ip) {
+		q.err = true
+		writeError(q.w, q.req, bandwidthQuotaExceeded, 429)
 		return 0, iodine.New(errors.New("Quota Met"), nil)
 	}
 	n, err := q.ReadCloser.Read(b)
@@ -76,17 +106,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 +126,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

pkg/api/quota/quota_handler.go

@@ -17,8 +17,11 @@
 package quota
 
 import (
+	"bytes"
 	"encoding/binary"
+	"encoding/xml"
 	"net"
+	"net/http"
 	"sync"
 	"time"
 )
@@ -92,3 +95,24 @@ func (p longIP) IptoUint32() (result uint32) {
 	}
 	return binary.BigEndian.Uint32(ip)
 }
+
+// copied from api, no cyclic deps allowed
+
+// ErrorResponse - error response format
+type ErrorResponse struct {
+	XMLName   xml.Name `xml:"Error" json:"-"`
+	Code      string
+	Message   string
+	Resource  string
+	RequestID string
+	HostID    string
+}
+
+func writeError(w http.ResponseWriter, req *http.Request, errorResponse ErrorResponse, status int) {
+	var buf bytes.Buffer
+	encoder := xml.NewEncoder(&buf)
+	w.WriteHeader(status)
+	encoder.Encode(errorResponse)
+	encoder.Flush()
+	w.Write(buf.Bytes())
+}