Merge pull request #1 from martinlevesque/wip

Wip

.gitignore

@@ -1,2 +1,3 @@
 configs-dev
-traefik-dev.yml
+traefik-dev.yml
+plugins-local/

.traefik.yml

@@ -1,5 +1,5 @@
 # The name of your plugin as displayed in the Traefik Pilot web UI.
-displayName: Hello world
+displayName: Http Shaping
 
 # For now, `middleware` is the only type available.
 type: middleware
@@ -8,7 +8,7 @@ type: middleware
 import: github.com/martinlevesque/http_shaping
 
 # A brief description of what your plugin is doing.
-summary: Description of what my plugin is doing
+summary: Limit in/out traffic given a time interval.
 
 # Medias associated to the plugin (optional)
 iconPath: foo/icon.png
@@ -18,5 +18,7 @@ bannerPath: foo/banner.png
 # This is mandatory,
 # and Traefik Pilot will try to execute the plugin with the data you provide as part of its startup validity tests.
 testData:
-  Headers:
-    Foo: Bar
+  InTrafficLimit: 200MiB
+  OutTrafficLimit: 200MiB
+  LoopInterval: 60
+  ConsiderLimits: true

http_shaping.go

@@ -1,4 +1,4 @@
-// Package plugindemo a demo plugin.
+
 package http_shaping
 
 import (
@@ -9,58 +9,84 @@ import (
 	"text/template"
 	"time"
 	"strconv"
-	"github.com/martinlevesque/http_shaping/bytefmt"
+	"errors"
+	"strings"
+	"unicode"
 )
 
 // Config the plugin configuration.
 type Config struct {
-	Headers map[string]string `json:"headers,omitempty"`
+	LoopInterval int64
+	InTrafficLimit string
+	OutTrafficLimit string
+	ConsiderLimits bool
 }
 
 // CreateConfig creates the default plugin configuration.
 func CreateConfig() *Config {
 	return &Config{
-		Headers: make(map[string]string),
+		LoopInterval: 0,
 	}
 }
 
-// Demo a Demo plugin.
-type Demo struct {
+type HttpShaping struct {
 	next     http.Handler
-	headers  map[string]string
 	name     string
 	template *template.Template
 	loopBeginTimestamp int64
 	loopSumInBytes int64
 	loopSumOutBytes int64
+	loopInterval int64
+	inTrafficLimit uint64
+	outTrafficLimit uint64
+	considerLimits bool
 }
 
-// New created a new Demo plugin.
+// New created a new HttpShaping plugin.
 func New(ctx context.Context, next http.Handler, config *Config, name string) (http.Handler, error) {
 	fmt.Print("Initiating http_shaping plugin ", name, "\n")
-	what := bytefmt.ToBytes("10.5GiB")
-	fmt.Print("what ", what, "\n")
+	fmt.Print("Config: ", config, "\n")
+
+	if config.LoopInterval <= 0 {
+		return nil, fmt.Errorf("LoopInterval should be positive")
+	}
+
+	inTrafficLimit, err := ToBytes(config.InTrafficLimit)
+
+	if err != nil {
+		return nil, fmt.Errorf("invalid inTrafficLimit")
+	}
+
+	outTrafficLimit, err := ToBytes(config.OutTrafficLimit)
 
-	if len(config.Headers) == 0 {
-		return nil, fmt.Errorf("headers cannot be empty")
+	if err != nil {
+		return nil, fmt.Errorf("invalid outTrafficLimit")
 	}
 
-	return &Demo{
-		headers:  config.Headers,
+	return &HttpShaping{
 		next:     next,
 		name:     name,
-		template: template.New("demo").Delims("[[", "]]"),
+		template: template.New("httpShaping").Delims("[[", "]]"),
 		loopBeginTimestamp: time.Now().Unix(),
 		loopSumInBytes: 0,
 		loopSumOutBytes: 0,
+		loopInterval: config.LoopInterval,
+		inTrafficLimit: inTrafficLimit,
+		outTrafficLimit: outTrafficLimit,
+		considerLimits: config.ConsiderLimits,
 	}, nil
 }
 
-func (a *Demo) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
-	fmt.Print("Serving http_shaping plugin \n")
+func (a *HttpShaping) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
+	fmt.Print("Serving http_shaping plugin, host=", req.Host, "\n")
 
-	fmt.Print("loopSumInBytes: ", a.loopSumInBytes, "\n")
-	fmt.Print("loopSumOutBytes: ", a.loopSumOutBytes, "\n")
+	// reset the begin loop if needed
+	if time.Now().Unix() - a.loopBeginTimestamp >= a.loopInterval {
+		fmt.Print("Reseting the loop interval, host=", req.Host, "\n")
+		a.loopBeginTimestamp = time.Now().Unix()
+		a.loopSumOutBytes = 0
+		a.loopSumInBytes = 0
+	}
 
 	requestBytesStr := req.Header.Get("content-length")
 
@@ -72,8 +98,23 @@ func (a *Demo) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		}
 	}
 
-	a.loopBeginTimestamp = time.Now().Unix()
+	// check if we reached any of the limit
+	fmt.Print("Current In: ", a.loopSumInBytes, " < ", a.inTrafficLimit, " host=", req.Host, "\n")
+
+	if a.considerLimits && a.loopSumInBytes >= int64(a.inTrafficLimit) {
+		http.Error(rw, "Bandwidth limit reached", http.StatusTooManyRequests)
+		return
+	}
+
+	fmt.Print("Current Out: ", a.loopSumOutBytes, " < ", a.outTrafficLimit, " host=", req.Host, "\n")
+
+	if a.considerLimits && a.loopSumOutBytes >= int64(a.outTrafficLimit) {
+		http.Error(rw, "Bandwidth limit reached", http.StatusTooManyRequests)
+		return
+	}
+
 
+	// forwarding the request
 	a.next.ServeHTTP(rw, req)
 
 	responseBytesStr := rw.Header().Get("content-length")
@@ -86,3 +127,59 @@ func (a *Demo) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		}
 	}
 }
