feat(redis): change Redis Architecture (#57)

* chore: go mod tidy -v

* style: remove blank line

* chore: do not name stash

* chore(integration): change git stash in build-integration

* fix(integration): fix makefile

* fix(integration): git stash

* feat(redis): set to nil when ExploitUniqueID is not filled

* fix: fix panic error

* feat(redis): fix unimplemented parts

* feat(redis): change redis architecture

* chore: add const fetchMetaKey

* feat(redis): add batch size option

* chore: check batch-size

* chore: do not name stash

* refactor(db): GetExploitAll fast and safe

* feat(redis): change from SMembers to SScan to avoid blocking server

* chore: remove unnecessary arg

* feat(redis): update Redis Architecture

* feat(redis): update Redis Architecture

* fix(redis): insert error

* fix(redis): fix get

* fix(fetcher): change ExploitUniqueID to searchable

* fix(fetcher): use base64.URLEncoding

* feat(models): remove unnecessary ExploitUniqueID

* fix(fetcher): use MD5 instead of base64.URLEncoding

* feat(redis): use pipeline in Get

* fix(redis): fix IsExploitModelV1

* style: fix log msg

* chore: use smembers

* feat(redis): use Hash

* style: change bar.Finish timing

* refactor(redis): map exists check

* chore: delete ExploitUniqueID field explicitly

* chore: delete ExploitUniqueID field explicitly

* style: add comment

* fix(redis): skip if cveid is empty string

* feat(redis): change redis key

* chore: add slice len check

* chore(redis): update Get

* chore(integration): update test

* chore(integration): move to integration/.gitignore

* chore: change integration/.gitignore

.gitignore

@@ -9,5 +9,4 @@ vendor/
 *.sqlite3-shm
 *.sqlite3-wal
 go-exploitdb
-
 .vscode

GNUmakefile

@@ -71,7 +71,7 @@ clean:
 PWD := $(shell pwd)
 BRANCH := $(shell git symbolic-ref --short HEAD)
 build-integration:
-	@ git stash save "build-integration"
+	@ git stash save
 	go build -ldflags "$(LDFLAGS)" -o integration/exploitdb.new
 	git checkout $(shell git describe --tags --abbrev=0)
 	@git reset --hard
@@ -83,6 +83,7 @@ clean-integration:
 	-pkill exploitdb.old
 	-pkill exploitdb.new
 	-rm integration/exploitdb.old integration/exploitdb.new integration/go-exploitdb.old.sqlite3 integration/go-exploitdb.new.sqlite3
+	-rm -rf integration/diff
 	-docker kill redis-old redis-new
 	-docker rm redis-old redis-new
 
@@ -108,14 +109,10 @@ fetch-redis:
 	integration/exploitdb.new fetch githubrepos --dbtype redis --dbpath "redis://127.0.0.1:6380/0"
 
 diff-cveid:
-	@ python integration/diff_server_mode.py cveid --sample_rate 0.01 awesomepoc
-	@ python integration/diff_server_mode.py cveid --sample_rate 0.01 exploitdb
-	@ python integration/diff_server_mode.py cveid --sample_rate 0.01 githubrepos
+	@ python integration/diff_server_mode.py cveid --sample_rate 0.01
 
 diff-uniqueid:
-	@ python integration/diff_server_mode.py uniqueid --sample_rate 0.01 awesomepoc
-	@ python integration/diff_server_mode.py uniqueid --sample_rate 0.01 exploitdb
-	@ python integration/diff_server_mode.py uniqueid --sample_rate 0.01 githubrepos
+	@ python integration/diff_server_mode.py uniqueid --sample_rate 0.01
 
 diff-server-rdb:
 	integration/exploitdb.old server --dbpath=$(PWD)/integration/go-exploitdb.old.sqlite3 --port 1325 > /dev/null 2>&1 & 

commands/fetch.go

@@ -18,6 +18,6 @@ func init() {
 	fetchCmd.PersistentFlags().Uint("expire", 0, "timeout to set for Redis keys in seconds. If set to 0, the key is persistent.")
 	_ = viper.BindPFlag("expire", fetchCmd.PersistentFlags().Lookup("expire"))
 
-	fetchCmd.PersistentFlags().Int("batch-size", 500, "The number of batch size to insert. NOTE: This Option does not work for dbtype: redis.")
+	fetchCmd.PersistentFlags().Int("batch-size", 500, "The number of batch size to insert.")
 	_ = viper.BindPFlag("batch-size", fetchCmd.PersistentFlags().Lookup("batch-size"))
 }

db/db.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 
 	"github.com/inconshreveable/log15"
-	"github.com/spf13/viper"
 	"github.com/vulsio/go-exploitdb/models"
 	"golang.org/x/xerrors"
 )
