github.com/google/osv-scalibr@v0.4.1/clients/datasource/insights_cache.go (about)

     1  // Copyright 2025 Google LLC
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package datasource
    16  
    17  import (
    18  	"time"
    19  
    20  	pb "deps.dev/api/v3"
    21  	"google.golang.org/protobuf/proto"
    22  )
    23  
    24  type depsdevAPICache struct {
    25  	Timestamp         *time.Time
    26  	PackageCache      map[packageKey][]byte
    27  	VersionCache      map[versionKey][]byte
    28  	RequirementsCache map[versionKey][]byte
    29  }
    30  
    31  func protoMarshalCache[K comparable, V proto.Message](protoMap map[K]V) (map[K][]byte, error) {
    32  	byteMap := make(map[K][]byte)
    33  	for k, v := range protoMap {
    34  		b, err := proto.Marshal(v)
    35  		if err != nil {
    36  			return nil, err
    37  		}
    38  		byteMap[k] = b
    39  	}
    40  
    41  	return byteMap, nil
    42  }
    43  
    44  func protoUnmarshalCache[K comparable, V any, PV interface {
    45  	proto.Message
    46  	*V
    47  }](byteMap map[K][]byte, protoMap *map[K]PV) error {
    48  	*protoMap = make(map[K]PV)
    49  	for k, b := range byteMap {
    50  		v := PV(new(V))
    51  		if err := proto.Unmarshal(b, v); err != nil {
    52  			return err
    53  		}
    54  		(*protoMap)[k] = v
    55  	}
    56  
    57  	return nil
    58  }
    59  
    60  // GobEncode encodes cache to bytes.
    61  func (c *CachedInsightsClient) GobEncode() ([]byte, error) {
    62  	var cache depsdevAPICache
    63  	c.mu.Lock()
    64  	defer c.mu.Unlock()
    65  
    66  	if c.cacheTimestamp == nil {
    67  		now := time.Now().UTC()
    68  		c.cacheTimestamp = &now
    69  	}
    70  
    71  	cache.Timestamp = c.cacheTimestamp
    72  	var err error
    73  	cache.PackageCache, err = protoMarshalCache(c.packageCache.GetMap())
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  	cache.VersionCache, err = protoMarshalCache(c.versionCache.GetMap())
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  	cache.RequirementsCache, err = protoMarshalCache(c.requirementsCache.GetMap())
    82  	if err != nil {
    83  		return nil, err
    84  	}
    85  
    86  	return gobMarshal(cache)
    87  }
    88  
    89  // GobDecode decodes bytes to cache.
    90  func (c *CachedInsightsClient) GobDecode(b []byte) error {
    91  	var cache depsdevAPICache
    92  	if err := gobUnmarshal(b, &cache); err != nil {
    93  		return err
    94  	}
    95  
    96  	if cache.Timestamp != nil && time.Since(*cache.Timestamp) >= cacheExpiry {
    97  		// Cache expired
    98  		return nil
    99  	}
   100  
   101  	c.mu.Lock()
   102  	defer c.mu.Unlock()
   103  
   104  	c.cacheTimestamp = cache.Timestamp
   105  
   106  	var pkgMap map[packageKey]*pb.Package
   107  	if err := protoUnmarshalCache(cache.PackageCache, &pkgMap); err != nil {
   108  		return err
   109  	}
   110  
   111  	var verMap map[versionKey]*pb.Version
   112  	if err := protoUnmarshalCache(cache.VersionCache, &verMap); err != nil {
   113  		return err
   114  	}
   115  
   116  	var reqMap map[versionKey]*pb.Requirements
   117  	if err := protoUnmarshalCache(cache.RequirementsCache, &reqMap); err != nil {
   118  		return err
   119  	}
   120  
   121  	c.packageCache.SetMap(pkgMap)
   122  	c.versionCache.SetMap(verMap)
   123  	c.requirementsCache.SetMap(reqMap)
   124  
   125  	return nil
   126  }