diff --git a/api/server/server.go b/api/server/server.go index cd7cf5978..fdb38168a 100644 --- a/api/server/server.go +++ b/api/server/server.go @@ -50,7 +50,7 @@ type Server struct { } // 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)) if err != nil { 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 diff --git a/api/server/server_options.go b/api/server/server_options.go index b3ed708a2..159636a4b 100644 --- a/api/server/server_options.go +++ b/api/server/server_options.go @@ -1,6 +1,12 @@ package server -import "context" +import ( + "context" + "fmt" + "net/http" + + "github.com/gin-gonic/gin" +) type ServerOption func(*Server) @@ -9,3 +15,38 @@ func EnableShutdownEndpoint(halt context.CancelFunc) ServerOption { 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) +}