+
+
+const (
+	BYTE = 1 << (10 * iota)
+	KILOBYTE
+	MEGABYTE
+	GIGABYTE
+	TERABYTE
+	PETABYTE
+	EXABYTE
+)
+
+var invalidByteQuantityError = errors.New("byte quantity must be a positive integer with a unit of measurement like M, MB, MiB, G, GiB, or GB")
+
+// ToBytes parses a string formatted by ByteSize as bytes. Note binary-prefixed and SI prefixed units both mean a base-2 units
+// KB = K = KiB	= 1024
+// MB = M = MiB = 1024 * K
+// GB = G = GiB = 1024 * M
+// TB = T = TiB = 1024 * G
+// PB = P = PiB = 1024 * T
+// EB = E = EiB = 1024 * P
+func ToBytes(s string) (uint64, error) {
+	s = strings.TrimSpace(s)
+	s = strings.ToUpper(s)
+
+	i := strings.IndexFunc(s, unicode.IsLetter)
+
+	if i == -1 {
+		return 0, invalidByteQuantityError
+	}
+
+	bytesString, multiple := s[:i], s[i:]
+	bytes, err := strconv.ParseFloat(bytesString, 64)
+	if err != nil || bytes < 0 {
+		return 0, invalidByteQuantityError
+	}
+
+	switch multiple {
+	case "E", "EB", "EIB":
+		return uint64(bytes * EXABYTE), nil
+	case "P", "PB", "PIB":
+		return uint64(bytes * PETABYTE), nil
+	case "T", "TB", "TIB":
+		return uint64(bytes * TERABYTE), nil
+	case "G", "GB", "GIB":
+		return uint64(bytes * GIGABYTE), nil
+	case "M", "MB", "MIB":
+		return uint64(bytes * MEGABYTE), nil
+	case "K", "KB", "KIB":
+		return uint64(bytes * KILOBYTE), nil
+	case "B":
+		return uint64(bytes), nil
+	default:
+		return 0, invalidByteQuantityError
+	}
+}