@@ -61,9 +60,32 @@ func NewDB(dbType string, dbPath string, debugSQL bool) (driver DB, locked bool,
 func newDB(dbType string) (DB, error) {
 	switch dbType {
 	case dialectSqlite3, dialectMysql, dialectPostgreSQL:
-		return &RDBDriver{name: dbType, batchSize: viper.GetInt("batch-size")}, nil
+		return &RDBDriver{name: dbType}, nil
 	case dialectRedis:
 		return &RedisDriver{name: dbType}, nil
 	}
 	return nil, fmt.Errorf("Invalid database dialect, %s", dbType)
 }
+
+// IndexChunk has a starting point and an ending point for Chunk
+type IndexChunk struct {
+	From, To int
+}
+
+func chunkSlice(length int, chunkSize int) <-chan IndexChunk {
+	ch := make(chan IndexChunk)
+
+	go func() {
+		defer close(ch)
+
+		for i := 0; i < length; i += chunkSize {
+			idx := IndexChunk{i, i + chunkSize}
+			if length < idx.To {
+				idx.To = length
+			}
+			ch <- idx
+		}
+	}()
+
+	return ch
+}

db/rdb.go

@@ -11,6 +11,7 @@ import (
 	"github.com/cheggaaa/pb/v3"
 	"github.com/inconshreveable/log15"
 	sqlite3 "github.com/mattn/go-sqlite3"
+	"github.com/spf13/viper"
 	"github.com/vulsio/go-exploitdb/config"
 	"github.com/vulsio/go-exploitdb/models"
 	"github.com/vulsio/go-exploitdb/util"
@@ -32,9 +33,8 @@ const (
 
 // RDBDriver :
 type RDBDriver struct {
-	name      string
-	conn      *gorm.DB
-	batchSize int
+	name string
+	conn *gorm.DB
 }
 
 // Name return db name
@@ -130,11 +130,11 @@ func (r *RDBDriver) MigrateDB() error {
 // InsertExploit :
 func (r *RDBDriver) InsertExploit(exploitType models.ExploitType, exploits []models.Exploit) (err error) {
 	log15.Info(fmt.Sprintf("Inserting %d Exploits", len(exploits)))
-	return r.deleteAndInsertExploit(r.conn, exploitType, exploits)
+	return r.deleteAndInsertExploit(exploitType, exploits)
 }
 
-func (r *RDBDriver) deleteAndInsertExploit(conn *gorm.DB, exploitType models.ExploitType, exploits []models.Exploit) (err error) {
-	tx := conn.Begin()
+func (r *RDBDriver) deleteAndInsertExploit(exploitType models.ExploitType, exploits []models.Exploit) (err error) {
+	tx := r.conn.Begin()
 	defer func() {
 		if err != nil {
 			tx.Rollback()
@@ -143,6 +143,11 @@ func (r *RDBDriver) deleteAndInsertExploit(conn *gorm.DB, exploitType models.Exp
 		tx.Commit()
 	}()
 
+	batchSize := viper.GetInt("batch-size")
+	if batchSize < 1 {
+		return xerrors.New("Failed to set batch-size. err: batch-size option is not set properly")
+	}
+
 	oldIDs := []int64{}
 	result := tx.Model(&models.Exploit{}).Select("id").Where("exploit_type = ?", exploitType).Find(&oldIDs)
 	if result.Error != nil && !errors.Is(result.Error, gorm.ErrRecordNotFound) {
@@ -152,7 +157,7 @@ func (r *RDBDriver) deleteAndInsertExploit(conn *gorm.DB, exploitType models.Exp
 	if result.RowsAffected > 0 {
 		log15.Info("Deleting old Exploits")
 		bar := pb.StartNew(len(oldIDs))
-		for idx := range chunkSlice(len(oldIDs), r.batchSize) {
+		for idx := range chunkSlice(len(oldIDs), batchSize) {
 			osIDs := []int64{}
 			if err := tx.Model(&models.OffensiveSecurity{}).Select("id").Where("exploit_id IN ?", oldIDs[idx.From:idx.To]).Find(&osIDs).Error; err != nil {
 				return xerrors.Errorf("Failed to select old OffensiveSecurity: %w", err)
@@ -184,7 +189,7 @@ func (r *RDBDriver) deleteAndInsertExploit(conn *gorm.DB, exploitType models.Exp
 
 	log15.Info("Inserting new Exploits")
 	bar := pb.StartNew(len(exploits))
-	for idx := range chunkSlice(len(exploits), r.batchSize) {
+	for idx := range chunkSlice(len(exploits), batchSize) {
 		if err = tx.Create(exploits[idx.From:idx.To]).Error; err != nil {
 			return fmt.Errorf("Failed to insert. err: %s", err)
 		}
@@ -231,43 +236,38 @@ func (r *RDBDriver) GetExploitByID(exploitUniqueID string) []models.Exploit {
 // GetExploitAll :
 func (r *RDBDriver) GetExploitAll() []models.Exploit {
 	es := []models.Exploit{}
-	docs := []models.Document{}
-	shells := []models.ShellCode{}
-	papers := []models.Paper{}
-	offensiveSecurities := []models.OffensiveSecurity{}
-	var errs util.Errors
 
-	errs = errs.Add(r.conn.Find(&es).Error)
-	errs = errs.Add(r.conn.Find(&offensiveSecurities).Error)
-	errs = errs.Add(r.conn.Find(&docs).Error)
-	errs = errs.Add(r.conn.Find(&shells).Error)
-	errs = errs.Add(r.conn.Find(&papers).Error)
-	if len(errs.GetErrors()) > 0 {
-		log15.Error("Failed to delete old records", "err", errs.Error())
+	rows, err := r.conn.Model(&models.Exploit{}).Rows()
+	if err != nil {
+		log15.Error("Failed to Rows", "err", err)
+		return nil
 	}
+	defer rows.Close()
 
-	for _, e := range es {
-		for _, o := range offensiveSecurities {
-			for _, d := range docs {
-				if o.ID == d.OffensiveSecurityID {
-					o.Document = &d
-				}
-			}
-			for _, s := range shells {
-				if o.ID == s.OffensiveSecurityID {
-					o.ShellCode = &s
-				}
-			}
-			for _, p := range papers {
-				if o.ID == p.OffensiveSecurityID {
-					o.Paper = &p
-				}
+	for rows.Next() {
+		exploit := models.Exploit{}
+		if err := r.conn.ScanRows(rows, &exploit); err != nil {
+			log15.Error("Failed to ScanRows", "err", err)
+			return nil
+		}
+		switch exploit.ExploitType {
+		case models.OffensiveSecurityType:
+			if err := r.conn.
+				Preload(clause.Associations).
+				Where(&models.OffensiveSecurity{ExploitID: exploit.ID}).
+				Take(&exploit.OffensiveSecurity).Error; err != nil {
+				log15.Error("Failed to Get OffensiveSecurity", "err", err)
+				return nil
 			}
-			if e.ID == o.ExploitID {
-				e.OffensiveSecurity = &o
+		case models.GitHubRepositoryType:
+			if err := r.conn.Where(&models.GitHubRepository{ExploitID: exploit.ID}).Take(&exploit.GitHubRepository).Error; err != nil {
+				log15.Error("Failed to Get GitHubRepository", "err", err)
+				return nil
 			}
 		}
+		es = append(es, exploit)
 	}
+
 	return es
 }
 
@@ -315,29 +315,6 @@ func (r *RDBDriver) GetExploitMultiByCveID(cveIDs []string) (exploits map[string
 	return exploits
 }
 
-// IndexChunk has a starting point and an ending point for Chunk
-type IndexChunk struct {
-	From, To int
-}
-
-func chunkSlice(length int, chunkSize int) <-chan IndexChunk {
-	ch := make(chan IndexChunk)
-
-	go func() {
-		defer close(ch)
-
-		for i := 0; i < length; i += chunkSize {
-			idx := IndexChunk{i, i + chunkSize}
-			if length < idx.To {
-				idx.To = length
-			}
-			ch <- idx
-		}
-	}()
-
-	return ch
-}
-
 // IsExploitModelV1 determines if the DB was created at the time of go-exploitdb Model v1
 func (r *RDBDriver) IsExploitModelV1() (bool, error) {
 	if r.conn.Migrator().HasTable(&models.FetchMeta{}) {