Skip to content

Commit

Permalink
Added redigo based IAttackStore implementation (#66)
Browse files Browse the repository at this point in the history
* Added redigo based IAttackStore implementation

* Added redigo based IAttackStore implementation

* Fix lint issues
  • Loading branch information
nitishm committed Jul 7, 2019
2 parents 6aa268e + 43a1f2c commit 3b2862f
Show file tree
Hide file tree
Showing 4 changed files with 129 additions and 6 deletions.
27 changes: 21 additions & 6 deletions cmd/server/main.go
Expand Up @@ -11,7 +11,9 @@ import (
"vegeta-server/models"
"vegeta-server/pkg/vegeta"

gin "github.com/gin-gonic/gin"
"github.com/gomodule/redigo/redis"

"github.com/gin-gonic/gin"

log "github.com/sirupsen/logrus"

Expand All @@ -25,10 +27,11 @@ var (
)

var (
ip = kingpin.Flag("ip", "Server IP Address.").Default("0.0.0.0").String()
port = kingpin.Flag("port", "Server Port.").Default("80").String()
v = kingpin.Flag("version", "Version Info").Short('v').Bool()
debug = kingpin.Flag("debug", "Enabled Debug").Bool()
ip = kingpin.Flag("ip", "Server IP Address.").Default("0.0.0.0").String()
port = kingpin.Flag("port", "Server Port.").Default("80").String()
redisHost = kingpin.Flag("redis", "Redis Server Address.").String()
v = kingpin.Flag("version", "Version Info").Short('v').Bool()
debug = kingpin.Flag("debug", "Enabled Debug").Bool()
)

func main() {
Expand Down Expand Up @@ -56,7 +59,19 @@ func main() {
quit := make(chan struct{})
defer close(quit)

db := models.NewTaskMap()
var db models.IAttackStore

if redisHost != nil && *redisHost != "" {
db = models.NewRedis(func() redis.Conn {
conn, err := redis.Dial("tcp", *redisHost)
if err != nil {
log.Fatalf("Failed to connect to redis-server @ %s", *redisHost)
}
return conn
})
} else {
db = models.NewTaskMap()
}

d := dispatcher.NewDispatcher(
db,
Expand Down
1 change: 1 addition & 0 deletions go.mod
Expand Up @@ -8,6 +8,7 @@ require (
github.com/gin-contrib/sse v0.0.0-20190125020943-a7658810eb74 // indirect
github.com/gin-gonic/gin v1.3.0
github.com/golang/protobuf v1.2.0 // indirect
github.com/gomodule/redigo v2.0.0+incompatible
github.com/influxdata/tdigest v0.0.0-20181121200506-bf2b5ad3c0a9 // indirect
github.com/json-iterator/go v1.1.5 // indirect
github.com/kr/pretty v0.1.0 // indirect
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Expand Up @@ -14,6 +14,8 @@ github.com/gin-gonic/gin v1.3.0 h1:kCmZyPklC0gVdL728E6Aj20uYBJV93nj/TkwBTKhFbs=
github.com/gin-gonic/gin v1.3.0/go.mod h1:7cKuhb5qV2ggCFctp2fJQ+ErvciLZrIeoOSOm6mUr7Y=
github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM=
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
github.com/gomodule/redigo v2.0.0+incompatible h1:K/R+8tc58AaqLkqG2Ol3Qk+DR/TlNuhuh457pBFPtt0=
github.com/gomodule/redigo v2.0.0+incompatible/go.mod h1:B4C85qUVwatsJoIUNIfCRsp7qO0iAmpGFZ4EELWSbC4=
github.com/google/go-cmp v0.2.0 h1:+dTQ8DZQJz0Mb/HjFlkptS1FeQ4cWSnN941F8aEG4SQ=
github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
github.com/influxdata/tdigest v0.0.0-20181121200506-bf2b5ad3c0a9 h1:MHTrDWmQpHq/hkq+7cw9oYAt2PqUw52TZazRA0N7PGE=
Expand Down
105 changes: 105 additions & 0 deletions models/db.go
@@ -1,8 +1,11 @@
package models

import (
"encoding/json"
"fmt"
"sync"

"github.com/gomodule/redigo/redis"
)

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

var mu sync.RWMutex

// Redis stores all Attack/Report information in a redis database
type Redis struct {
connFn func() redis.Conn
}

func NewRedis(f func() redis.Conn) Redis {
return Redis{
f,
}
}

func (r Redis) Add(attack AttackDetails) error {
conn := r.connFn()
defer conn.Close()

v, err := json.Marshal(attack)
if err != nil {
return err
}
_, err = conn.Do("SET", attack.ID, v)
if err != nil {
return err
}
return nil
}

func (r Redis) GetAll(filterParams FilterParams) []AttackDetails {
attacks := make([]AttackDetails, 0)

filters := createFilterChain(filterParams)
conn := r.connFn()
defer conn.Close()

res, err := conn.Do("KEYS", "*")
if err != nil {
return nil
}
attackIDs, ok := res.([]interface{})
if !ok {
return nil
}

for _, attackID := range attackIDs {
var attack AttackDetails

res, err := conn.Do("GET", attackID)
if err != nil {
return nil
}

err = json.Unmarshal(res.([]byte), &attack)
if err != nil {
return nil
}
for _, filter := range filters {
if !filter(attack) {
goto skip
}
}
attacks = append(attacks, attack)
skip:
}

return attacks
}

func (r Redis) GetByID(id string) (AttackDetails, error) {
var attack AttackDetails
conn := r.connFn()
defer conn.Close()

res, err := conn.Do("GET", id)
if err != nil {
return attack, err
}

err = json.Unmarshal(res.([]byte), &attack)
if err != nil {
return attack, err
}

return attack, err
}

func (r Redis) Update(id string, attack AttackDetails) error {
if attack.ID != id {
return fmt.Errorf("update ID %s and attack ID %s do not match", id, attack.ID)
}
return r.Add(attack)
}

func (r Redis) Delete(id string) error {
conn := r.connFn()
defer conn.Close()

_, err := conn.Do("DEL", id)
if err != nil {
return err
}
return nil
}

// TaskMap is a map of attack ID's to their AttackDetails
type TaskMap map[string]AttackDetails

Expand Down

0 comments on commit 3b2862f

Please sign in to comment.