diff --git a/go.mod b/go.mod index 68d13ba1..4a0f47ed 100644 --- a/go.mod +++ b/go.mod @@ -119,7 +119,7 @@ require ( go.uber.org/multierr v1.11.0 // indirect golang.org/x/arch v0.8.0 // indirect golang.org/x/mod v0.34.0 // indirect - golang.org/x/sync v0.20.0 // indirect + golang.org/x/sync v0.20.0 golang.org/x/term v0.42.0 // indirect golang.org/x/tools v0.43.0 // indirect golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect diff --git a/internal/util/resolver/cache.go b/internal/util/resolver/cache.go index 95352562..07725116 100644 --- a/internal/util/resolver/cache.go +++ b/internal/util/resolver/cache.go @@ -15,9 +15,11 @@ const ( defaultTTL = 60 * time.Second ) +// CacheKey identifies a cached DNS response by question name, class, and type. type CacheKey string -// NewCacheKey generates resolver cache key from question of dns query. +// NewCacheKey generates a resolver cache key from a DNS question. +// Returns an empty key if q is nil. func NewCacheKey(q *dns.Question) CacheKey { if q == nil { return "" @@ -32,41 +34,62 @@ type cacheItem struct { ttl time.Duration } +// Cache stores DNS responses with TTL-based expiration and optional size bounding. +// A zero maxSize means unbounded (backward compatible). +// +// Load returns: +// - (nil, 0) on cache miss +// - (msg, ttl > 0) on fresh cache hit +// - (msg, ttl <= 0) on expired cache hit (for stale-while-revalidate) type Cache struct { - m sync.Map - logger logger.Logger + mu sync.RWMutex + entries map[CacheKey]*cacheItem + maxSize int // 0 = unlimited + logger logger.Logger } +// NewCache creates an unbounded DNS cache. func NewCache() *Cache { - return &Cache{} + return &Cache{entries: make(map[CacheKey]*cacheItem)} } +// WithLogger sets the logger for cache debug messages. func (c *Cache) WithLogger(logger logger.Logger) *Cache { c.logger = logger return c } +// WithMaxSize sets the maximum number of entries. 0 means unlimited. +// When the limit is reached, expired entries are evicted first, then the oldest. +func (c *Cache) WithMaxSize(n int) *Cache { + c.maxSize = n + return c +} + +// Load returns a cached DNS message and its remaining TTL. +// The returned message is a deep copy safe for caller mutation. func (c *Cache) Load(ctx context.Context, key CacheKey) (msg *dns.Msg, ttl time.Duration) { - v, ok := c.m.Load(key) + var elapsed time.Duration + c.mu.RLock() + item, ok := c.entries[key] + if ok { + msg = item.msg.Copy() + elapsed = time.Since(item.ts) + ttl = item.ttl - elapsed + } + c.mu.RUnlock() if !ok { return } - item, ok := v.(*cacheItem) - if !ok { - return - } - - msg = item.msg.Copy() for i := range msg.Answer { - d := uint32(time.Since(item.ts).Seconds()) + d := uint32(elapsed.Seconds()) if msg.Answer[i].Header().Ttl > d { msg.Answer[i].Header().Ttl -= d } else { msg.Answer[i].Header().Ttl = 1 } } - ttl = item.ttl - time.Since(item.ts) if log := c.logger; log.IsLevelEnabled(logger.DebugLevel) { if sid := ctxvalue.SidFromContext(ctx); sid != "" { @@ -80,6 +103,9 @@ func (c *Cache) Load(ctx context.Context, key CacheKey) (msg *dns.Msg, ttl time. return } +// Store adds a DNS response to the cache. Negative ttl means "do not cache" and is +// silently skipped. Zero ttl uses the minimum TTL from the answer records (or +// defaultTTL if none have a positive TTL). Positive ttl overrides all answer TTLs. func (c *Cache) Store(ctx context.Context, key CacheKey, mr *dns.Msg, ttl time.Duration) { if key == "" || mr == nil || ttl < 0 { return @@ -95,17 +121,38 @@ func (c *Cache) Store(ctx context.Context, key CacheKey, mr *dns.Msg, ttl time.D if ttl == 0 { ttl = defaultTTL } - } else { - for i := range mr.Answer { - mr.Answer[i].Header().Ttl = uint32(ttl.Seconds()) + } + + cp := mr.Copy() + if ttl > 0 { + for i := range cp.Answer { + cp.Answer[i].Header().Ttl = uint32(ttl.Seconds()) } } - c.m.Store(key, &cacheItem{ - msg: mr.Copy(), + c.mu.Lock() + // Evict if over capacity. + if c.maxSize > 0 { + c.cleanupLocked() + if len(c.entries) >= c.maxSize { + // Evict the oldest entry. + var oldestKey CacheKey + var oldestTS time.Time + for k, item := range c.entries { + if oldestKey == "" || item.ts.Before(oldestTS) { + oldestKey = k + oldestTS = item.ts + } + } + delete(c.entries, oldestKey) + } + } + c.entries[key] = &cacheItem{ + msg: cp, ts: time.Now(), ttl: ttl, - }) + } + c.mu.Unlock() if log := c.logger; log.IsLevelEnabled(logger.DebugLevel) { if sid := ctxvalue.SidFromContext(ctx); sid != "" { @@ -117,15 +164,22 @@ func (c *Cache) Store(ctx context.Context, key CacheKey, mr *dns.Msg, ttl time.D } } +// RefreshTTL resets the timestamp on a cached entry, extending its TTL. func (c *Cache) RefreshTTL(key CacheKey) { - v, ok := c.m.Load(key) - if !ok { - return + c.mu.Lock() + item, ok := c.entries[key] + if ok { + item.ts = time.Now() + } + c.mu.Unlock() +} + +// cleanupLocked removes expired entries. Caller must hold c.mu.Lock. +func (c *Cache) cleanupLocked() { + now := time.Now() + for key, item := range c.entries { + if now.Sub(item.ts) > item.ttl { + delete(c.entries, key) + } } - - item, ok := v.(*cacheItem) - if !ok { - return - } - item.ts = time.Now() } diff --git a/resolver/exchanger/exchanger.go b/resolver/exchanger/exchanger.go index 46cbe1f4..0baf9a1e 100644 --- a/resolver/exchanger/exchanger.go +++ b/resolver/exchanger/exchanger.go @@ -204,7 +204,7 @@ func (ex *exchanger) exchange(ctx context.Context, msg []byte) ([]byte, error) { } conn := &dns.Conn{ - UDPSize: 1024, + UDPSize: 4096, Conn: c, } diff --git a/resolver/resolver.go b/resolver/resolver.go index 2cc4ab4a..03f4c210 100644 --- a/resolver/resolver.go +++ b/resolver/resolver.go @@ -4,6 +4,7 @@ package resolver import ( "context" + "fmt" "net" "strings" "time" @@ -15,6 +16,7 @@ import ( resolver_util "github.com/go-gost/x/internal/util/resolver" "github.com/go-gost/x/resolver/exchanger" "github.com/miekg/dns" + "golang.org/x/sync/singleflight" ) // NameServer describes a DNS nameserver configuration including its address, @@ -54,11 +56,15 @@ func LoggerOption(logger logger.Logger) Option { } } +const maxAsyncRefresh = 32 // max concurrent background refresh goroutines + // localResolver resolves hostnames via DNS using configured nameservers and a cache. type localResolver struct { - servers []NameServer - cache *resolver_util.Cache - options options + servers []NameServer + cache *resolver_util.Cache + options options + refreshSem chan struct{} // bounds concurrent async refresh goroutines + sfGroup singleflight.Group } // NewResolver creates a Resolver that queries the given nameservers in order, @@ -94,8 +100,12 @@ func NewResolver(nameservers []NameServer, opts ...Option) (resolver.Resolver, e server.exchanger = ex switch server.Only { - case "ip4", "ipv4", "ip6", "ipv6": - server.Prefer = server.Only + case "ip4", "ipv4": + server.Only = "ipv4" + server.Prefer = "ipv4" + case "ip6", "ipv6": + server.Only = "ipv6" + server.Prefer = "ipv6" default: server.Only = "" } @@ -108,12 +118,26 @@ func NewResolver(nameservers []NameServer, opts ...Option) (resolver.Resolver, e WithLogger(options.logger) return &localResolver{ - servers: servers, - cache: cache, - options: options, + servers: servers, + cache: cache, + options: options, + refreshSem: make(chan struct{}, maxAsyncRefresh), }, nil } +// callerNetworkPref maps the caller's network hint to a DNS query preference. +// Returns "ipv4" for "ip4", "ipv6" for "ip6", or "" for no preference. +func callerNetworkPref(network string) string { + switch network { + case "ip4": + return "ipv4" + case "ip6": + return "ipv6" + default: + return "" + } +} + func (r *localResolver) Resolve(ctx context.Context, network, host string, opts ...resolver.Option) (ips []net.IP, err error) { if ip := net.ParseIP(host); ip != nil { return []net.IP{ip}, nil @@ -124,11 +148,13 @@ func (r *localResolver) Resolve(ctx context.Context, network, host string, opts host = host + "." + r.options.domain } + callerPref := callerNetworkPref(network) + for _, server := range r.servers { if server.Async { - ips, err = r.resolveAsync(ctx, &server, host) + ips, err = r.resolveAsync(ctx, &server, host, callerPref) } else { - ips, err = r.resolve(ctx, &server, host) + ips, err = r.resolve(ctx, &server, host, callerPref) } if err != nil { r.options.logger.Error(err) @@ -145,38 +171,63 @@ func (r *localResolver) Resolve(ctx context.Context, network, host string, opts return } -func (r *localResolver) resolve(ctx context.Context, server *NameServer, host string) (ips []net.IP, err error) { +// resolvePreference determines the effective IP preference. +// Precedence: server.Only (hard constraint) > callerPref > server.Prefer > default (ipv4-first). +func resolvePreference(server *NameServer, callerPref string) string { + if server.Only != "" { + return server.Only + } + if callerPref != "" { + return callerPref + } + if server.Prefer != "" { + return server.Prefer + } + return "" // default: ipv4-first +} + +func (r *localResolver) resolve(ctx context.Context, server *NameServer, host string, callerPref string) (ips []net.IP, err error) { if server == nil { return } - if server.Prefer == "ipv6" { // prefer ipv6 - if ips, err = r.resolve6(ctx, server, host); len(ips) > 0 || server.Only == "ipv6" { + pref := resolvePreference(server, callerPref) + + if pref == "ipv6" { + if ips, err = r.resolve6(ctx, server, host); len(ips) > 0 || server.Only == "ipv6" || callerPref == "ipv6" { return } return r.resolve4(ctx, server, host) } - if ips, err = r.resolve4(ctx, server, host); len(ips) > 0 || server.Only == "ipv4" { + if ips, err = r.resolve4(ctx, server, host); len(ips) > 0 || server.Only == "ipv4" || callerPref == "ipv4" { return } return r.resolve6(ctx, server, host) } -func (r *localResolver) resolveAsync(ctx context.Context, server *NameServer, host string) (ips []net.IP, err error) { - ips, ttl, ok := r.lookupCache(ctx, server, host) +func (r *localResolver) resolveAsync(ctx context.Context, server *NameServer, host string, callerPref string) (ips []net.IP, err error) { + ips, ttl, ok := r.lookupCache(ctx, server, host, callerPref) if !ok { - return r.resolve(ctx, server, host) + return r.resolve(ctx, server, host, callerPref) } if ttl <= 0 { r.options.logger.Debugf("async resolve %s via %s", host, server.exchanger.String()) - go r.resolve(context.WithoutCancel(ctx), server, host) + select { + case r.refreshSem <- struct{}{}: + go func() { + defer func() { <-r.refreshSem }() + r.resolve(context.WithoutCancel(ctx), server, host, callerPref) + }() + default: + r.options.logger.Debugf("async refresh skipped: semaphore full for %s", host) + } } return } -func (r *localResolver) lookupCache(ctx context.Context, server *NameServer, host string) (ips []net.IP, ttl time.Duration, ok bool) { +func (r *localResolver) lookupCache(ctx context.Context, server *NameServer, host string, callerPref string) (ips []net.IP, ttl time.Duration, ok bool) { lookup := func(t uint16, host string) (ips []net.IP, ttl time.Duration, ok bool) { mq := dns.Msg{} mq.SetQuestion(dns.Fqdn(host), t) @@ -184,32 +235,24 @@ func (r *localResolver) lookupCache(ctx context.Context, server *NameServer, hos if mr == nil { return } - ok = true - - for _, ans := range mr.Answer { - if ar, _ := ans.(*dns.AAAA); ar != nil { - ips = append(ips, ar.AAAA) - } - if ar, _ := ans.(*dns.A); ar != nil { - ips = append(ips, ar.A) - } - } + ips = extractIPs(mr) return } - if server.Prefer == "ipv6" { + pref := resolvePreference(server, callerPref) + + if pref == "ipv6" { ips, ttl, ok = lookup(dns.TypeAAAA, host) - if len(ips) > 0 || server.Only == "ipv6" { + if len(ips) > 0 || server.Only == "ipv6" || callerPref == "ipv6" { return } - ips, ttl, ok = lookup(dns.TypeA, host) return } ips, ttl, ok = lookup(dns.TypeA, host) - if len(ips) > 0 || server.Only == "ipv4" { + if len(ips) > 0 || server.Only == "ipv4" || callerPref == "ipv4" { return } return lookup(dns.TypeAAAA, host) @@ -233,30 +276,50 @@ func (r *localResolver) resolveIPs(ctx context.Context, server *NameServer, mq * } key := resolver_util.NewCacheKey(&mq.Question[0]) - mr, ttl := r.cache.Load(ctx, key) - if ttl <= 0 { + + result, err, _ := r.sfGroup.Do(string(key), func() (any, error) { + // Double-check cache after winning the singleflight race. + mr, ttl := r.cache.Load(ctx, key) + if ttl > 0 { + return mr, nil + } + + // Apply EDNS0 subnet option before exchange. mq is fresh per resolve4/resolve6 + // call, so mutation is safe — it will not be shared across singleflight callers. resolver_util.AddSubnetOpt(mq, server.ClientIP) - mr, err = r.exchange(ctx, server.exchanger, mq) + mr, err := r.exchange(ctx, server.exchanger, mq) if err != nil { - return + return nil, err } r.cache.Store(ctx, key, mr, server.TTL) if r.options.logger.IsLevelEnabled(logger.TraceLevel) { r.options.logger.Trace(mr.String()) } + return mr, nil + }) + if err != nil { + return nil, err } + return extractIPs(result.(*dns.Msg)), nil +} + +// extractIPs returns all A and AAAA records from a DNS message. +func extractIPs(mr *dns.Msg) []net.IP { + if mr == nil { + return nil + } + ips := make([]net.IP, 0, len(mr.Answer)) for _, ans := range mr.Answer { - if ar, _ := ans.(*dns.AAAA); ar != nil { - ips = append(ips, ar.AAAA) - } - if ar, _ := ans.(*dns.A); ar != nil { + if ar, ok := ans.(*dns.A); ok { ips = append(ips, ar.A) } + if ar, ok := ans.(*dns.AAAA); ok { + ips = append(ips, ar.AAAA) + } } - - return + return ips } func (r *localResolver) exchange(ctx context.Context, ex exchanger.Exchanger, mq *dns.Msg) (mr *dns.Msg, err error) { @@ -270,7 +333,13 @@ func (r *localResolver) exchange(ctx context.Context, ex exchanger.Exchanger, mq } mr = &dns.Msg{} - err = mr.Unpack(reply) + if err = mr.Unpack(reply); err != nil { + return + } + + if mr.Rcode != dns.RcodeSuccess { + return nil, fmt.Errorf("DNS %s for %s", dns.RcodeToString[mr.Rcode], mq.Question[0].Name) + } return } diff --git a/resolver/resolver_test.go b/resolver/resolver_test.go index e1238b7d..b57c295e 100644 --- a/resolver/resolver_test.go +++ b/resolver/resolver_test.go @@ -3,6 +3,7 @@ package resolver import ( "context" "net" + "sync" "sync/atomic" "testing" "time" @@ -79,9 +80,10 @@ func newTestResolver(servers []NameServer, opts ...Option) *localResolver { } cache := resolver_util.NewCache().WithLogger(xlogger.Nop()) return &localResolver{ - servers: servers, - cache: cache, - options: o, + servers: servers, + cache: cache, + options: o, + refreshSem: make(chan struct{}, maxAsyncRefresh), } } @@ -259,7 +261,7 @@ func TestResolve_CacheHit(t *testing.T) { func TestResolve_NilServer(t *testing.T) { r := &localResolver{options: nopLog()} - ips, err := r.resolve(context.Background(), nil, "example.com") + ips, err := r.resolve(context.Background(), nil, "example.com", "") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -417,3 +419,157 @@ func TestLoggerOption(t *testing.T) { t.Error("logger option not applied") } } + +func TestExchange_DNSRcodeError(t *testing.T) { + // Return a message with NXDOMAIN rcode + mr := new(dns.Msg) + mr.SetRcode(&dns.Msg{MsgHdr: dns.MsgHdr{Id: 1}}, dns.RcodeNameError) + packed, err := mr.Pack() + if err != nil { + t.Fatal(err) + } + + ex := &rcodeMockExchanger{response: packed} + r := newTestResolver([]NameServer{{exchanger: ex, TTL: time.Minute}}) + + ips, err := r.Resolve(context.Background(), "ip", "nonexistent.example.com") + if err == nil { + t.Fatal("expected error for NXDOMAIN") + } + if len(ips) != 0 { + t.Fatalf("expected no IPs for NXDOMAIN, got %v", ips) + } +} + +func TestExchange_DNSRcodeServfail_Fallback(t *testing.T) { + // First server returns SERVFAIL, second returns success + servfailMr := new(dns.Msg) + servfailMr.SetRcode(&dns.Msg{MsgHdr: dns.MsgHdr{Id: 1}}, dns.RcodeServerFailure) + packed, err := servfailMr.Pack() + if err != nil { + t.Fatal(err) + } + + okEx := &mockExchanger{response: newDNSMsgWithA("10.0.0.1")} + servfailEx := &rcodeMockExchanger{response: packed} + + r := newTestResolver([]NameServer{ + {exchanger: servfailEx, TTL: time.Minute}, + {exchanger: okEx, TTL: time.Minute}, + }) + + ips, err := r.Resolve(context.Background(), "ip", "example.com") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(ips) == 0 { + t.Fatal("expected fallback to second server") + } +} + +func TestResolve_NetworkIP4(t *testing.T) { + // When caller asks for ip4, only A queries should be made (no AAAA fallback). + // Use a mock that returns empty response for A queries. + ex := &mockExchanger{response: new(dns.Msg)} // empty response, no records + r := newTestResolver([]NameServer{{exchanger: ex, TTL: time.Minute}}) + + ips, err := r.Resolve(context.Background(), "ip4", "example.com") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Should get empty result (no A records), and NOT fall back to IPv6 + if len(ips) != 0 { + t.Errorf("expected no IPs when only ip4 requested and no A records, got %v", ips) + } +} + +func TestResolve_NetworkIP6(t *testing.T) { + ex := &mockExchanger{response: newDNSMsgWithAAAA("2001:db8::1")} + r := newTestResolver([]NameServer{{exchanger: ex, TTL: time.Minute}}) + + ips, err := r.Resolve(context.Background(), "ip6", "example.com") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(ips) == 0 { + t.Fatal("expected IPv6 result for ip6 network") + } +} + +func TestResolve_NetworkIP4_ServerOnlyOverrides(t *testing.T) { + // Server configured with Only="ipv6", caller asks for "ip4" → server wins + ex := &mockExchanger{response: newDNSMsgWithAAAA("2001:db8::1")} + r := newTestResolver([]NameServer{{exchanger: ex, Only: "ipv6", Prefer: "ipv6", TTL: time.Minute}}) + + ips, err := r.Resolve(context.Background(), "ip4", "example.com") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Server Only=ipv6 should override caller's ip4 request + if len(ips) == 0 { + t.Fatal("expected IPv6 result (server Only overrides caller)") + } + for _, ip := range ips { + if ip.To4() != nil { + t.Errorf("expected only IPv6 from server Only constraint, got %v", ip) + } + } +} + +func TestResolve_ConcurrentDedup(t *testing.T) { + ex := &mockExchanger{response: newDNSMsgWithA("10.0.0.1")} + r := newTestResolver([]NameServer{{exchanger: ex, TTL: time.Minute}}) + + const n = 10 + var wg sync.WaitGroup + results := make([][]net.IP, n) + errs := make([]error, n) + + for i := 0; i < n; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + results[idx], errs[idx] = r.Resolve(context.Background(), "ip", "example.com") + }(i) + } + wg.Wait() + + for i, err := range errs { + if err != nil { + t.Fatalf("goroutine %d: unexpected error: %v", i, err) + } + if len(results[i]) == 0 { + t.Fatalf("goroutine %d: expected IPs", i) + } + } + + // singleflight should have deduplicated the exchange calls + calls := ex.calls.Load() + if calls > 2 { // allow small margin for race, but not 10 + t.Errorf("expected ~1 exchange call from singleflight, got %d", calls) + } +} + +func TestResolve_IP_UnchangedBehavior(t *testing.T) { + ex := &mockExchanger{response: newDNSMsgWithA("10.0.0.1")} + r := newTestResolver([]NameServer{{exchanger: ex, TTL: time.Minute}}) + + ips, err := r.Resolve(context.Background(), "ip", "example.com") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(ips) == 0 { + t.Fatal("expected result for network=ip (default behavior)") + } +} + +// rcodeMockExchanger returns a pre-packed response, used for DNS rcode testing. +type rcodeMockExchanger struct { + response []byte +} + +func (m *rcodeMockExchanger) Exchange(_ context.Context, _ []byte) ([]byte, error) { + return m.response, nil +} + +func (m *rcodeMockExchanger) String() string { return "mock://rcode" }