god-eye/internal/store/memory.go

package store

import (
	"context"
	"sort"
	"sync"
	"time"
)

// MemoryStore is the default in-memory Store implementation. Thread-safe,
// suitable for single-process scans. Persistent backends (BoltDB for ASM /
// resume workflows) land in Fase 5; they will implement the same Store
// interface so callers need no changes.
type MemoryStore struct {
	mu    sync.RWMutex
	hosts map[string]*Host
	// perHostLocks serializes Upsert mutations per-host without blocking
	// independent hosts. It's populated lazily and never cleared — the number
	// of subdomains per scan is bounded (thousands, not millions).
	perHostLocks map[string]*sync.Mutex
	locksMu      sync.Mutex
}

// NewMemoryStore creates an empty MemoryStore.
func NewMemoryStore() *MemoryStore {
	return &MemoryStore{
		hosts:        make(map[string]*Host),
		perHostLocks: make(map[string]*sync.Mutex),
	}
}

// lockFor returns the mutex that protects mutations to subdomain, creating
// it lazily if needed.
func (s *MemoryStore) lockFor(subdomain string) *sync.Mutex {
	s.locksMu.Lock()
	defer s.locksMu.Unlock()
	l, ok := s.perHostLocks[subdomain]
	if !ok {
		l = &sync.Mutex{}
		s.perHostLocks[subdomain] = l
	}
	return l
}

// Upsert creates or updates the record for subdomain, invoking mutate under
// a per-host lock. Concurrent callers mutating different subdomains proceed
// in parallel; concurrent mutations of the same subdomain are serialized.
func (s *MemoryStore) Upsert(ctx context.Context, subdomain string, mutate func(*Host)) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	if subdomain == "" {
		return nil
	}

	hostLock := s.lockFor(subdomain)
	hostLock.Lock()
	defer hostLock.Unlock()

	s.mu.Lock()
	h, existed := s.hosts[subdomain]
	if !existed {
		h = &Host{
			Subdomain: subdomain,
			FirstSeen: time.Now(),
		}
		s.hosts[subdomain] = h
	}
	s.mu.Unlock()

	if mutate != nil {
		mutate(h)
	}
	h.LastUpdated = time.Now()
	return nil
}

// Get returns a deep-enough copy of the record so the caller cannot
// accidentally mutate store state. Slice fields are copied; nested struct
// pointers (TLSFingerprint, Takeover) are shallow-copied — callers MUST treat
// the result as read-only.
func (s *MemoryStore) Get(ctx context.Context, subdomain string) (*Host, bool) {
	s.mu.RLock()
	h, ok := s.hosts[subdomain]
	s.mu.RUnlock()
	if !ok {
		return nil, false
	}
	hostLock := s.lockFor(subdomain)
	hostLock.Lock()
	defer hostLock.Unlock()
	return cloneHost(h), true
}

// All returns every host, sorted by subdomain. Each returned Host is a copy;
// mutations to the slice or its elements do not affect the store.
func (s *MemoryStore) All(ctx context.Context) []*Host {
	s.mu.RLock()
	names := make([]string, 0, len(s.hosts))
	for name := range s.hosts {
		names = append(names, name)
	}
	s.mu.RUnlock()
	sort.Strings(names)

	out := make([]*Host, 0, len(names))
	for _, n := range names {
		if h, ok := s.Get(ctx, n); ok {
			out = append(out, h)
		}
	}
	return out
}

// Count returns the number of hosts in the store.
func (s *MemoryStore) Count(ctx context.Context) int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.hosts)
}

// Close is a no-op for MemoryStore; implemented to satisfy Store interface.
func (s *MemoryStore) Close() error { return nil }

// cloneHost returns a deep-enough copy that slice/map fields are detached.
func cloneHost(h *Host) *Host {
	if h == nil {
		return nil
	}
	c := *h
	c.IPs = cloneStrings(h.IPs)
	c.Technologies = cloneStrings(h.Technologies)
	c.TLSAltNames = cloneStrings(h.TLSAltNames)
	c.DiscoveredVia = cloneStrings(h.DiscoveredVia)
	c.Ports = cloneInts(h.Ports)
	c.Headers = cloneStringMap(h.Headers)

	if h.TLSFingerprint != nil {
		fp := *h.TLSFingerprint
		fp.InternalHosts = cloneStrings(h.TLSFingerprint.InternalHosts)
		c.TLSFingerprint = &fp
	}
	if h.Takeover != nil {
		t := *h.Takeover
		c.Takeover = &t
	}

	c.Vulnerabilities = cloneVulns(h.Vulnerabilities)
	c.Secrets = cloneSecrets(h.Secrets)
	c.CVEs = cloneCVEs(h.CVEs)
	c.AIFindings = cloneAIFindings(h.AIFindings)
	return &c
}

func cloneStrings(in []string) []string {
	if len(in) == 0 {
		return nil
	}
	out := make([]string, len(in))
	copy(out, in)
	return out
}

func cloneInts(in []int) []int {
	if len(in) == 0 {
		return nil
	}
	out := make([]int, len(in))
	copy(out, in)
	return out
}

func cloneStringMap(in map[string]string) map[string]string {
	if len(in) == 0 {
		return nil
	}
	out := make(map[string]string, len(in))
	for k, v := range in {
		out[k] = v
	}
	return out
}

func cloneVulns(in []Vulnerability) []Vulnerability {
	if len(in) == 0 {
		return nil
	}
	out := make([]Vulnerability, len(in))
	for i, v := range in {
		v.CVEs = cloneStrings(v.CVEs)
		out[i] = v
	}
	return out
}

func cloneSecrets(in []Secret) []Secret {
	if len(in) == 0 {
		return nil
	}
	out := make([]Secret, len(in))
	copy(out, in)
	return out
}

func cloneCVEs(in []CVE) []CVE {
	if len(in) == 0 {
		return nil
	}
	out := make([]CVE, len(in))
	copy(out, in)
	return out
}

func cloneAIFindings(in []AIFinding) []AIFinding {
	if len(in) == 0 {
		return nil
	}
	out := make([]AIFinding, len(in))
	for i, f := range in {
		f.CVEs = cloneStrings(f.CVEs)
		out[i] = f
	}
	return out
}

// AppendUnique helpers — exported for modules that want to append slice
// fields without introducing duplicates. Keeps mutation semantics in one place.

// AddDiscoveryMethod appends method to h.DiscoveredVia if not already present.
func AddDiscoveryMethod(h *Host, method string) {
	for _, m := range h.DiscoveredVia {
		if m == method {
			return
		}
	}
	h.DiscoveredVia = append(h.DiscoveredVia, method)
}

// AddIPs appends new IPs (dedup, in-place).
func AddIPs(h *Host, ips []string) {
	seen := make(map[string]bool, len(h.IPs))
	for _, ip := range h.IPs {
		seen[ip] = true
	}
	for _, ip := range ips {
		if ip == "" || seen[ip] {
			continue
		}
		seen[ip] = true
		h.IPs = append(h.IPs, ip)
	}
}

// AddTechnologies appends new technologies (dedup, in-place).
func AddTechnologies(h *Host, tech []string) {
	seen := make(map[string]bool, len(h.Technologies))
	for _, t := range h.Technologies {
		seen[t] = true
	}
	for _, t := range tech {
		if t == "" || seen[t] {
			continue
		}
		seen[t] = true
		h.Technologies = append(h.Technologies, t)
	}
}