add server option to limit request size (#320)

we're going to want to do this in our service version of this thing, but
adding this here so that it's usable by everyone. just an option, can add it
to server configuration, but response is nicely formatted, etc.

closes #277
This commit is contained in:
Reed Allman
2017-09-18 22:34:19 -07:00
committed by Travis Reeder
parent c6f9b50afe
commit f335d34636
2 changed files with 44 additions and 3 deletions

View File

@@ -50,7 +50,7 @@ type Server struct {
} }
// NewFromEnv creates a new Functions server based on env vars. // NewFromEnv creates a new Functions server based on env vars.
func NewFromEnv(ctx context.Context) *Server { func NewFromEnv(ctx context.Context, opts ...ServerOption) *Server {
ds, err := datastore.New(viper.GetString(EnvDBURL)) ds, err := datastore.New(viper.GetString(EnvDBURL))
if err != nil { if err != nil {
logrus.WithError(err).Fatalln("Error initializing datastore.") logrus.WithError(err).Fatalln("Error initializing datastore.")
@@ -69,7 +69,7 @@ func NewFromEnv(ctx context.Context) *Server {
} }
} }
return New(ctx, ds, mq, logDB) return New(ctx, ds, mq, logDB, opts...)
} }
// New creates a new Functions server with the passed in datastore, message queue and API URL // New creates a new Functions server with the passed in datastore, message queue and API URL

View File

@@ -1,6 +1,12 @@
package server package server
import "context" import (
"context"
"fmt"
"net/http"
"github.com/gin-gonic/gin"
)
type ServerOption func(*Server) type ServerOption func(*Server)
@@ -9,3 +15,38 @@ func EnableShutdownEndpoint(halt context.CancelFunc) ServerOption {
s.Router.GET("/shutdown", s.handleShutdown(halt)) s.Router.GET("/shutdown", s.handleShutdown(halt))
} }
} }
func LimitRequestBody(max int64) ServerOption {
return func(s *Server) {
s.Router.Use(limitRequestBody(max))
}
}
func limitRequestBody(max int64) func(c *gin.Context) {
return func(c *gin.Context) {
cl := int64(c.Request.ContentLength)
if cl > max {
// try to deny this quickly, instead of just letting it get lopped off
handleErrorResponse(c, errTooBig{cl, max})
c.Abort()
return
}
// if no Content-Length specified, limit how many bytes we read and error
// if we hit the max (intercontinental anti-air missile defense system).
// read http.MaxBytesReader for gritty details..
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, max)
c.Next()
}
}
// models.APIError
type errTooBig struct {
n, max int64
}
func (e errTooBig) Code() int { return http.StatusRequestEntityTooLarge }
func (e errTooBig) Error() string {
return fmt.Sprintf("Content-Length too large for this server, %d > max %d", e.n, e.max)
}