AdGuardDNS/internal/dnsmsg/optcloner.go
Andrey Meshkov 87137bddcf Sync v2.10.0
2024-11-08 16:26:22 +03:00

120 lines
2.4 KiB
Go

package dnsmsg
import (
"net"
"github.com/AdguardTeam/golibs/syncutil"
"github.com/miekg/dns"
)
// optCloner is a pool that can clone common parts of DNS messages of type OPT
// with fewer allocations.
type optCloner struct {
// Top-level structures.
rr *syncutil.Pool[dns.OPT]
// Options.
cookie *syncutil.Pool[dns.EDNS0_COOKIE]
ede *syncutil.Pool[dns.EDNS0_EDE]
subnet *syncutil.Pool[dns.EDNS0_SUBNET]
}
// newOPTCloner returns a new properly initialized *optCloner.
func newOPTCloner() (c *optCloner) {
return &optCloner{
rr: syncutil.NewPool(func() (v *dns.OPT) {
return &dns.OPT{}
}),
cookie: syncutil.NewPool(func() (v *dns.EDNS0_COOKIE) {
return &dns.EDNS0_COOKIE{}
}),
ede: syncutil.NewPool(func() (v *dns.EDNS0_EDE) {
return &dns.EDNS0_EDE{}
}),
subnet: syncutil.NewPool(func() (v *dns.EDNS0_SUBNET) {
return &dns.EDNS0_SUBNET{
// Use the IPv6 length to increase the effectiveness of the
// pool.
Address: make(net.IP, net.IPv6len),
}
}),
}
}
// clone returns a deep clone of rr. full is true if rr was cloned entirely
// without the use of [dns.Copy].
func (c *optCloner) clone(rr *dns.OPT) (clone *dns.OPT, full bool) {
if rr == nil {
return nil, true
}
clone = c.rr.Get()
clone.Hdr = rr.Hdr
if rr.Option == nil {
clone.Option = nil
return clone, true
}
clone.Option = clone.Option[:0]
for _, orig := range rr.Option {
var optClone dns.EDNS0
switch orig := orig.(type) {
case *dns.EDNS0_COOKIE:
opt := c.cookie.Get()
*opt = *orig
optClone = opt
case *dns.EDNS0_EDE:
opt := c.ede.Get()
*opt = *orig
optClone = opt
case *dns.EDNS0_SUBNET:
opt := c.subnet.Get()
opt.Code = orig.Code
opt.Family = orig.Family
opt.SourceNetmask = orig.SourceNetmask
opt.SourceScope = orig.SourceScope
opt.Address = append(opt.Address[:0], orig.Address...)
optClone = opt
// TODO(a.garipov): Add more if necessary.
default:
return dns.Copy(rr).(*dns.OPT), false
}
clone.Option = append(clone.Option, optClone)
}
return clone, true
}
// put returns structures from rr into c's pools.
func (c *optCloner) put(rr *dns.OPT) {
if rr == nil {
return
}
for _, opt := range rr.Option {
switch opt := opt.(type) {
case *dns.EDNS0_COOKIE:
c.cookie.Put(opt)
case *dns.EDNS0_SUBNET:
c.subnet.Put(opt)
case *dns.EDNS0_EDE:
c.ede.Put(opt)
default:
// Go on.
}
}
c.rr.Put(rr)
}