feat: initial

Author: appleboy

.travis.yml

@@ -0,0 +1,28 @@
+language: go
+sudo: false
+
+go:
+  - 1.6.x
+  - 1.7.x
+  - 1.8.x
+  - tip
+
+install:
+  - go get -u github.com/kardianos/govendor
+  - go get github.com/campoy/embedmd
+  - govendor sync
+
+script:
+  - embedmd -d *.md
+  - go test -v -covermode=atomic -coverprofile=coverage.out
+
+after_success:
+  - bash <(curl -s https://codecov.io/bash)
+
+notifications:
+  webhooks:
+    urls:
+      - https://webhooks.gitter.im/e/acc2c57482e94b44f557
+    on_success: change
+    on_failure: always
+    on_start: false

README.md

@@ -1,2 +1,30 @@
 # size
+
 Limit size of POST requests for Gin framework
+
+## Example
+
+[embedmd]:# (example/main.go go)
+```go
+package main
+
+import (
+	"github.com/gin-contrib/size"
+	"github.com/gin-gonic/gin"
+)
+
+func handler(ctx *gin.Context) {
+	val := ctx.PostForm("b")
+	if len(ctx.Errors) > 0 {
+		return
+	}
+	ctx.String(http.StatusOK, "got %s\n", val)
+}
+
+func main() {
+	rtr := gin.Default()
+	rtr.Use(ratelimit.RateLimiter(10))
+	rtr.POST("/", handler)
+	rtr.Run(":8080")
+}
+```

example/main.go

@@ -0,0 +1,21 @@
+package main
+
+import (
+	"github.com/gin-contrib/size"
+	"github.com/gin-gonic/gin"
+)
+
+func handler(ctx *gin.Context) {
+	val := ctx.PostForm("b")
+	if len(ctx.Errors) > 0 {
+		return
+	}
+	ctx.String(http.StatusOK, "got %s\n", val)
+}
+
+func main() {
+	rtr := gin.Default()
+	rtr.Use(ratelimit.RateLimiter(10))
+	rtr.POST("/", handler)
+	rtr.Run(":8080")
+}

size.go

@@ -0,0 +1,90 @@
+package limits
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+)
+
+type maxBytesReader struct {
+	ctx        *gin.Context
+	rdr        io.ReadCloser
+	remaining  int64
+	wasAborted bool
+	sawEOF     bool
+}
+
+func (mbr *maxBytesReader) tooLarge() (n int, err error) {
+	n, err = 0, fmt.Errorf("HTTP request too large")
+
+	if !mbr.wasAborted {
+		mbr.wasAborted = true
+		ctx := mbr.ctx
+		ctx.Error(err)
+		ctx.Header("connection", "close")
+		ctx.String(http.StatusRequestEntityTooLarge, "request too large")
+		ctx.AbortWithStatus(http.StatusRequestEntityTooLarge)
+	}
+	return
+}
+
+func (mbr *maxBytesReader) Read(p []byte) (n int, err error) {
+	toRead := mbr.remaining
+	if mbr.remaining == 0 {
+		if mbr.sawEOF {
+			return mbr.tooLarge()
+		}
+		// The underlying io.Reader may not return (0, io.EOF)
+		// at EOF if the requested size is 0, so read 1 byte
+		// instead. The io.Reader docs are a bit ambiguous
+		// about the return value of Read when 0 bytes are
+		// requested, and {bytes,strings}.Reader gets it wrong
+		// too (it returns (0, nil) even at EOF).
+		toRead = 1
+	}
+	if int64(len(p)) > toRead {
+		p = p[:toRead]
+	}
+	n, err = mbr.rdr.Read(p)
+	if err == io.EOF {
+		mbr.sawEOF = true
+	}
+	if mbr.remaining == 0 {
+		// If we had zero bytes to read remaining (but hadn't seen EOF)
+		// and we get a byte here, that means we went over our limit.
+		if n > 0 {
+			return mbr.tooLarge()
+		}
+		return 0, err
+	}
+	mbr.remaining -= int64(n)
+	if mbr.remaining < 0 {
+		mbr.remaining = 0
+	}
+	return
+}
+
+func (mbr *maxBytesReader) Close() error {
+	return mbr.rdr.Close()
+}
+
+// RateLimiter returns a middleware that limits the size of request
+// When a request is over the limit, the following will happen:
+// * Error will be added to the context
+// * Connection: close header will be set
+// * Error 413 will be send to client (http.StatusRequestEntityTooLarge)
+// * Current context will be aborted
+func RateLimiter(limit int64) gin.HandlerFunc {
+	return func(ctx *gin.Context) {
+		ctx.Request.Body = &maxBytesReader{
+			ctx:        ctx,
+			rdr:        ctx.Request.Body,
+			remaining:  limit,
+			wasAborted: false,
+			sawEOF:     false,
+		}
+		ctx.Next()
+	}
+}

