Skip to content

Commit ef810b2

Browse files
LostAttractorLostAttractor
committed
feat: allow mapping multiple domain names to single ip
1 parent 47b8fed commit ef810b2

File tree

3 files changed

+157
-35
lines changed

3 files changed

+157
-35
lines changed

‎control/control_plane.go

+10-9
Original file line numberDiff line numberDiff line change
@@ -432,28 +432,29 @@ func NewControlPlane(
432432
if plane.dnsController, err = NewDnsController(dnsUpstream, &DnsControllerOption{
433433
Log: log,
434434
CacheAccessCallback: func(cache *DnsCache) (err error) {
435-
// Write mappings into eBPF map:
436-
// IP record (from dns lookup) -> domain routing
437-
if err = core.BatchUpdateDomainRouting(cache); err != nil {
438-
return fmt.Errorf("BatchUpdateDomainRouting: %w", err)
439-
}
440435
return nil
441436
},
442437
CacheRemoveCallback: func(cache *DnsCache) (err error) {
443438
// Write mappings into eBPF map:
444439
// IP record (from dns lookup) -> domain routing
445-
if err = core.BatchRemoveDomainRouting(cache); err != nil {
446-
return fmt.Errorf("BatchUpdateDomainRouting: %w", err)
440+
if err = core.BatchRemoveDomain(cache); err != nil {
441+
return fmt.Errorf("BatchRemoveDomain: %w", err)
447442
}
448443
return nil
449444
},
450445
NewCache: func(fqdn string, answers []dnsmessage.RR, deadline time.Time, originalDeadline time.Time) (cache *DnsCache, err error) {
451-
return &DnsCache{
446+
cache = &DnsCache{
452447
DomainBitmap: plane.routingMatcher.domainMatcher.MatchDomainBitmap(fqdn),
453448
Answer: answers,
454449
Deadline: deadline,
455450
OriginalDeadline: originalDeadline,
456-
}, nil
451+
}
452+
// Write mappings into eBPF map:
453+
// IP record (from dns lookup) -> domain routing
454+
if err = core.BatchNewDomain(cache); err != nil {
455+
return cache, fmt.Errorf("BatchNewDomain: %w", err)
456+
}
457+
return cache, nil
457458
},
458459
BestDialerChooser: plane.chooseBestDnsDialer,
459460
TimeoutExceedCallback: func(dialArgument *dialArgument, err error) {

‎control/control_plane_core.go

+121-24
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,10 @@ type controlPlaneCore struct {
4545
isReload bool
4646
bpfEjected bool
4747

48+
domainBumpMap map[netip.Addr][]uint32
49+
domainRoutingMap map[netip.Addr][]uint32
50+
bumpMapMu sync.Mutex
51+
4852
closed context.Context
4953
close context.CancelFunc
5054
}
@@ -64,16 +68,18 @@ func newControlPlaneCore(log *logrus.Logger,
6468
}
6569
closed, toClose := context.WithCancel(context.Background())
6670
return &controlPlaneCore{
67-
log: log,
68-
deferFuncs: deferFuncs,
69-
bpf: bpf,
70-
outboundId2Name: outboundId2Name,
71-
kernelVersion: kernelVersion,
72-
flip: coreFlip,
73-
isReload: isReload,
74-
bpfEjected: false,
75-
closed: closed,
76-
close: toClose,
71+
log: log,
72+
deferFuncs: deferFuncs,
73+
bpf: bpf,
74+
outboundId2Name: outboundId2Name,
75+
kernelVersion: kernelVersion,
76+
flip: coreFlip,
77+
isReload: isReload,
78+
bpfEjected: false,
79+
domainBumpMap: make(map[netip.Addr][]uint32),
80+
domainRoutingMap: make(map[netip.Addr][]uint32),
81+
closed: closed,
82+
close: toClose,
7783
}
7884
}
7985

@@ -618,9 +624,8 @@ func (c *controlPlaneCore) bindDaens() (err error) {
618624
return
619625
}
620626

621-
// BatchUpdateDomainRouting update bpf map domain_routing. Since one IP may have multiple domains, this function should
622-
// be invoked every A/AAAA-record lookup.
623-
func (c *controlPlaneCore) BatchUpdateDomainRouting(cache *DnsCache) error {
627+
// BatchNewDomain update bpf map domain_bump and domain_routing. This function should be invoked every new cache.
628+
func (c *controlPlaneCore) BatchNewDomain(cache *DnsCache) error {
624629
// Parse ips from DNS resp answers.
625630
var ips []netip.Addr
626631
for _, ans := range cache.Answer {
@@ -646,27 +651,70 @@ func (c *controlPlaneCore) BatchUpdateDomainRouting(cache *DnsCache) error {
646651
// Update bpf map.
647652
// Construct keys and vals, and BpfMapBatchUpdate.
648653
var keys [][4]uint32
649-
var vals []bpfDomainRouting
654+
var vals_bump []bpfDomainRouting
655+
var vals_routing []bpfDomainRouting
656+
657+
c.bumpMapMu.Lock()
658+
defer c.bumpMapMu.Unlock()
659+
650660
for _, ip := range ips {
651661
ip6 := ip.As16()
652662
keys = append(keys, common.Ipv6ByteSliceToUint32Array(ip6[:]))
663+
653664
r := bpfDomainRouting{}
654-
if len(cache.DomainBitmap) != len(r.Bitmap) {
665+
666+
if consts.MaxMatchSetLen/32 != len(r.Bitmap) || len(cache.DomainBitmap) != len(r.Bitmap) {
655667
return fmt.Errorf("domain bitmap length not sync with kern program")
656668
}
657-
copy(r.Bitmap[:], cache.DomainBitmap)
658-
vals = append(vals, r)
669+
670+
newBumpMap, exists := c.domainBumpMap[ip]
671+
if !exists {
672+
newBumpMap = make([]uint32, consts.MaxMatchSetLen)
673+
}
674+
for index := 0; index < consts.MaxMatchSetLen; index++ {
675+
newBumpMap[index] += cache.DomainBitmap[index/32] >> (index % 32) & 1
676+
}
677+
for index, val := range newBumpMap {
678+
if val > 0 {
679+
r.Bitmap[index/32] |= 1 << (index % 32)
680+
}
681+
}
682+
c.domainBumpMap[ip] = newBumpMap
683+
vals_bump = append(vals_bump, r)
684+
685+
if !exists {
686+
// New IP, init routingMap
687+
c.domainRoutingMap[ip] = cache.DomainBitmap
688+
} else {
689+
// Old IP, Update routingMap
690+
for index := 0; index < consts.MaxMatchSetLen; index++ {
691+
if (cache.DomainBitmap[index/32]>>(index%32)&1) == 1 && (c.domainRoutingMap[ip][index/32]>>(index%32)&1) == 1 {
692+
// If this domain matches the current rule, all previous domains also match the current rule, then it still matches
693+
c.domainRoutingMap[ip][index/32] |= 1 << (index % 32)
694+
} else {
695+
// Otherwise, it does not match
696+
c.domainRoutingMap[ip][index/32] &^= 1 << (index % 32)
697+
}
698+
}
699+
}
700+
copy(r.Bitmap[:], c.domainRoutingMap[ip])
701+
vals_routing = append(vals_routing, r)
702+
}
703+
if _, err := BpfMapBatchUpdate(c.bpf.DomainBumpMap, keys, vals_bump, &ebpf.BatchOptions{
704+
ElemFlags: uint64(ebpf.UpdateAny),
705+
}); err != nil {
706+
return err
659707
}
660-
if _, err := BpfMapBatchUpdate(c.bpf.DomainRoutingMap, keys, vals, &ebpf.BatchOptions{
708+
if _, err := BpfMapBatchUpdate(c.bpf.DomainRoutingMap, keys, vals_routing, &ebpf.BatchOptions{
661709
ElemFlags: uint64(ebpf.UpdateAny),
662710
}); err != nil {
663711
return err
664712
}
665713
return nil
666714
}
667715

668-
// BatchRemoveDomainRouting remove bpf map domain_routing.
669-
func (c *controlPlaneCore) BatchRemoveDomainRouting(cache *DnsCache) error {
716+
// BatchRemoveDomainBump update or remove bpf map domain_bump and domain_routing.
717+
func (c *controlPlaneCore) BatchRemoveDomain(cache *DnsCache) error {
670718
// Parse ips from DNS resp answers.
671719
var ips []netip.Addr
672720
for _, ans := range cache.Answer {
@@ -690,15 +738,64 @@ func (c *controlPlaneCore) BatchRemoveDomainRouting(cache *DnsCache) error {
690738
}
691739

692740
// Update bpf map.
693-
// Construct keys and vals, and BpfMapBatchUpdate.
694-
var keys [][4]uint32
741+
// Update and determine whether to delete
742+
var keys_del [][4]uint32
743+
var keys_modify [][4]uint32
744+
var vals_modify_bump []bpfDomainRouting
745+
var vals_modify_routing []bpfDomainRouting
746+
747+
c.bumpMapMu.Lock()
748+
defer c.bumpMapMu.Unlock()
749+
695750
for _, ip := range ips {
696751
ip6 := ip.As16()
697-
keys = append(keys, common.Ipv6ByteSliceToUint32Array(ip6[:]))
752+
newBumpMapVal := c.domainBumpMap[ip]
753+
for index := 0; index < consts.MaxMatchSetLen; index++ {
754+
newBumpMapVal[index] -= cache.DomainBitmap[index/32] >> (index % 32) & 1
755+
}
756+
757+
bumpMap := bpfDomainRouting{}
758+
routingMap := bpfDomainRouting{}
759+
copy(routingMap.Bitmap[:], c.domainRoutingMap[ip])
760+
761+
del := true
762+
for index, val := range newBumpMapVal {
763+
if val > 0 {
764+
del = false // This IP refers to some domain name that matches the domain_set, so there is no need to delete
765+
bumpMap.Bitmap[index/32] |= 1 << (index % 32)
766+
} else {
767+
// This IP no longer refers to any domain name that matches the domain_set
768+
routingMap.Bitmap[index/32] &^= 1 << (index % 32)
769+
}
770+
}
771+
if del {
772+
delete(c.domainBumpMap, ip)
773+
delete(c.domainRoutingMap, ip)
774+
keys_del = append(keys_del, common.Ipv6ByteSliceToUint32Array(ip6[:]))
775+
} else {
776+
c.domainBumpMap[ip] = newBumpMapVal
777+
keys_modify = append(keys_modify, common.Ipv6ByteSliceToUint32Array(ip6[:]))
778+
vals_modify_bump = append(vals_modify_bump, bumpMap)
779+
vals_modify_routing = append(vals_modify_routing, routingMap)
780+
}
781+
}
782+
if _, err := BpfMapBatchDelete(c.bpf.DomainBumpMap, keys_del); err != nil {
783+
return err
784+
}
785+
if _, err := BpfMapBatchDelete(c.bpf.DomainRoutingMap, keys_del); err != nil {
786+
return err
787+
}
788+
if _, err := BpfMapBatchUpdate(c.bpf.DomainBumpMap, keys_modify, vals_modify_bump, &ebpf.BatchOptions{
789+
ElemFlags: uint64(ebpf.UpdateAny),
790+
}); err != nil {
791+
return err
698792
}
699-
if _, err := BpfMapBatchDelete(c.bpf.DomainRoutingMap, keys); err != nil {
793+
if _, err := BpfMapBatchUpdate(c.bpf.DomainRoutingMap, keys_modify, vals_modify_routing, &ebpf.BatchOptions{
794+
ElemFlags: uint64(ebpf.UpdateAny),
795+
}); err != nil {
700796
return err
701797
}
798+
702799
return nil
703800
}
704801

‎control/kern/tproxy.c

+26-2
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,15 @@ struct {
342342
// __uint(pinning, LIBBPF_PIN_BY_NAME);
343343
} domain_routing_map SEC(".maps");
344344

345+
struct {
346+
__uint(type, BPF_MAP_TYPE_LRU_HASH);
347+
__type(key, __be32[4]);
348+
__type(value, struct domain_routing);
349+
__uint(max_entries, MAX_DOMAIN_ROUTING_NUM);
350+
/// NOTICE: No persistence.
351+
// __uint(pinning, LIBBPF_PIN_BY_NAME);
352+
} domain_bump_map SEC(".maps");
353+
345354
struct ip_port_proto {
346355
__u32 ip[4];
347356
__be16 port;
@@ -647,6 +656,8 @@ static int route_loop_cb(__u32 index, void *data)
647656
// proxy Subrule is like: domain(suffix:baidu.com, suffix:google.com) Match
648657
// set is like: suffix:baidu.com
649658
struct domain_routing *domain_routing;
659+
struct domain_routing *domain_bump;
660+
bool need_control_plane_routing = false;
650661

651662
if (unlikely(index / 32 >= MAX_MATCH_SET_LEN / 32)) {
652663
ctx->result = -EFAULT;
@@ -753,8 +764,21 @@ static int route_loop_cb(__u32 index, void *data)
753764

754765
// We use key instead of k to pass checker.
755766
if (domain_routing &&
756-
(domain_routing->bitmap[index / 32] >> (index % 32)) & 1)
767+
(domain_routing->bitmap[index / 32] >> (index % 32)) & 1) {
768+
// All domains mapeed by the current IP address are matched.
757769
ctx->goodsubrule = true;
770+
} else {
771+
// Get domain bump bitmap.
772+
domain_bump = bpf_map_lookup_elem(&domain_bump_map,
773+
ctx->params->daddr);
774+
if (domain_bump &&
775+
(domain_bump->bitmap[index / 32] >> (index % 32)) & 1) {
776+
ctx->goodsubrule = true;
777+
// The current IP has mapped domains that match this rule, but not all of them do.
778+
// jump to control plane.
779+
need_control_plane_routing = true;
780+
}
781+
}
758782
break;
759783
case MatchType_ProcessName:
760784
#ifdef __DEBUG_ROUTING
@@ -830,7 +854,7 @@ static int route_loop_cb(__u32 index, void *data)
830854
} else {
831855
bool must = ctx->must || match_set->must;
832856

833-
if (!must && ctx->isdns) {
857+
if ((!must && ctx->isdns) || need_control_plane_routing) {
834858
ctx->result =
835859
(__s64)OUTBOUND_CONTROL_PLANE_ROUTING |
836860
((__s64)match_set->mark << 8) |

0 commit comments

Comments
 (0)