Skip to content

Commit

Permalink
feat: initial
Browse files Browse the repository at this point in the history
  • Loading branch information
appleboy committed Mar 26, 2017
1 parent bf7de1d commit 3926445
Show file tree
Hide file tree
Showing 6 changed files with 288 additions and 0 deletions.
28 changes: 28 additions & 0 deletions .travis.yml
Original file line number Diff line number Diff line change
@@ -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
28 changes: 28 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -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")
}
```
21 changes: 21 additions & 0 deletions example/main.go
Original file line number Diff line number Diff line change
@@ -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")
}
90 changes: 90 additions & 0 deletions size.go
Original file line number Diff line number Diff line change
@@ -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()
}
}
98 changes: 98 additions & 0 deletions size_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
}
23 changes: 23 additions & 0 deletions test_server.tmpl
Original file line number Diff line number Diff line change
@@ -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}}")
}

0 comments on commit 3926445

Please sign in to comment.