Skip to content

Commit 3b2862f

Browse files
nitishmnitishm
authored
Added redigo based IAttackStore implementation (#66)
* Added redigo based IAttackStore implementation * Added redigo based IAttackStore implementation * Fix lint issues
2 parents 6aa268e + 43a1f2c commit 3b2862f

File tree

4 files changed

+129
-6
lines changed

4 files changed

+129
-6
lines changed

‎cmd/server/main.go‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,9 @@ import (
1111
"vegeta-server/models"
1212
"vegeta-server/pkg/vegeta"
1313

14-
gin "github.com/gin-gonic/gin"
14+
"github.com/gomodule/redigo/redis"
15+
16+
"github.com/gin-gonic/gin"
1517

1618
log "github.com/sirupsen/logrus"
1719

@@ -25,10 +27,11 @@ var (
2527
)
2628

2729
var (
28-
ip = kingpin.Flag("ip", "Server IP Address.").Default("0.0.0.0").String()
29-
port = kingpin.Flag("port", "Server Port.").Default("80").String()
30-
v = kingpin.Flag("version", "Version Info").Short('v').Bool()
31-
debug = kingpin.Flag("debug", "Enabled Debug").Bool()
30+
ip = kingpin.Flag("ip", "Server IP Address.").Default("0.0.0.0").String()
31+
port = kingpin.Flag("port", "Server Port.").Default("80").String()
32+
redisHost = kingpin.Flag("redis", "Redis Server Address.").String()
33+
v = kingpin.Flag("version", "Version Info").Short('v').Bool()
34+
debug = kingpin.Flag("debug", "Enabled Debug").Bool()
3235
)
3336

3437
func main() {
@@ -56,7 +59,19 @@ func main() {
5659
quit := make(chan struct{})
5760
defer close(quit)
5861

59-
db := models.NewTaskMap()
62+
var db models.IAttackStore
63+
64+
if redisHost != nil && *redisHost != "" {
65+
db = models.NewRedis(func() redis.Conn {
66+
conn, err := redis.Dial("tcp", *redisHost)
67+
if err != nil {
68+
log.Fatalf("Failed to connect to redis-server @ %s", *redisHost)
69+
}
70+
return conn
71+
})
72+
} else {
73+
db = models.NewTaskMap()
74+
}
6075

6176
d := dispatcher.NewDispatcher(
6277
db,

‎go.mod‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ require (
88
github.com/gin-contrib/sse v0.0.0-20190125020943-a7658810eb74 // indirect
99
github.com/gin-gonic/gin v1.3.0
1010
github.com/golang/protobuf v1.2.0 // indirect
11+
github.com/gomodule/redigo v2.0.0+incompatible
1112
github.com/influxdata/tdigest v0.0.0-20181121200506-bf2b5ad3c0a9 // indirect
1213
github.com/json-iterator/go v1.1.5 // indirect
1314
github.com/kr/pretty v0.1.0 // indirect

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ github.com/gin-gonic/gin v1.3.0 h1:kCmZyPklC0gVdL728E6Aj20uYBJV93nj/TkwBTKhFbs=
1414
github.com/gin-gonic/gin v1.3.0/go.mod h1:7cKuhb5qV2ggCFctp2fJQ+ErvciLZrIeoOSOm6mUr7Y=
1515
github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM=
1616
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
17+
github.com/gomodule/redigo v2.0.0+incompatible h1:K/R+8tc58AaqLkqG2Ol3Qk+DR/TlNuhuh457pBFPtt0=
18+
github.com/gomodule/redigo v2.0.0+incompatible/go.mod h1:B4C85qUVwatsJoIUNIfCRsp7qO0iAmpGFZ4EELWSbC4=
1719
github.com/google/go-cmp v0.2.0 h1:+dTQ8DZQJz0Mb/HjFlkptS1FeQ4cWSnN941F8aEG4SQ=
1820
github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
1921
github.com/influxdata/tdigest v0.0.0-20181121200506-bf2b5ad3c0a9 h1:MHTrDWmQpHq/hkq+7cw9oYAt2PqUw52TZazRA0N7PGE=

‎models/db.go‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
package models
22

33
import (
4+
"encoding/json"
45
"fmt"
56
"sync"
7+
8+
"github.com/gomodule/redigo/redis"
69
)
710

811
// IAttackStore captures all methods related to storing and retrieving attack details
@@ -26,6 +29,108 @@ type IAttackStore interface {
2629

2730
var mu sync.RWMutex
2831

32+
// Redis stores all Attack/Report information in a redis database
33+
type Redis struct {
34+
connFn func() redis.Conn
35+
}
36+
37+
func NewRedis(f func() redis.Conn) Redis {
38+
return Redis{
39+
f,
40+
}
41+
}
42+
43+
func (r Redis) Add(attack AttackDetails) error {
44+
conn := r.connFn()
45+
defer conn.Close()
46+
47+
v, err := json.Marshal(attack)
48+
if err != nil {
49+
return err
50+
}
51+
_, err = conn.Do("SET", attack.ID, v)
52+
if err != nil {
53+
return err
54+
}
55+
return nil
56+
}
57+
58+
func (r Redis) GetAll(filterParams FilterParams) []AttackDetails {
59+
attacks := make([]AttackDetails, 0)
60+
61+
filters := createFilterChain(filterParams)
62+
conn := r.connFn()
63+
defer conn.Close()
64+
65+
res, err := conn.Do("KEYS", "*")
66+
if err != nil {
67+
return nil
68+
}
69+
attackIDs, ok := res.([]interface{})
70+
if !ok {
71+
return nil
72+
}
73+
74+
for _, attackID := range attackIDs {
75+
var attack AttackDetails
76+
77+
res, err := conn.Do("GET", attackID)
78+
if err != nil {
79+
return nil
80+
}
81+
82+
err = json.Unmarshal(res.([]byte), &attack)
83+
if err != nil {
84+
return nil
85+
}
86+
for _, filter := range filters {
87+
if !filter(attack) {
88+
goto skip
89+
}
90+
}
91+
attacks = append(attacks, attack)
92+
skip:
93+
}
94+
95+
return attacks
96+
}
97+
98+
func (r Redis) GetByID(id string) (AttackDetails, error) {
99+
var attack AttackDetails
100+
conn := r.connFn()
101+
defer conn.Close()
102+
103+
res, err := conn.Do("GET", id)
104+
if err != nil {
105+
return attack, err
106+
}
107+
108+
err = json.Unmarshal(res.([]byte), &attack)
109+
if err != nil {
110+
return attack, err
111+
}
112+
113+
return attack, err
114+
}
115+
116+
func (r Redis) Update(id string, attack AttackDetails) error {
117+
if attack.ID != id {
118+
return fmt.Errorf("update ID %s and attack ID %s do not match", id, attack.ID)
119+
}
120+
return r.Add(attack)
121+
}
122+
123+
func (r Redis) Delete(id string) error {
124+
conn := r.connFn()
125+
defer conn.Close()
126+
127+
_, err := conn.Do("DEL", id)
128+
if err != nil {
129+
return err
130+
}
131+
return nil
132+
}
133+
29134
// TaskMap is a map of attack ID's to their AttackDetails
30135
type TaskMap map[string]AttackDetails
31136

0 commit comments

Comments
 (0)