From 36e8a3ffc1a5eecca76e7e0045a1b639e695925a Mon Sep 17 00:00:00 2001 From: nmalhotra Date: Sat, 6 Jul 2019 21:37:09 -0400 Subject: [PATCH 1/3] Added redigo based IAttackStore implementation --- cmd/server/main.go | 27 ++++++++++--- go.mod | 1 + go.sum | 2 + models/db.go | 95 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 119 insertions(+), 6 deletions(-) diff --git a/cmd/server/main.go b/cmd/server/main.go index fddbe07..61d8b82 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -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" @@ -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() { @@ -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, diff --git a/go.mod b/go.mod index ccdddd4..6cbebf6 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 08279ba..51838fa 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/models/db.go b/models/db.go index 3f7c658..56fc855 100644 --- a/models/db.go +++ b/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 @@ -26,6 +29,98 @@ 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() + 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 { + var attacks []AttackDetails + filters := createFilterChain(filterParams) + conn := r.connFn() + res, err := conn.Do("KEYS", "*") + if err != nil { + return nil + } + if attackIDs, ok := res.([]interface{}); !ok { + return nil + } else { + for _, attackID := range attackIDs { + var attack AttackDetails + + res, err := conn.Do("GET", attackID) + if err != nil { + return nil + } + + if err := json.Unmarshal(res.([]byte), &attack); err != nil { + return nil + } else { + 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() + 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() + _, 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 From ccb9bb1f137dcdd1a1b2cc3c707714860bb0c7f8 Mon Sep 17 00:00:00 2001 From: nmalhotra Date: Sat, 6 Jul 2019 21:37:09 -0400 Subject: [PATCH 2/3] Added redigo based IAttackStore implementation --- cmd/server/main.go | 27 +++++++++--- go.mod | 1 + go.sum | 2 + models/db.go | 104 +++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 128 insertions(+), 6 deletions(-) diff --git a/cmd/server/main.go b/cmd/server/main.go index fddbe07..61d8b82 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -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" @@ -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() { @@ -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, diff --git a/go.mod b/go.mod index ccdddd4..6cbebf6 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 08279ba..51838fa 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/models/db.go b/models/db.go index 3f7c658..9d009dd 100644 --- a/models/db.go +++ b/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 @@ -26,6 +29,107 @@ 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 { + var attacks []AttackDetails + 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 From 43a1f2c29abb55950018044f1ec957401ba53a96 Mon Sep 17 00:00:00 2001 From: nmalhotra Date: Sat, 6 Jul 2019 21:53:02 -0400 Subject: [PATCH 3/3] Fix lint issues --- models/db.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/models/db.go b/models/db.go index 9d009dd..8c77610 100644 --- a/models/db.go +++ b/models/db.go @@ -56,7 +56,8 @@ func (r Redis) Add(attack AttackDetails) error { } func (r Redis) GetAll(filterParams FilterParams) []AttackDetails { - var attacks []AttackDetails + attacks := make([]AttackDetails, 0) + filters := createFilterChain(filterParams) conn := r.connFn() defer conn.Close()