fix(DNSWhoamiService): implement cache expiration

internal/webconnectivityalgo/dnswhoami.go

@@ -14,14 +14,21 @@ import (
 
 	"github.com/ooni/probe-cli/v3/internal/model"
 	"github.com/ooni/probe-cli/v3/internal/netxlite"
+	"github.com/ooni/probe-cli/v3/internal/optional"
 )
 
 // DNSWhoamiInfoEntry contains an entry for DNSWhoamiInfo.
 type DNSWhoamiInfoEntry struct {
-	// Address is the IP address
+	// Address is the IP address used by the resolver.
 	Address string `json:"address"`
 }
 
+// dnsWhoamiInfoTimedEntry keeps an address and the time we created the entry together.
+type dnsWhoamiInfoTimedEntry struct {
+	Addr string
+	T    time.Time
+}
+
 // TODO(bassosimone): this code needs refining before we can merge it inside
 // master. For one, we already have systemv4 info. Additionally, it would
 // be neat to avoid additional AAAA queries. Furthermore, we should also see
@@ -30,27 +37,25 @@ type DNSWhoamiInfoEntry struct {
 // TODO(bassosimone): consider factoring this code and keeping state
 // on disk rather than on memory.
 
-// TODO(bassosimone): we should periodically invalidate the whoami lookup results.
-
 // DNSWhoamiService is a service that performs DNS whoami lookups.
 //
 // The zero value of this struct is invalid. Please, construct using
 // the [NewDNSWhoamiService] factory function.
 type DNSWhoamiService struct {
-	// logger is the logger
+	// entries contains the entries.
+	entries map[string]*dnsWhoamiInfoTimedEntry
+
+	// logger is the logger.
 	logger model.Logger
 
-	// mu provides mutual exclusion
+	// mu provides mutual exclusion.
 	mu *sync.Mutex
 
-	// netx is the underlying network we're using
+	// netx is the underlying network we're using.
 	netx *netxlite.Netx
 
-	// systemv4 contains systemv4 results
-	systemv4 []DNSWhoamiInfoEntry
-
-	// udpv4 contains udpv4 results
-	udpv4 map[string][]DNSWhoamiInfoEntry
+	// timeNow allows to get the current time.
+	timeNow func() time.Time
 
 	// whoamiDomain is the whoamiDomain to query for.
 	whoamiDomain string
@@ -59,53 +64,115 @@ type DNSWhoamiService struct {
 // NewDNSWhoamiService constructs a new [*DNSWhoamiService].
 func NewDNSWhoamiService(logger model.Logger) *DNSWhoamiService {
 	return &DNSWhoamiService{
+		entries:      map[string]*dnsWhoamiInfoTimedEntry{},
 		logger:       logger,
 		mu:           &sync.Mutex{},
 		netx:         &netxlite.Netx{Underlying: nil},
-		systemv4:     []DNSWhoamiInfoEntry{},
-		udpv4:        map[string][]DNSWhoamiInfoEntry{},
+		timeNow:      time.Now,
 		whoamiDomain: "whoami.v4.powerdns.org",
 	}
 }
 
 // SystemV4 returns the results of querying using the system resolver and IPv4.
 func (svc *DNSWhoamiService) SystemV4(ctx context.Context) ([]DNSWhoamiInfoEntry, bool) {
-	svc.mu.Lock()
-	defer svc.mu.Unlock()
-	if len(svc.systemv4) <= 0 {
-		ctx, cancel := context.WithTimeout(ctx, 4*time.Second)
-		defer cancel()
-		reso := svc.netx.NewStdlibResolver(svc.logger)
-		addrs, err := reso.LookupHost(ctx, svc.whoamiDomain)
-		if err != nil || len(addrs) < 1 {
-			return nil, false
-		}
-		svc.systemv4 = []DNSWhoamiInfoEntry{{
-			Address: addrs[0],
-		}}
+	spec := &dnsWhoamiResolverSpec{
+		name: "system:///",
+		factory: func(logger model.Logger, netx *netxlite.Netx) model.Resolver {
+			return svc.netx.NewStdlibResolver(svc.logger)
+		},
 	}
-	return svc.systemv4, len(svc.systemv4) > 0
+	v := svc.lookup(ctx, spec)
+	return v, len(v) > 0
 }
 
 // UDPv4 returns the results of querying a given UDP resolver and IPv4.
 func (svc *DNSWhoamiService) UDPv4(ctx context.Context, address string) ([]DNSWhoamiInfoEntry, bool) {
+	spec := &dnsWhoamiResolverSpec{
+		name: address,
+		factory: func(logger model.Logger, netx *netxlite.Netx) model.Resolver {
+			dialer := svc.netx.NewDialerWithResolver(svc.logger, svc.netx.NewStdlibResolver(svc.logger))
+			return svc.netx.NewParallelUDPResolver(svc.logger, dialer, address)
+		},
+	}
+	v := svc.lookup(ctx, spec)
+	return v, len(v) > 0
+}
+
+type dnsWhoamiResolverSpec struct {
+	name    string
+	factory func(logger model.Logger, netx *netxlite.Netx) model.Resolver
+}
+
+func (svc *DNSWhoamiService) lookup(ctx context.Context, spec *dnsWhoamiResolverSpec) []DNSWhoamiInfoEntry {
+	// get the current time
+	now := svc.timeNow()
+
+	// possibly use cache
+	mentry := svc.lockAndGet(now, spec.name)
+	if !mentry.IsNone() {
+		return []DNSWhoamiInfoEntry{mentry.Unwrap()}
+	}
+
+	// perform lookup
+	ctx, cancel := context.WithTimeout(ctx, 4*time.Second)
+	defer cancel()
+	reso := spec.factory(svc.logger, svc.netx)
+	addrs, err := reso.LookupHost(ctx, svc.whoamiDomain)
+	if err != nil || len(addrs) < 1 {
+		return nil
+	}
+
+	// update cache
+	svc.lockAndUpdate(now, spec.name, addrs[0])
+
+	// return to the caller
+	return []DNSWhoamiInfoEntry{{Address: addrs[0]}}
+}
+
+func (svc *DNSWhoamiService) lockAndGet(now time.Time, serverAddr string) optional.Value[DNSWhoamiInfoEntry] {
+	// ensure there's mutual exclusion
+	defer svc.mu.Unlock()
+	svc.mu.Lock()
+
+	// see if there's an entry
+	entry, found := svc.entries[serverAddr]
+	if !found {
+		return optional.None[DNSWhoamiInfoEntry]()
+	}
+
+	// make sure the entry has not expired
+	const validity = 45 * time.Second
+	if now.Sub(entry.T) > validity {
+		return optional.None[DNSWhoamiInfoEntry]()
+	}
+
+	// return a copy of the value
+	return optional.Some(DNSWhoamiInfoEntry{
+		Address: entry.Addr,
+	})
+}
+
+func (svc *DNSWhoamiService) lockAndUpdate(now time.Time, serverAddr, whoamiAddr string) {
+	// ensure there's mutual exclusion
+	defer svc.mu.Unlock()
 	svc.mu.Lock()
+
+	// insert into the table
+	svc.entries[serverAddr] = &dnsWhoamiInfoTimedEntry{
+		Addr: whoamiAddr,
+		T:    now,
+	}
+}
+
+func (svc *DNSWhoamiService) cloneEntries() map[string]*dnsWhoamiInfoTimedEntry {
 	defer svc.mu.Unlock()
-	if len(svc.udpv4[address]) <= 0 {
-		ctx, cancel := context.WithTimeout(ctx, 4*time.Second)
-		defer cancel()
-		dialer := svc.netx.NewDialerWithResolver(svc.logger, svc.netx.NewStdlibResolver(svc.logger))
-		reso := svc.netx.NewParallelUDPResolver(svc.logger, dialer, address)
-		// TODO(bassosimone): this should actually only send an A query. Sending an AAAA
-		// query is _way_ unnecessary since we know that only A is going to work.
-		addrs, err := reso.LookupHost(ctx, svc.whoamiDomain)
-		if err != nil || len(addrs) < 1 {
-			return nil, false
+	svc.mu.Lock()
+	output := make(map[string]*dnsWhoamiInfoTimedEntry)
+	for key, value := range svc.entries {
+		output[key] = &dnsWhoamiInfoTimedEntry{
+			Addr: value.Addr,
+			T:    value.T,
 		}
-		svc.udpv4[address] = []DNSWhoamiInfoEntry{{
-			Address: addrs[0],
-		}}
 	}
-	value := svc.udpv4[address]
-	return value, len(value) > 0
+	return output
 }