diff --git a/main.go b/main.go index 49f4dee..8b500e4 100644 --- a/main.go +++ b/main.go @@ -1,7 +1,32 @@ package main -import "fmt" +import ( + "github.com/labstack/echo/v4" + "github.com/labstack/echo/v4/middleware" +) func main() { - fmt.Println("Hello, Bounty Hunter!") + e := echo.New() + + // Configure the BodyCache middleware + bodyCacheConfig := middleware.BodyCacheConfig{ + MaxBodySize: 1024 * 1024, // 1MB + } + + // Use the BodyCache middleware + e.Use(middleware.BodyCache(bodyCacheConfig)) + + // Other middlewares and routes + e.Use(middleware.Logger()) + e.Use(middleware.Recover()) + + e.POST("/bind", func(c echo.Context) error { + var data map[string]interface{} + if err := c.Bind(&data); err!= nil { + return err + } + return c.JSON(http.StatusOK, data) + }) + + e.Start(":8080") } diff --git a/middleware/body_cache.go b/middleware/body_cache.go new file mode 100644 index 0000000..9e1e82c --- /dev/null +++ b/middleware/body_cache.go @@ -0,0 +1,68 @@ +package middleware + +import ( + "bytes" + "io" + "net/http" + + "github.com/labstack/echo/v4" +) + +// BodyCacheConfig defines the config for BodyCache middleware. +type BodyCacheConfig struct { + // Skipper defines a function to skip middleware. + Skipper echo.Skipper + MaxBodySize int64 // Maximum size of the body to cache in bytes +} + +// DefaultBodyCacheConfig is the default config for BodyCache middleware. +var DefaultBodyCacheConfig = BodyCacheConfig{ + Skipper: echo.DefaultSkipper, + MaxBodySize: 1024 * 1024, // 1MB +} + +// BodyCache returns a middleware that caches the request body. +func BodyCache(config BodyCacheConfig) echo.MiddlewareFunc { + if config.Skipper == nil { + config.Skipper = DefaultBodyCacheConfig.Skipper + } + if config.MaxBodySize == 0 { + config.MaxBodySize = DefaultBodyCacheConfig.MaxBodySize + } + + return func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + if config.Skipper(c) { + return next(c) + } + + req := c.Request() + if req.Body == nil { + return next(c) + } + + var body []byte + buf := &bytes.Buffer{} + tee := io.TeeReader(req.Body, buf) + + // Read the body into the buffer + _, err := io.Copy(io.Discard, tee) + if err!= nil && err!= io.EOF { + return err + } + + body = buf.Bytes() + + // Check if the body size exceeds the maximum allowed size + if int64(len(body)) > config.MaxBodySize { + return echo.NewHTTPError(http.StatusRequestEntityTooLarge, "Request body too large") + } + + // Replace the request body with a new ReadCloser + req.Body = io.NopCloser(bytes.NewReader(body)) + + // Call the next handler + return next(c) + } + } +}