From bbe0b9fa421e86b0408e3c395b662af6101bbc86 Mon Sep 17 00:00:00 2001 From: MaineK00n Date: Fri, 27 Sep 2024 15:46:37 +0900 Subject: [PATCH] refactor: use std slices, maps package (#278) --- db/arch.go | 7 +++-- db/db.go | 23 --------------- db/debian.go | 7 +++-- db/microsoft.go | 22 +++++++------- db/rdb.go | 4 +-- db/redhat.go | 7 +++-- db/redis.go | 72 ++++++++++++++++++++++----------------------- db/ubuntu.go | 7 +++-- go.mod | 2 +- models/debian.go | 7 +++-- models/microsoft.go | 3 +- util/util.go | 5 ++-- 12 files changed, 74 insertions(+), 92 deletions(-) diff --git a/db/arch.go b/db/arch.go index 983b7b9..a94ba1c 100644 --- a/db/arch.go +++ b/db/arch.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "os" + "slices" "github.com/cheggaaa/pb/v3" "github.com/spf13/viper" @@ -82,11 +83,11 @@ func (r *RDBDriver) deleteAndInsertArch(advs []models.ArchADV) (err error) { return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly") } - for idx := range chunkSlice(len(advs), batchSize) { - if err = tx.Create(advs[idx.From:idx.To]).Error; err != nil { + for chunk := range slices.Chunk(advs, batchSize) { + if err = tx.Create(chunk).Error; err != nil { return xerrors.Errorf("Failed to insert. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() diff --git a/db/db.go b/db/db.go index cbc683c..21a18ac 100644 --- a/db/db.go +++ b/db/db.go @@ -91,26 +91,3 @@ func newDB(dbType string) (DB, error) { } return nil, fmt.Errorf("Invalid database dialect. dbType: %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 -} diff --git a/db/debian.go b/db/debian.go index c4dd970..1bc54fa 100644 --- a/db/debian.go +++ b/db/debian.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "os" + "slices" "github.com/cheggaaa/pb/v3" "github.com/spf13/viper" @@ -90,11 +91,11 @@ func (r *RDBDriver) deleteAndInsertDebian(cves []models.DebianCVE) (err error) { return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly") } - for idx := range chunkSlice(len(cves), batchSize) { - if err = tx.Create(cves[idx.From:idx.To]).Error; err != nil { + for chunk := range slices.Chunk(cves, batchSize) { + if err = tx.Create(chunk).Error; err != nil { return xerrors.Errorf("Failed to insert. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() diff --git a/db/microsoft.go b/db/microsoft.go index ccbde92..b2917ef 100644 --- a/db/microsoft.go +++ b/db/microsoft.go @@ -4,15 +4,15 @@ import ( "errors" "fmt" "io" + "maps" "os" + "slices" "strconv" "strings" "github.com/cheggaaa/pb/v3" "github.com/inconshreveable/log15" "github.com/spf13/viper" - "golang.org/x/exp/maps" - "golang.org/x/exp/slices" "golang.org/x/xerrors" "gorm.io/gorm" @@ -69,7 +69,7 @@ func (r *RDBDriver) GetExpandKB(applied []string, unapplied []string) ([]string, uniqUnappliedKBIDs[kbID] = struct{}{} delete(uniqAppliedKBIDs, kbID) } - applied = maps.Keys(uniqAppliedKBIDs) + applied = slices.Collect(maps.Keys(uniqAppliedKBIDs)) if len(applied) > 0 { relations := []models.MicrosoftKBRelation{} @@ -102,7 +102,7 @@ func (r *RDBDriver) GetExpandKB(applied []string, unapplied []string) ([]string, if err := r.conn. Preload("SupersededBy"). - Where("kb_id IN ?", maps.Keys(uniqUnappliedKBIDs)). + Where("kb_id IN ?", slices.Collect(maps.Keys(uniqUnappliedKBIDs))). Find(&relations).Error; err != nil { return nil, nil, xerrors.Errorf("Failed to get KB Relation by unapplied KBID: %q. err: %w", unapplied, err) } @@ -114,7 +114,7 @@ func (r *RDBDriver) GetExpandKB(applied []string, unapplied []string) ([]string, } } - return applied, maps.Keys(uniqUnappliedKBIDs), nil + return applied, slices.Collect(maps.Keys(uniqUnappliedKBIDs)), nil } // GetRelatedProducts : @@ -284,11 +284,11 @@ func (r *RDBDriver) deleteAndInsertMicrosoft(cves []models.MicrosoftCVE) (err er return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly") } - for idx := range chunkSlice(len(cves), batchSize) { - if err = tx.Create(cves[idx.From:idx.To]).Error; err != nil { + for chunk := range slices.Chunk(cves, batchSize) { + if err = tx.Create(chunk).Error; err != nil { return xerrors.Errorf("Failed to insert. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() @@ -324,11 +324,11 @@ func (r *RDBDriver) deleteAndInsertMicrosoftKBRelation(kbs []models.MicrosoftKBR return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly") } - for idx := range chunkSlice(len(kbs), batchSize) { - if err = tx.Create(kbs[idx.From:idx.To]).Error; err != nil { + for chunk := range slices.Chunk(kbs, batchSize) { + if err = tx.Create(chunk).Error; err != nil { return xerrors.Errorf("Failed to insert. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() return nil diff --git a/db/rdb.go b/db/rdb.go index a6f7fc3..d737953 100644 --- a/db/rdb.go +++ b/db/rdb.go @@ -186,9 +186,7 @@ func (r *RDBDriver) MigrateDB() error { } } case dialectMysql, dialectPostgreSQL: - if err != nil { - return xerrors.Errorf("Failed to migrate. err: %w", err) - } + return xerrors.Errorf("Failed to migrate. err: %w", err) default: return xerrors.Errorf("Not Supported DB dialects. r.name: %s", r.name) } diff --git a/db/redhat.go b/db/redhat.go index 061ec67..c40c02d 100644 --- a/db/redhat.go +++ b/db/redhat.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "os" + "slices" "strings" "time" @@ -222,11 +223,11 @@ func (r *RDBDriver) deleteAndInsertRedhat(cves []models.RedhatCVE) (err error) { return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly") } - for idx := range chunkSlice(len(cves), batchSize) { - if err = tx.Create(cves[idx.From:idx.To]).Error; err != nil { + for chunk := range slices.Chunk(cves, batchSize) { + if err = tx.Create(chunk).Error; err != nil { return xerrors.Errorf("Failed to insert. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() diff --git a/db/redis.go b/db/redis.go index 6a205d1..e688c97 100644 --- a/db/redis.go +++ b/db/redis.go @@ -6,7 +6,9 @@ import ( "errors" "fmt" "io" + "maps" "os" + "slices" "strconv" "strings" "time" @@ -15,8 +17,6 @@ import ( "github.com/go-redis/redis/v8" "github.com/inconshreveable/log15" "github.com/spf13/viper" - "golang.org/x/exp/maps" - "golang.org/x/exp/slices" "golang.org/x/xerrors" "github.com/vulsio/gost/config" @@ -445,8 +445,8 @@ func (r *RedisDriver) getCvesUbuntuWithFixStatus(major, pkgName string, fixStatu } m := map[string]models.UbuntuCVE{} - for idx := range chunkSlice(len(cveIDs), 20) { - res, err := r.GetUbuntuMulti(cveIDs[idx.From:idx.To]) + for chunk := range slices.Chunk(cveIDs, 20) { + res, err := r.GetUbuntuMulti(chunk) if err != nil { return nil, xerrors.Errorf("Failed to GetUbuntuMulti. err: %w", err) } @@ -578,7 +578,7 @@ func (r *RedisDriver) GetExpandKB(applied []string, unapplied []string) ([]strin uniqUnappliedKBIDs[kbID] = struct{}{} delete(uniqAppliedKBIDs, kbID) } - applied = maps.Keys(uniqAppliedKBIDs) + applied = slices.Collect(maps.Keys(uniqAppliedKBIDs)) pipe := r.conn.Pipeline() for _, kbID := range applied { @@ -626,7 +626,7 @@ func (r *RedisDriver) GetExpandKB(applied []string, unapplied []string) ([]strin } } - return applied, maps.Keys(uniqUnappliedKBIDs), nil + return applied, slices.Collect(maps.Keys(uniqUnappliedKBIDs)), nil } // GetRelatedProducts : @@ -920,10 +920,10 @@ func (r *RedisDriver) InsertRedhat(cves []models.RedhatCVE) (err error) { } return os.Stderr }()) - for idx := range chunkSlice(len(cves), batchSize) { + for chunk := range slices.Chunk(cves, batchSize) { pipe := r.conn.Pipeline() cvekey := fmt.Sprintf(cveKeyFormat, redhatName) - for _, cve := range cves[idx.From:idx.To] { + for _, cve := range chunk { j, err := json.Marshal(cve) if err != nil { return xerrors.Errorf("Failed to marshal json. err: %w", err) @@ -950,7 +950,7 @@ func (r *RedisDriver) InsertRedhat(cves []models.RedhatCVE) (err error) { if _, err := pipe.Exec(ctx); err != nil { return xerrors.Errorf("Failed to exec pipeline. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() @@ -961,10 +961,10 @@ func (r *RedisDriver) InsertRedhat(cves []models.RedhatCVE) (err error) { } return os.Stderr }()) - keys := maps.Keys(advs) - for idx := range chunkSlice(len(keys), batchSize) { + keys := slices.Collect(maps.Keys(advs)) + for chunk := range slices.Chunk(keys, batchSize) { pipe := r.conn.Pipeline() - for _, adv := range keys[idx.From:idx.To] { + for _, adv := range chunk { var aj []byte if aj, err = json.Marshal(advs[adv]); err != nil { return xerrors.Errorf("Failed to marshal json. err: %w", err) @@ -979,7 +979,7 @@ func (r *RedisDriver) InsertRedhat(cves []models.RedhatCVE) (err error) { if _, err = pipe.Exec(ctx); err != nil { return xerrors.Errorf("Failed to exec pipeline. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() @@ -1040,10 +1040,10 @@ func (r *RedisDriver) InsertDebian(cves []models.DebianCVE) error { } return os.Stderr }()) - for idx := range chunkSlice(len(cves), batchSize) { + for chunk := range slices.Chunk(cves, batchSize) { pipe := r.conn.Pipeline() cvekey := fmt.Sprintf(cveKeyFormat, debianName) - for _, cve := range cves[idx.From:idx.To] { + for _, cve := range chunk { j, err := json.Marshal(cve) if err != nil { return xerrors.Errorf("Failed to marshal json. err: %w", err) @@ -1070,7 +1070,7 @@ func (r *RedisDriver) InsertDebian(cves []models.DebianCVE) error { if _, err := pipe.Exec(ctx); err != nil { return xerrors.Errorf("Failed to exec pipeline. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() @@ -1136,10 +1136,10 @@ func (r *RedisDriver) InsertUbuntu(cves []models.UbuntuCVE) (err error) { } return os.Stderr }()) - for idx := range chunkSlice(len(cves), batchSize) { + for chunk := range slices.Chunk(cves, batchSize) { pipe := r.conn.Pipeline() cvekey := fmt.Sprintf(cveKeyFormat, ubuntuName) - for _, cve := range cves[idx.From:idx.To] { + for _, cve := range chunk { j, err := json.Marshal(cve) if err != nil { return xerrors.Errorf("Failed to marshal json. err: %w", err) @@ -1166,7 +1166,7 @@ func (r *RedisDriver) InsertUbuntu(cves []models.UbuntuCVE) (err error) { if _, err := pipe.Exec(ctx); err != nil { return xerrors.Errorf("Failed to exec pipeline. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() @@ -1177,10 +1177,10 @@ func (r *RedisDriver) InsertUbuntu(cves []models.UbuntuCVE) (err error) { } return os.Stderr }()) - keys := maps.Keys(advs) - for idx := range chunkSlice(len(keys), batchSize) { + keys := slices.Collect(maps.Keys(advs)) + for chunk := range slices.Chunk(keys, batchSize) { pipe := r.conn.Pipeline() - for _, adv := range keys[idx.From:idx.To] { + for _, adv := range chunk { var aj []byte if aj, err = json.Marshal(advs[adv]); err != nil { return xerrors.Errorf("Failed to marshal json. err: %w", err) @@ -1195,7 +1195,7 @@ func (r *RedisDriver) InsertUbuntu(cves []models.UbuntuCVE) (err error) { if _, err = pipe.Exec(ctx); err != nil { return xerrors.Errorf("Failed to exec pipeline. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() @@ -1280,10 +1280,10 @@ func (r *RedisDriver) InsertMicrosoft(cves []models.MicrosoftCVE, relations []mo } return os.Stderr }()) - for idx := range chunkSlice(len(cves), batchSize) { + for chunk := range slices.Chunk(cves, batchSize) { pipe := r.conn.Pipeline() cvekey := fmt.Sprintf(cveKeyFormat, microsoftName) - for _, cve := range cves[idx.From:idx.To] { + for _, cve := range chunk { j, err := json.Marshal(cve) if err != nil { return xerrors.Errorf("Failed to marshal json. err: %w", err) @@ -1324,7 +1324,7 @@ func (r *RedisDriver) InsertMicrosoft(cves []models.MicrosoftCVE, relations []mo if _, err := pipe.Exec(ctx); err != nil { return xerrors.Errorf("Failed to exec pipeline. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() @@ -1335,10 +1335,10 @@ func (r *RedisDriver) InsertMicrosoft(cves []models.MicrosoftCVE, relations []mo } return os.Stderr }()) - keys := maps.Keys(advs) - for idx := range chunkSlice(len(keys), batchSize) { + keys := slices.Collect(maps.Keys(advs)) + for chunk := range slices.Chunk(keys, batchSize) { pipe := r.conn.Pipeline() - for _, adv := range keys[idx.From:idx.To] { + for _, adv := range chunk { var aj []byte if aj, err = json.Marshal(advs[adv]); err != nil { return xerrors.Errorf("Failed to marshal json. err: %w", err) @@ -1353,7 +1353,7 @@ func (r *RedisDriver) InsertMicrosoft(cves []models.MicrosoftCVE, relations []mo if _, err = pipe.Exec(ctx); err != nil { return xerrors.Errorf("Failed to exec pipeline. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() @@ -1364,9 +1364,9 @@ func (r *RedisDriver) InsertMicrosoft(cves []models.MicrosoftCVE, relations []mo } return os.Stderr }()) - for idx := range chunkSlice(len(relations), batchSize) { + for chunk := range slices.Chunk(relations, batchSize) { pipe := r.conn.Pipeline() - for _, relation := range relations[idx.From:idx.To] { + for _, relation := range chunk { key := fmt.Sprintf(pkgKeyFormat, microsoftName, fmt.Sprintf("K#%s", relation.KBID)) if _, ok := newDeps["relations"][relation.KBID]; !ok { newDeps["relations"][relation.KBID] = map[string]struct{}{} @@ -1385,7 +1385,7 @@ func (r *RedisDriver) InsertMicrosoft(cves []models.MicrosoftCVE, relations []mo if _, err := pipe.Exec(ctx); err != nil { return xerrors.Errorf("Failed to exec pipeline. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() @@ -1461,9 +1461,9 @@ func (r *RedisDriver) InsertArch(advs []models.ArchADV) error { } return os.Stderr }()) - for idx := range chunkSlice(len(advs), batchSize) { + for chunk := range slices.Chunk(advs, batchSize) { pipe := r.conn.Pipeline() - for _, adv := range advs[idx.From:idx.To] { + for _, adv := range chunk { j, err := json.Marshal(adv) if err != nil { return xerrors.Errorf("Failed to marshal json. err: %w", err) @@ -1511,7 +1511,7 @@ func (r *RedisDriver) InsertArch(advs []models.ArchADV) error { if _, err := pipe.Exec(ctx); err != nil { return xerrors.Errorf("Failed to exec pipeline. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() diff --git a/db/ubuntu.go b/db/ubuntu.go index 9826099..d87556b 100644 --- a/db/ubuntu.go +++ b/db/ubuntu.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "os" + "slices" "strings" "github.com/cheggaaa/pb/v3" @@ -114,11 +115,11 @@ func (r *RDBDriver) deleteAndInsertUbuntu(cves []models.UbuntuCVE) (err error) { return xerrors.New("Failed to set batch-size. err: batch-size option is not set properly") } - for idx := range chunkSlice(len(cves), batchSize) { - if err = tx.Create(cves[idx.From:idx.To]).Error; err != nil { + for chunk := range slices.Chunk(cves, batchSize) { + if err = tx.Create(chunk).Error; err != nil { return xerrors.Errorf("Failed to insert. err: %w", err) } - bar.Add(idx.To - idx.From) + bar.Add(len(chunk)) } bar.Finish() diff --git a/go.mod b/go.mod index 05de780..ed17256 100644 --- a/go.mod +++ b/go.mod @@ -21,7 +21,6 @@ require ( github.com/spf13/cobra v1.8.1 github.com/spf13/viper v1.19.0 go.etcd.io/bbolt v1.3.11 - golang.org/x/exp v0.0.0-20230905200255-921286631fa9 golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 gorm.io/driver/mysql v1.5.5 gorm.io/driver/postgres v1.5.7 @@ -87,6 +86,7 @@ require ( go.uber.org/atomic v1.9.0 // indirect go.uber.org/multierr v1.9.0 // indirect golang.org/x/crypto v0.22.0 // indirect + golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect golang.org/x/mod v0.12.0 // indirect golang.org/x/net v0.24.0 // indirect golang.org/x/sync v0.6.0 // indirect diff --git a/models/debian.go b/models/debian.go index 322c2e6..64e6d09 100644 --- a/models/debian.go +++ b/models/debian.go @@ -1,6 +1,9 @@ package models -import "golang.org/x/exp/maps" +import ( + "maps" + "slices" +) // DebianJSON : type DebianJSON map[string]DebianCveMap @@ -87,5 +90,5 @@ func ConvertDebian(cveJSONs DebianJSON) []DebianCVE { } } } - return maps.Values(uniqCve) + return slices.Collect(maps.Values(uniqCve)) } diff --git a/models/microsoft.go b/models/microsoft.go index bc27b5a..b6df6bc 100644 --- a/models/microsoft.go +++ b/models/microsoft.go @@ -1,10 +1,9 @@ package models import ( + "slices" "strings" "time" - - "golang.org/x/exp/slices" ) // MicrosoftVulnerability : diff --git a/util/util.go b/util/util.go index 32ce911..68936fa 100644 --- a/util/util.go +++ b/util/util.go @@ -4,10 +4,12 @@ import ( "bytes" "fmt" "io" + "maps" "os" "os/exec" "path/filepath" "runtime" + "slices" "strings" "time" @@ -16,7 +18,6 @@ import ( "github.com/inconshreveable/log15" "github.com/parnurzeal/gorequest" "github.com/spf13/viper" - "golang.org/x/exp/maps" "golang.org/x/xerrors" ) @@ -26,7 +27,7 @@ func Unique[T comparable](s []T) []T { for _, v := range s { m[v] = struct{}{} } - return maps.Keys(m) + return slices.Collect(maps.Keys(m)) } // GenWorkers generate workers