container-census/internal/vulnerability/cache.go

package vulnerability

import (
	"fmt"
	"sync"
	"time"
)

// Cache manages the vulnerability scan cache
type Cache struct {
	mu      sync.RWMutex
	scans   map[string]*VulnerabilityScan // imageID -> scan
	ttl     time.Duration
	storage VulnerabilityStorage
}

// VulnerabilityStorage is the interface for database operations
type VulnerabilityStorage interface {
	GetVulnerabilityScan(imageID string) (*VulnerabilityScan, error)
	SaveVulnerabilityScan(scan *VulnerabilityScan, vulnerabilities []Vulnerability) error
}

// NewCache creates a new scan cache
func NewCache(storage VulnerabilityStorage, ttl time.Duration) *Cache {
	return &Cache{
		scans:   make(map[string]*VulnerabilityScan),
		ttl:     ttl,
		storage: storage,
	}
}

// Get retrieves a scan from cache or database
func (c *Cache) Get(imageID string) (*VulnerabilityScan, error) {
	// Check memory cache first
	c.mu.RLock()
	scan, exists := c.scans[imageID]
	c.mu.RUnlock()

	if exists {
		// Check if cached scan is still valid
		if time.Since(scan.ScannedAt) < c.ttl {
			return scan, nil
		}
		// Cache expired, remove from memory
		c.mu.Lock()
		delete(c.scans, imageID)
		c.mu.Unlock()
	}

	// Try to load from database
	scan, err := c.storage.GetVulnerabilityScan(imageID)
	if err != nil {
		return nil, err
	}

	// Check if database scan is still valid
	if scan != nil && time.Since(scan.ScannedAt) < c.ttl {
		// Add to memory cache
		c.mu.Lock()
		c.scans[imageID] = scan
		c.mu.Unlock()
		return scan, nil
	}

	// No valid scan found
	return nil, fmt.Errorf("no valid scan found for image %s", imageID)
}

// Set stores a scan in cache and database
func (c *Cache) Set(scan *VulnerabilityScan, vulnerabilities []Vulnerability) error {
	// Save to database first
	err := c.storage.SaveVulnerabilityScan(scan, vulnerabilities)
	if err != nil {
		return fmt.Errorf("failed to save scan to database: %w", err)
	}

	// Update memory cache
	c.mu.Lock()
	c.scans[scan.ImageID] = scan
	c.mu.Unlock()

	return nil
}

// IsValid checks if a scan exists and is still valid
func (c *Cache) IsValid(imageID string) bool {
	scan, err := c.Get(imageID)
	if err != nil || scan == nil {
		return false
	}
	return time.Since(scan.ScannedAt) < c.ttl
}

// NeedsScan determines if an image needs to be scanned
func (c *Cache) NeedsScan(imageID string) bool {
	return !c.IsValid(imageID)
}

// Invalidate removes a scan from the cache
func (c *Cache) Invalidate(imageID string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.scans, imageID)
}

// Clear removes all scans from memory cache
func (c *Cache) Clear() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.scans = make(map[string]*VulnerabilityScan)
}

// Size returns the number of scans in memory cache
func (c *Cache) Size() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.scans)
}

// SetTTL updates the cache TTL
func (c *Cache) SetTTL(ttl time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.ttl = ttl
}

// GetTTL returns the current cache TTL
func (c *Cache) GetTTL() time.Duration {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.ttl
}

// PruneExpired removes expired scans from memory cache
func (c *Cache) PruneExpired() int {
	c.mu.Lock()
	defer c.mu.Unlock()

	removed := 0
	for imageID, scan := range c.scans {
		if time.Since(scan.ScannedAt) >= c.ttl {
			delete(c.scans, imageID)
			removed++
		}
	}
	return removed
}