size_test.go

@@ -0,0 +1,98 @@
+package limits
+
+import (
+	"bytes"
+	"fmt"
+	"net/http"
+	"os"
+	"os/exec"
+	"testing"
+	"text/template"
+	"time"
+)
+
+var (
+	params = struct {
+		Size int
+		Port int
+	}{10, 9388}
+
+	codeFile  = "/tmp/ratelimit_test_server.go"
+	serverURL string
+)
+
+func init() {
+	tmpl := template.Must(template.ParseFiles("test_server.tmpl"))
+	fp, err := os.Create(codeFile)
+	if err != nil {
+		panic(fmt.Errorf("can't open %s - %s", codeFile, err))
+	}
+	err = tmpl.Execute(fp, params)
+	if err != nil {
+		panic(fmt.Errorf("can't create %s - %s", codeFile, err))
+	}
+	serverURL = fmt.Sprintf("http://localhost:%d", params.Port)
+}
+
+func waitForServer() error {
+	timeout := 30 * time.Second
+	ch := make(chan bool)
+	go func() {
+		for {
+			_, err := http.Post(serverURL, "text/plain", nil)
+			if err == nil {
+				ch <- true
+			}
+			time.Sleep(10 * time.Millisecond)
+		}
+	}()
+
+	select {
+	case <-ch:
+		return nil
+	case <-time.After(timeout):
+		return fmt.Errorf("server did not reply after %v", timeout)
+	}
+
+}
+
+func runServer() (*exec.Cmd, error) {
+	cmd := exec.Command("go", "run", codeFile)
+	cmd.Start()
+	if err := waitForServer(); err != nil {
+		return nil, err
+	}
+	return cmd, nil
+}
+
+func doPost(val string) (*http.Response, error) {
+	cmd, err := runServer()
+	if err != nil {
+		return nil, err
+	}
+	defer cmd.Process.Kill()
+
+	var buf bytes.Buffer
+	fmt.Fprintf(&buf, "big=%s", val)
+	return http.Post(serverURL, "application/x-www-form-urlencoded", &buf)
+}
+
+func TestRateLimiterOK(t *testing.T) {
+	resp, err := doPost("abc")
+	if err != nil {
+		t.Fatalf("error posting - %s", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("bad status - %d", resp.StatusCode)
+	}
+}
+
+func TestRateLimiterOver(t *testing.T) {
+	resp, err := doPost("abcdefghijklmnop")
+	if err != nil {
+		t.Fatalf("error posting - %s", err)
+	}
+	if resp.StatusCode != http.StatusRequestEntityTooLarge {
+		t.Fatalf("bad status - %d", resp.StatusCode)
+	}
+}

test_server.tmpl

@@ -0,0 +1,23 @@
+package main
+
+import (
+  "net/http"
+
+  "github.com/gin-contrib/limits"
+  "github.com/gin-gonic/gin"
+)
+
+func handler(ctx *gin.Context) {
+  val := ctx.PostForm("b")
+  if len(ctx.Errors) > 0 {
+    return
+  }
+  ctx.String(http.StatusOK, val)
+}
+
+func main() {
+  rtr := gin.Default()
+  rtr.Use(limits.RateLimiter({{.Size}}))
+  rtr.POST("/", handler)
+  rtr.Run(":{{.Port}}")
+}