|
|
|
@ -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 = "aReader{ |
|
|
|
|
ReadCloser: req.Body, |
|
|
|
|
quotas: h.quotas, |
|
|
|
|
ip: longIP, |
|
|
|
|
w: w, |
|
|
|
|
req: req, |
|
|
|
|
} |
|
|
|
|
w = quotaWriter{ |
|
|
|
|
w = "aWriter{ |
|
|
|
|
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 |
|
|
|
|
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 |
|
|
|
|