Skip to content

Commit

Permalink
feat!(db/redis): use CPE#Titles instead of CPE#Cache#Titles
Browse files Browse the repository at this point in the history
  • Loading branch information
@MaineK00n
MaineK00n committed Nov 5, 2024
1 parent 6920f82 commit 515af55
Showing 1 changed file with 80 additions and 27 deletions.
107 changes: 80 additions & 27 deletions db/redis.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"errors"
"fmt"
"io"
"maps"
"os"
"slices"
"strconv"
Expand All @@ -29,7 +30,7 @@ import (
│NO │ KEY │ VALUE │ PURPOSE │
└───┴──────────────────┴──────────┴──────────────────────────────────────────────────┘
┌───┬──────────────────┬──────────┬──────────────────────────────────────────────────┐
│ 1 │ CPE#Cache#Titles │ JSON │ TO CACHE CPE#Titles │
│ 1 │ CPE#Titles │ JSON │ Get ALL Titles │
└───┴──────────────────┴──────────┴──────────────────────────────────────────────────┘
- Sets
Expand All @@ -45,8 +46,6 @@ import (
├─────────────────────────────┼───────────────────────┼────────────────────────────────────┤
│ CPE#DeprecatedCPEs │ CPEURI │ Get DeprecatedCPEs │
├─────────────────────────────┼───────────────────────┼────────────────────────────────────┤
│ CPE#Titles │ Title │ Get ALL Titles │
├─────────────────────────────┼───────────────────────┼────────────────────────────────────┤
│ CPE#Title#{title} │ CPEURI │ Get CPEURI by title │
└─────────────────────────────┴───────────────────────┴────────────────────────────────────┘
Expand Down Expand Up @@ -256,20 +255,38 @@ func (r *RedisDriver) GetSimilarCpesByTitle(query string, n int, algorithm edlib
}

ctx := context.Background()
t, err := r.conn.Type(ctx, titleListKey).Result()
if err != nil {
return nil, xerrors.Errorf("Failed to TYPE CPE#Titles. err: %w", err)
}

var ts []string
bs, err := r.conn.Get(ctx, titleListCacheKey).Bytes()
if err == nil {
switch t {
case "string":
bs, err := r.conn.Get(ctx, titleListKey).Bytes()
if err != nil {
return nil, xerrors.Errorf("Failed to Get Titles. err: %w", err)
}
if err := json.Unmarshal(bs, &ts); err != nil {
return nil, xerrors.Errorf("Failed to Unmarshal JSON. err: %w", err)
}
} else {
if !xerrors.Is(err, redis.Nil) {
return nil, xerrors.Errorf("Failed to Get Titles. err: %w", err)
}
ts, err = r.conn.SMembers(ctx, titleListKey).Result()
if err != nil {
return nil, xerrors.Errorf("Failed to SMembers Titles. err: %w", err)
case "set": // backward compatibility: https://github.com/vulsio/go-cpe-dictionary/pull/186
bs, err := r.conn.Get(ctx, titleListCacheKey).Bytes()
if err == nil {
if err := json.Unmarshal(bs, &ts); err != nil {
return nil, xerrors.Errorf("Failed to Unmarshal JSON. err: %w", err)
}
} else {
if !xerrors.Is(err, redis.Nil) {
return nil, xerrors.Errorf("Failed to Get Titles. err: %w", err)
}
ts, err = r.conn.SMembers(ctx, titleListKey).Result()
if err != nil {
return nil, xerrors.Errorf("Failed to SMembers Titles. err: %w", err)
}
}
default:
return nil, xerrors.Errorf("unexpected CPE#Titles type. expected: %q, actual: %q", []string{"string", "set"}, t)
}

if len(ts) < n {
Expand Down Expand Up @@ -310,7 +327,6 @@ func (r *RedisDriver) InsertCpes(fetchType models.FetchType, cpes models.Fetched
"VendorProducts": {},
"DeprecatedVendorProducts": {},
"DeprecatedCPEs": {},
"Titles": {},
"Title": {},
}
oldDepsStr, err := r.conn.HGet(ctx, depKey, string(fetchType)).Result()
Expand All @@ -323,7 +339,6 @@ func (r *RedisDriver) InsertCpes(fetchType models.FetchType, cpes models.Fetched
"VendorProducts": {},
"DeprecatedVendorProducts": {},
"DeprecatedCPEs": {},
"Titles": {},
"Title": {}
}`
}
Expand All @@ -332,6 +347,50 @@ func (r *RedisDriver) InsertCpes(fetchType models.FetchType, cpes models.Fetched
return xerrors.Errorf("Failed to unmarshal JSON. err: %w", err)
}

titles := make(map[string]struct{})
exists, err := r.conn.Exists(ctx, titleListKey).Result()
if err != nil {
return xerrors.Errorf("Failed to Exists CPE#Titles. err: %w", err)
}
if exists > 0 {
t, err := r.conn.Type(ctx, titleListKey).Result()
if err != nil {
return xerrors.Errorf("Failed to TYPE CPE#Titles. err: %w", err)
}

var ts []string
switch t {
case "string":
bs, err := r.conn.Get(ctx, titleListKey).Bytes()
if err != nil {
return xerrors.Errorf("Failed to Get Titles. err: %w", err)
}
if err := json.Unmarshal(bs, &ts); err != nil {
return xerrors.Errorf("Failed to Unmarshal JSON. err: %w", err)
}
case "set": // backward compatibility: https://github.com/vulsio/go-cpe-dictionary/pull/186
bs, err := r.conn.Get(ctx, titleListCacheKey).Bytes()
if err == nil {
if err := json.Unmarshal(bs, &ts); err != nil {
return xerrors.Errorf("Failed to Unmarshal JSON. err: %w", err)
}
} else {
if !xerrors.Is(err, redis.Nil) {
return xerrors.Errorf("Failed to Get Titles. err: %w", err)
}
ts, err = r.conn.SMembers(ctx, titleListKey).Result()
if err != nil {
return xerrors.Errorf("Failed to SMembers Titles. err: %w", err)
}
}
default:
return xerrors.Errorf("unexpected CPE#Titles type. expected: %q, actual: %q", []string{"string", "set"}, t)
}
for _, t := range ts {
titles[t] = struct{}{}
}
}

bar := pb.StartNew(len(cpes.CPEs) + len(cpes.Deprecated) + 1).SetWriter(func() io.Writer {
if viper.GetBool("log-json") {
return io.Discard
Expand Down Expand Up @@ -368,6 +427,7 @@ func (r *RedisDriver) InsertCpes(fetchType models.FetchType, cpes models.Fetched
newDeps["VendorProducts"][vendorProductStr] = map[string]struct{}{}
delete(oldDeps["VendorProducts"], vendorProductStr)
}

_ = pipe.SAdd(ctx, fmt.Sprintf(vpKeyFormat, c.Vendor, c.Product), c.CpeURI)
if _, ok := newDeps["VP"][vendorProductStr]; !ok {
newDeps["VP"][vendorProductStr] = map[string]struct{}{}
Expand All @@ -379,9 +439,8 @@ func (r *RedisDriver) InsertCpes(fetchType models.FetchType, cpes models.Fetched
delete(oldDeps["VP"], vendorProductStr)
}
}
_ = pipe.SAdd(ctx, titleListKey, c.Title)
newDeps["Titles"][c.Title] = map[string]struct{}{}
delete(oldDeps["Titles"], c.Title)

titles[c.Title] = struct{}{}

_ = pipe.SAdd(ctx, fmt.Sprintf(titleKeyFormat, c.Title), c.CpeURI)
if _, ok := newDeps["Title"][c.Title]; !ok {
Expand All @@ -402,16 +461,12 @@ func (r *RedisDriver) InsertCpes(fetchType models.FetchType, cpes models.Fetched
}
}

ts, err := r.conn.SMembers(ctx, titleListKey).Result()
if err != nil {
return xerrors.Errorf("Failed to SMembers Titles. err: %w", err)
}
bs, err := json.Marshal(ts)
bs, err := json.Marshal(slices.Collect(maps.Keys(titles)))
if err != nil {
return xerrors.Errorf("Failed to Marshal JSON. err: %w", err)
}
if err := r.conn.Set(ctx, titleListCacheKey, string(bs), 0).Err(); err != nil {
return xerrors.Errorf("Failed to SET Titles Cache. err: %w", err)
if err := r.conn.Set(ctx, titleListKey, string(bs), 0).Err(); err != nil {
return xerrors.Errorf("Failed to SET Titles. err: %w", err)
}
bar.Increment()

Expand All @@ -433,9 +488,7 @@ func (r *RedisDriver) InsertCpes(fetchType models.FetchType, cpes models.Fetched
for cpeURI := range oldDeps["DeprecatedCPEs"] {
_ = pipe.SRem(ctx, deprecatedCPEsKey, cpeURI)
}
for title := range oldDeps["Titles"] {
_ = pipe.SRem(ctx, titleListKey, title)
}
_ = pipe.Del(ctx, titleListCacheKey)
for title, cpeURIs := range oldDeps["Title"] {
for cpeURI := range cpeURIs {
_ = pipe.SRem(ctx, fmt.Sprintf(titleKeyFormat, title), cpeURI)
Expand Down

0 comments on commit 515af55

Please sign in to comment.