add redis timeout option (#67)

* add redis timeout option * delete golint from make
vulsio · Oct 28, 2021 · 1ebf9c4 · 1ebf9c4
1 parent f8c6979
commit 1ebf9c4
Show file tree

Hide file tree

Showing 11 changed files with 27 additions and 19 deletions.
diff --git a/.gitignore b/.gitignore
@@ -8,5 +8,6 @@ vendor/
 *.sqlite3
 *.sqlite3-shm
 *.sqlite3-wal
+*.sqlite3-journal
 go-exploitdb
-.vscode
+.vscode
diff --git a/GNUmakefile b/GNUmakefile
@@ -3,7 +3,6 @@
 	install \
 	all \
 	vendor \
-	lint \
 	vet \
 	fmt \
 	fmtcheck \
@@ -39,10 +38,6 @@ install: main.go
 
 all: test
 
-lint:
-	@ go get -v golang.org/x/lint/golint
-	$(foreach file,$(SRCS),golint $(file) || exit;)
-
 vet:
 	$(foreach pkg,$(PKGS),go vet $(pkg);)
 
@@ -52,7 +47,7 @@ fmt:
 fmtcheck:
 	$(foreach file,$(SRCS),gofmt -d $(file);)
 
-pretest: lint vet fmtcheck
+pretest: vet fmtcheck
 
 test: pretest
 	$(foreach pkg,$(PKGS),go test -v $(pkg) || exit;)

diff --git a/commands/fetch-awesomepoc.go b/commands/fetch-awesomepoc.go
@@ -31,6 +31,7 @@ func fetchAwesomePoc(cmd *cobra.Command, args []string) (err error) {
 		viper.GetString("dbtype"),
 		viper.GetString("dbpath"),
 		viper.GetBool("debug-sql"),
+		db.Option{},
 	)
 	if err != nil {
 		if locked {

diff --git a/commands/fetch-exploitdb.go b/commands/fetch-exploitdb.go
@@ -31,6 +31,7 @@ func fetchExploitDB(cmd *cobra.Command, args []string) (err error) {
 		viper.GetString("dbtype"),
 		viper.GetString("dbpath"),
 		viper.GetBool("debug-sql"),
+		db.Option{},
 	)
 	if err != nil {
 		if locked {

diff --git a/commands/fetch-githubrepos.go b/commands/fetch-githubrepos.go
@@ -37,6 +37,7 @@ func fetchGitHubRepos(cmd *cobra.Command, args []string) (err error) {
 		viper.GetString("dbtype"),
 		viper.GetString("dbpath"),
 		viper.GetBool("debug-sql"),
+		db.Option{},
 	)
 	if err != nil {
 		if locked {

diff --git a/commands/search.go b/commands/search.go
@@ -46,6 +46,7 @@ func searchExploit(cmd *cobra.Command, args []string) error {
 		viper.GetString("dbtype"),
 		viper.GetString("dbpath"),
 		viper.GetBool("debug-sql"),
+		db.Option{},
 	)
 	if err != nil {
 		if locked {

diff --git a/commands/server.go b/commands/server.go
@@ -36,7 +36,7 @@ func executeServer(cmd *cobra.Command, args []string) (err error) {
 		return xerrors.Errorf("Failed to SetLogger. err: %w", err)
 	}
 
-	driver, locked, err := db.NewDB(viper.GetString("dbtype"), viper.GetString("dbpath"), viper.GetBool("debug-sql"))
+	driver, locked, err := db.NewDB(viper.GetString("dbtype"), viper.GetString("dbpath"), viper.GetBool("debug-sql"), db.Option{})
 	if err != nil {
 		if locked {
 			return xerrors.Errorf("Failed to initialize DB. Close DB connection before fetching. err: %w", err)

diff --git a/db/db.go b/db/db.go
@@ -2,6 +2,7 @@ package db
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/vulsio/go-exploitdb/models"
 	"golang.org/x/xerrors"
@@ -10,7 +11,7 @@ import (
 // DB :
 type DB interface {
 	Name() string
-	OpenDB(dbType, dbPath string, debugSQL bool) (bool, error)
+	OpenDB(dbType, dbPath string, debugSQL bool, option Option) (bool, error)
 	CloseDB() error
 	MigrateDB() error
 	GetExploitByID(string) ([]models.Exploit, error)
@@ -25,13 +26,17 @@ type DB interface {
 	UpsertFetchMeta(*models.FetchMeta) error
 }
 
+type Option struct {
+	RedisTimeout time.Duration
+}
+
 // NewDB :
-func NewDB(dbType string, dbPath string, debugSQL bool) (driver DB, locked bool, err error) {
+func NewDB(dbType string, dbPath string, debugSQL bool, option Option) (driver DB, locked bool, err error) {
 	if driver, err = newDB(dbType); err != nil {
 		return driver, false, xerrors.Errorf("Failed to new db. err: %w", err)
 	}
 
-	if locked, err := driver.OpenDB(dbType, dbPath, debugSQL); err != nil {
+	if locked, err := driver.OpenDB(dbType, dbPath, debugSQL, option); err != nil {
 		if locked {
 			return nil, true, err
 		}

diff --git a/db/rdb.go b/db/rdb.go
@@ -42,7 +42,7 @@ func (r *RDBDriver) Name() string {
 }
 
 // OpenDB opens Database
-func (r *RDBDriver) OpenDB(dbType, dbPath string, debugSQL bool) (locked bool, err error) {
+func (r *RDBDriver) OpenDB(dbType, dbPath string, debugSQL bool, option Option) (locked bool, err error) {
 	gormConfig := gorm.Config{
 		DisableForeignKeyConstraintWhenMigrating: true,
 		Logger: logger.New(

diff --git a/db/redis.go b/db/redis.go
@@ -61,8 +61,8 @@ func (r *RedisDriver) Name() string {
 }
 
 // OpenDB opens Database
-func (r *RedisDriver) OpenDB(dbType, dbPath string, debugSQL bool) (locked bool, err error) {
-	if err = r.connectRedis(dbPath); err != nil {
+func (r *RedisDriver) OpenDB(dbType, dbPath string, debugSQL bool, option Option) (locked bool, err error) {
+	if err = r.connectRedis(dbPath, option); err != nil {
 		err = xerrors.Errorf("Failed to open DB. dbtype: %s, dbpath: %s, err: %w", dbType, dbPath, err)
 	}
 	return
@@ -79,12 +79,15 @@ func (r *RedisDriver) CloseDB() (err error) {
 	return
 }
 
-func (r *RedisDriver) connectRedis(dbPath string) error {
-	option, err := redis.ParseURL(dbPath)
+func (r *RedisDriver) connectRedis(dbPath string, option Option) error {
+	opt, err := redis.ParseURL(dbPath)
 	if err != nil {
 		return xerrors.Errorf("Failed to parse url. err: %w", err)
 	}
-	r.conn = redis.NewClient(option)
+	if 0 < option.RedisTimeout.Seconds() {
+		opt.ReadTimeout = option.RedisTimeout
+	}
+	r.conn = redis.NewClient(opt)
 	return r.conn.Ping(context.Background()).Err()
 }
 

diff --git a/go.mod b/go.mod
@@ -4,7 +4,7 @@ go 1.15
 
 require (
 	github.com/Microsoft/go-winio v0.5.0 // indirect
-	github.com/asaskevich/govalidator v0.0.0-20200907205600-7a23bdc65eef
+	github.com/asaskevich/govalidator v0.0.0-20200907205600-7a23bdc65eef // indirect
 	github.com/cheggaaa/pb/v3 v3.0.5
 	github.com/elazarl/goproxy v0.0.0-20210110162100-a92cc753f88e // indirect
 	github.com/go-redis/redis/v8 v8.4.10
@@ -17,7 +17,7 @@ require (
 	github.com/mattn/go-sqlite3 v1.14.7
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/parnurzeal/gorequest v0.2.16
-	github.com/pkg/errors v0.9.1
+	github.com/pkg/errors v0.9.1 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0
 	github.com/sergi/go-diff v1.2.0 // indirect
 	github.com/spf13/cast v1.4.0 // indirect