Skip to content

Commit

Permalink
Merge pull request #973 from Fenny/master
Browse files Browse the repository at this point in the history 
🧽 update limiter & filesystem
  • Loading branch information
@Fenny
Fenny authored Oct 27, 2020
2 parents 8b392e8 + 32fdbf0 commit 0d53f94
Show file tree
Hide file tree
Showing 9 changed files with 172 additions and 306 deletions.
3 changes: 1 addition & 2 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ module github.com/gofiber/fiber/v2
go 1.14

require (
github.com/philhofer/fwd v1.1.0
github.com/valyala/fasthttp v1.16.0
golang.org/x/sys v0.0.0-20201020230747-6e5568b54d1a
golang.org/x/sys v0.0.0-20201026173827-119d4633e4d1
)
6 changes: 2 additions & 4 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,6 @@ github.com/andybalholm/brotli v1.0.0 h1:7UCwP93aiSfvWpapti8g88vVVGp2qqtGyePsSuDa
github.com/andybalholm/brotli v1.0.0/go.mod h1:loMXtMfwqflxFJPmdbJO0a3KNoPuLBgiu3qAvBg8x/Y=
github.com/klauspost/compress v1.10.7 h1:7rix8v8GpI3ZBb0nSozFRgbtXKv+hOe+qfEpZqybrAg=
github.com/klauspost/compress v1.10.7/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
github.com/philhofer/fwd v1.1.0 h1:PAdZw9+/BCf4gc/kA2L/PbGPkFe72Kl2GLZXTG8HpU8=
github.com/philhofer/fwd v1.1.0/go.mod h1:gk3iGcWd9+svBvR0sR+KPcfE+RNWozjowpeBVG3ZVNU=
github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
github.com/valyala/fasthttp v1.16.0 h1:9zAqOYLl8Tuy3E5R6ckzGDJ1g8+pw15oQp2iL9Jl6gQ=
Expand All @@ -19,7 +17,7 @@ golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd h1:xhmwyvizuTgC2qz7ZlMluP20u
golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20200602225109-6fdc65e7d980 h1:OjiUf46hAmXblsZdnoSXsEUSKU8r1UEzcL5RVZ4gO9Y=
golang.org/x/sys v0.0.0-20200602225109-6fdc65e7d980/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20201020230747-6e5568b54d1a h1:e3IU37lwO4aq3uoRKINC7JikojFmE5gO7xhfxs8VC34=
golang.org/x/sys v0.0.0-20201020230747-6e5568b54d1a/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20201026173827-119d4633e4d1 h1:/DtoiOYKoQCcIFXQjz07RnWNPRCbqmSXSpgEzhC9ZHM=
golang.org/x/sys v0.0.0-20201026173827-119d4633e4d1/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
26 changes: 17 additions & 9 deletions middleware/filesystem/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,10 @@ app.Use(filesystem.New(filesystem.Config{
// Or extend your config for customization
app.Use(filesystem.New(filesystem.Config{
Root: http.Dir("./assets"),
Index: "index.html",
Browse: true,
NotFoundFile: "404.html"
Index: "index.html",
NotFoundFile: "404.html",
MaxAge: 3600,
}))
```

Expand Down Expand Up @@ -179,22 +180,28 @@ type Config struct {
// to a collection of files and directories.
//
// Required. Default: nil
Root http.FileSystem
Root http.FileSystem `json:"-"`

// Enable directory browsing.
//
// Optional. Default: false
Browse bool `json:"browse"`

// Index file for serving a directory.
//
// Optional. Default: "index.html"
Index string
Index string `json:"index"`

// Enable directory browsing.
// The value for the Cache-Control HTTP-header
// that is set on the file response. MaxAge is defined in seconds.
//
// Optional. Default: false
Browse bool
// Optional. Default value 0.
MaxAge int `json:"max_age"`

// File to return if path is not found. Useful for SPA's.
//
// Optional. Default: ""
NotFoundFile string
NotFoundFile string `json:"not_found_file"`
}
```

Expand All @@ -203,7 +210,8 @@ type Config struct {
var ConfigDefault = Config{
Next: nil,
Root: nil,
Index: "/index.html",
Browse: false,
Index: "/index.html",
MaxAge: 0,
}
```
26 changes: 19 additions & 7 deletions middleware/filesystem/filesystem.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package filesystem
import (
"net/http"
"os"
"strconv"
"strings"
"sync"

Expand All @@ -20,30 +21,37 @@ type Config struct {
// to a collection of files and directories.
//
// Required. Default: nil
Root http.FileSystem
Root http.FileSystem `json:"-"`

// Enable directory browsing.
//
// Optional. Default: false
Browse bool `json:"browse"`

// Index file for serving a directory.
//
// Optional. Default: "index.html"
Index string
Index string `json:"index"`

// Enable directory browsing.
// The value for the Cache-Control HTTP-header
// that is set on the file response. MaxAge is defined in seconds.
//
// Optional. Default: false
Browse bool
// Optional. Default value 0.
MaxAge int `json:"max_age"`

// File to return if path is not found. Useful for SPA's.
//
// Optional. Default: ""
NotFoundFile string
NotFoundFile string `json:"not_found_file"`
}

// ConfigDefault is the default config
var ConfigDefault = Config{
Next: nil,
Root: nil,
Index: "/index.html",
Browse: false,
Index: "/index.html",
MaxAge: 0,
}

// New creates a new middleware handler
Expand Down Expand Up @@ -73,6 +81,7 @@ func New(config ...Config) fiber.Handler {

var once sync.Once
var prefix string
var cacheControlStr = "public, max-age=" + strconv.Itoa(cfg.MaxAge)

// Return new handler
return func(c *fiber.Ctx) (err error) {
Expand Down Expand Up @@ -153,6 +162,9 @@ func New(config ...Config) fiber.Handler {
}

if method == fiber.MethodGet {
if cfg.MaxAge > 0 {
c.Set(fiber.HeaderCacheControl, cacheControlStr)
}
c.Response().SetBodyStream(file, contentLength)
return nil
}
Expand Down
133 changes: 69 additions & 64 deletions middleware/limiter/limiter.go
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package limiter

import (
"fmt"
"strconv"
"sync"
"sync/atomic"
Expand All @@ -24,10 +25,13 @@ type Config struct {
// Default: 5
Max int

// Duration is the time on how long to keep records of requests in memory
// DEPRECATED: Use Expiration instead
Duration time.Duration

// Expiration is the time on how long to keep records of requests in memory
//
// Default: 1 * time.Minute
Duration time.Duration
Expiration time.Duration

// Key allows you to generate custom keys, by default c.IP() is used
//
Expand All @@ -50,26 +54,21 @@ type Config struct {

// Internally used - if true, the simpler method of two maps is used in order to keep
// execution time down.
usingCustomStore bool
defaultStore bool
}

// ConfigDefault is the default config
var ConfigDefault = Config{
Next: nil,
Max: 5,
Duration: 1 * time.Minute,
Next: nil,
Max: 5,
Expiration: 1 * time.Minute,
Key: func(c *fiber.Ctx) string {
return c.IP()
},
LimitReached: func(c *fiber.Ctx) error {
return c.SendStatus(fiber.StatusTooManyRequests)
},
}

// trackedSession is the type used for session tracking
type trackedSession struct {
Hits int
ResetTime uint64
defaultStore: true,
}

// X-RateLimit-* headers
Expand All @@ -95,28 +94,36 @@ func New(config ...Config) fiber.Handler {
if cfg.Max <= 0 {
cfg.Max = ConfigDefault.Max
}
if int(cfg.Duration.Seconds()) <= 0 {
cfg.Duration = ConfigDefault.Duration
if int(cfg.Duration.Seconds()) <= 0 && int(cfg.Expiration.Seconds()) <= 0 {
cfg.Expiration = ConfigDefault.Expiration
}
if int(cfg.Duration.Seconds()) > 0 {
fmt.Println("[LIMITER] Duration is deprecated, please use Expiration")
if cfg.Expiration != ConfigDefault.Expiration {
cfg.Expiration = cfg.Duration
}
}
if cfg.Key == nil {
cfg.Key = ConfigDefault.Key
}
if cfg.LimitReached == nil {
cfg.LimitReached = ConfigDefault.LimitReached
}
if cfg.Store != nil {
cfg.usingCustomStore = true
if cfg.Store == nil {
cfg.defaultStore = true
}
}

// Limiter settings
var max = strconv.Itoa(cfg.Max)
var sessions = make(map[string]trackedSession)
var timestamp = uint64(time.Now().Unix())
var duration = uint64(cfg.Duration.Seconds())
var (
// Limiter settings
max = strconv.Itoa(cfg.Max)
timestamp = uint64(time.Now().Unix())
expiration = uint64(cfg.Expiration.Seconds())

// mutex for parallel read and write access
mux := &sync.Mutex{}
// Default store logic (if no Store is provided)
data = make(map[string]Entry)
mux = &sync.RWMutex{}
)

// Update timestamp every second
go func() {
Expand All @@ -136,80 +143,71 @@ func New(config ...Config) fiber.Handler {
// Get key (default is the remote IP)
key := cfg.Key(c)

// Lock mux (prevents values changing between retrieval and reassignment, which can and does
// break things)
mux.Lock()
// Create new entry
entry := Entry{}

var session trackedSession
// Lock entry
mux.Lock()

if cfg.usingCustomStore {
// Check if we need to use the default in-memory storage
if cfg.defaultStore {
entry = data[key]
} else {
// Load data from store
fromStore, err := cfg.Store.Get(key)
storeEntry, err := cfg.Store.Get(key)
if err != nil {
return err
}

if len(fromStore) == 0 {
// Assume this means item not found.
session = trackedSession{}
} else {
// Only decode if we found an entry
if len(storeEntry) > 0 {
// Decode bytes using msgp
_, err := session.UnmarshalMsg(fromStore)
if err != nil {
if _, err := entry.UnmarshalMsg(storeEntry); err != nil {
return err
}
}
} else {
// Load data from in-memory map
session = sessions[key]
}

// Set unix timestamp if not exist
ts := atomic.LoadUint64(&timestamp)
if session.ResetTime == 0 {
session.ResetTime = ts + duration
} else if ts >= session.ResetTime {
session.Hits = 0
session.ResetTime = ts + duration
if entry.Exp == 0 {
entry.Exp = ts + expiration
} else if ts >= entry.Exp {
entry.Hits = 0
entry.Exp = ts + expiration
}

// Increment key hits
session.Hits++

if cfg.usingCustomStore {
// Convert session struct into bytes

data, err := session.MarshalMsg(nil)
// Increment hits
entry.Hits++

// Check if we need to use the default in-memory storage
if cfg.defaultStore {
data[key] = entry
} else {
// Encode Entry to bytes using msgp
data, err := entry.MarshalMsg(nil)
if err != nil {
return err
}

// Store those bytes
err = cfg.Store.Set(key, data, cfg.Duration)
if err != nil {
// Pass bytes to Storage
if err = cfg.Store.Set(key, data, cfg.Expiration); err != nil {
return err
}
} else {
sessions[key] = session
}

// Get current hits
hitCount := session.Hits
mux.Unlock()

// Calculate when it resets in seconds
resetTime := session.ResetTime - ts
expire := entry.Exp - ts

// Set how many hits we have left
remaining := cfg.Max - hitCount

mux.Unlock()
remaining := cfg.Max - entry.Hits

// Check if hits exceed the cfg.Max
if remaining < 0 {
// Return response with Retry-After header
// https://tools.ietf.org/html/rfc6584
c.Set(fiber.HeaderRetryAfter, strconv.FormatUint(resetTime, 10))
c.Set(fiber.HeaderRetryAfter, strconv.FormatUint(expire, 10))

// Call LimitReached handler
return cfg.LimitReached(c)
Expand All @@ -218,9 +216,16 @@ func New(config ...Config) fiber.Handler {
// We can continue, update RateLimit headers
c.Set(xRateLimitLimit, max)
c.Set(xRateLimitRemaining, strconv.Itoa(remaining))
c.Set(xRateLimitReset, strconv.FormatUint(resetTime, 10))
c.Set(xRateLimitReset, strconv.FormatUint(expire, 10))

// Continue stack
return c.Next()
}
}

// replacer for strconv.FormatUint
// func appendInt(buf *bytebufferpool.ByteBuffer, v int) (int, error) {
// old := len(buf.B)
// buf.B = fasthttp.AppendUint(buf.B, v)
// return len(buf.B) - old, nil
// }
Loading

0 comments on commit 0d53f94

Please sign in to comment.