forked from albertrdixon/romulus
/
connlimit.go
70 lines (59 loc) · 1.92 KB
/
connlimit.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
package connlimit
import (
"fmt"
"net/http"
"github.com/codegangsta/cli"
"github.com/mailgun/oxy/connlimit"
"github.com/mailgun/oxy/utils"
"github.com/timelinelabs/vulcand/plugin"
)
const Type = "connlimit"
func GetSpec() *plugin.MiddlewareSpec {
return &plugin.MiddlewareSpec{
Type: Type,
FromOther: FromOther,
FromCli: FromCli,
CliFlags: CliFlags(),
}
}
// Control simultaneous connections for a location per some variable.
type ConnLimit struct {
Connections int64
Variable string // Variable defines how the limiting should be done. e.g. 'client.ip' or 'request.header.X-My-Header'
}
// Returns vulcan library compatible middleware
func (c *ConnLimit) NewHandler(next http.Handler) (http.Handler, error) {
extract, err := utils.NewExtractor(c.Variable)
if err != nil {
return nil, err
}
return connlimit.New(next, extract, c.Connections)
}
func NewConnLimit(connections int64, variable string) (*ConnLimit, error) {
if _, err := utils.NewExtractor(variable); err != nil {
return nil, err
}
if connections < 0 {
return nil, fmt.Errorf("connections should be > 0, got %d", connections)
}
return &ConnLimit{
Connections: connections,
Variable: variable,
}, nil
}
func (c *ConnLimit) String() string {
return fmt.Sprintf("connections=%d, variable=%s", c.Connections, c.Variable)
}
func FromOther(c ConnLimit) (plugin.Middleware, error) {
return NewConnLimit(c.Connections, c.Variable)
}
// Constructs the middleware from the command line
func FromCli(c *cli.Context) (plugin.Middleware, error) {
return NewConnLimit(int64(c.Int("connections")), c.String("var"))
}
func CliFlags() []cli.Flag {
return []cli.Flag{
cli.StringFlag{Name: "variable, var", Value: "client.ip", Usage: "variable to rate against, e.g. client.ip, request.host or request.header.X-Header"},
cli.IntFlag{Name: "connections, conns", Value: 1, Usage: "amount of simultaneous connections allowed per variable value"},
}
}