-
Notifications
You must be signed in to change notification settings - Fork 16
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
6 changed files
with
288 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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") | ||
} | ||
``` |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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") | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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}}") | ||
} |