Skip to content

Commit

Permalink
Add rejected DNS response cache support
Browse files Browse the repository at this point in the history
  • Loading branch information
nekohasekai committed Feb 14, 2024
1 parent 8dc88e9 commit 3974c99
Showing 1 changed file with 31 additions and 3 deletions.
34 changes: 31 additions & 3 deletions client.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,20 +19,28 @@ import (
const DefaultTTL = 600

var (
ErrNoRawSupport = E.New("no raw query support by current transport")
ErrNotCached = E.New("not cached")
ErrResponseRejected = E.New("response rejected")
ErrNoRawSupport = E.New("no raw query support by current transport")
ErrNotCached = E.New("not cached")
ErrResponseRejected = E.New("response rejected")
ErrResponseRejectedCached = E.Extend(ErrResponseRejected, "cached")
)

type Client struct {
disableCache bool
disableExpire bool
independentCache bool
rdrc RDRCStore
logger logger.ContextLogger
cache *cache.LruCache[dns.Question, *dns.Msg]
transportCache *cache.LruCache[transportCacheKey, *dns.Msg]
}

type RDRCStore interface {
LoadRDRC(transportName string, qName string) (rejected bool)
SaveRDRC(transportName string, qName string) error
SaveRDRCAsync(transportName string, qName string, logger logger.Logger)
}

type transportCacheKey struct {
dns.Question
transportName string
Expand All @@ -42,6 +50,7 @@ type ClientOptions struct {
DisableCache bool
DisableExpire bool
IndependentCache bool
RDRC RDRCStore
Logger logger.ContextLogger
}

Expand All @@ -50,6 +59,7 @@ func NewClient(options ClientOptions) *Client {
disableCache: options.DisableCache,
disableExpire: options.DisableExpire,
independentCache: options.IndependentCache,
rdrc: options.RDRC,
logger: options.Logger,
}
if !client.disableCache {
Expand Down Expand Up @@ -121,11 +131,20 @@ func (c *Client) ExchangeWithResponseCheck(ctx context.Context, transport Transp
if loaded {
SetClientSubnet(message, clientSubnet, true)
}
if responseChecker != nil && c.rdrc != nil {
rejected := c.rdrc.LoadRDRC(transport.Name(), question.Name)
if rejected {
return nil, ErrResponseRejectedCached
}
}
response, err := transport.Exchange(ctx, message)
if err != nil {
return nil, err
}
if responseChecker != nil && !responseChecker(response) {
if c.rdrc != nil {
c.rdrc.SaveRDRCAsync(transport.Name(), question.Name, c.logger)
}
return response, ErrResponseRejected
}
var timeToLive int
Expand Down Expand Up @@ -238,12 +257,21 @@ func (c *Client) LookupWithResponseCheck(ctx context.Context, transport Transpor
}
}
}
if responseChecker != nil && c.rdrc != nil {
rejected := c.rdrc.LoadRDRC(transport.Name(), dnsName)
if rejected {
return nil, ErrResponseRejectedCached
}
}
var rCode int
response, err := transport.Lookup(ctx, domain, strategy)
if err != nil {
return nil, wrapError(err)
}
if responseChecker != nil && !responseChecker(response) {
if c.rdrc != nil {
c.rdrc.SaveRDRCAsync(transport.Name(), dnsName, c.logger)
}
return response, ErrResponseRejected
}
header := dns.MsgHdr{
Expand Down

0 comments on commit 3974c99

Please sign in to comment.