feat: add simple downlink rebalancing to secondary UPF

cmd/config/config.go

@@ -113,6 +113,9 @@ type UpfConfig struct {
 	TraceMaxDumpPackets     int            `mapstructure:"trace_max_packets" json:"trace_max_packets"`
 	IP6RaSupport            bool           `mapstructure:"ip6_ra_support" json:"ip6_ra_support"`
 	Ip6RaPrefix             string         `mapstructure:"ip6_ra_prefix" validate:"cidrv6" json:"ip6_ra_prefix"`
+	Ip4UEPrefixes           []string       `mapstructure:"ip4_ue_prefixes" validate:"omitempty,dive,cidrv4" json:"ip4_ue_prefixes"`
+	DLRebalanceAddress      string         `mapstructure:"rebalance_address" validate:"omitempty,ipv4" json:"rebalance_address"`
+	DLRebalanceLocalAddress string         `mapstructure:"rebalance_local_address" validate:"omitempty,ipv4" json:"rebalance_local_address"`
 
 	n3AdvertisedAddress5G net.IP
 }
@@ -183,6 +186,9 @@ func defineFlags() {
 	pflag.Int("tracefpackets", 100000, "Maximum number of packets in one trace dump file")
 	pflag.Bool("ip6ra", false, "Enable or disable IPv6 Router Advertisement support")
 	pflag.String("ip6rapref", "2a03:d000:29a0:509::/64", "Subscriber IPv6 prefix")
+	pflag.StringArray("ip4prefs", []string{}, "Subscriber IPv4 prefixes")
+	pflag.String("rbaddr", "", "Downlink rebalace address")
+	pflag.String("rbladdr", "", "Downlink rebalace local address")
 
 	pflag.Parse()
 }
@@ -264,6 +270,9 @@ func initCommonConfig() {
 	_ = commonConfigV.BindPFlag("trace_max_packets", pflag.Lookup("tracefpackets"))
 	_ = commonConfigV.BindPFlag("ip6_ra_support", pflag.Lookup("ip6ra"))
 	_ = commonConfigV.BindPFlag("ip6_ra_prefix", pflag.Lookup("ip6rapref"))
+	_ = commonConfigV.BindPFlag("ip4_ue_prefixes", pflag.Lookup("ip4prefs"))
+	_ = commonConfigV.BindPFlag("rebalance_address", pflag.Lookup("rbaddr"))
+	_ = commonConfigV.BindPFlag("rebalance_local_address", pflag.Lookup("rbladdr"))
 
 	commonConfigV.SetDefault("n9_address", commonConfigV.GetString("n3_address"))
 	commonConfigV.SetDefault("n3_adv_address", commonConfigV.GetString("n3_address"))

cmd/ebpf/objects.go

@@ -86,8 +86,9 @@ func (bpfObjects *BpfObjects) Load() error {
 		"far_map":              bpfObjects.farMapSize,
 		"pdr_map_downlink_ip4": bpfObjects.pdrMapSize,
 		"pdr_map_downlink_ip6": bpfObjects.pdrMapSize,
-		"pdr_map_teid_ip4":   bpfObjects.pdrMapSize,
+		"pdr_map_teid_ip4":     bpfObjects.pdrMapSize,
 		"urr_map":              bpfObjects.urrMapSize,
+		"ue_rebalance_lpm_map": 16,
 	}
 
 	replacements := make(map[string]*ebpf.Map)
@@ -127,6 +128,15 @@ func (bpfObjects *BpfObjects) Load() error {
 			m.Close()
 		}
 
+		var ve *ebpf.VerifierError
+		if errors.As(err, &ve) {
+			for i, line := range ve.Log {
+				if line != "" {
+					log.Info().Msgf("[%03d] %s", i, line)
+				}
+			}
+		}
+
 		log.Warn().Msgf("Failed to load objects: %s", err)
 		return err
 	}

cmd/ebpf/objects_test.go

@@ -1134,6 +1134,113 @@ func TestEntrypoint(t *testing.T) {
 
 }
 
+var (
+	rebalanceIP      = net.ParseIP("10.99.0.2")
+	rebalanceLocalIP = net.ParseIP("10.99.0.1")
+)
+
+func testDownlinkRebalancePacket(bpfObjects *BpfObjects, dstIP net.IP) (int, []byte, error) {
+	packet := gopacket.NewSerializeBuffer()
+	if err := gopacket.SerializeLayers(packet, gopacket.SerializeOptions{},
+		&layers.Ethernet{
+			SrcMAC:       n9MAC,
+			DstMAC:       n3MAC,
+			EthernetType: layers.EthernetTypeIPv4,
+		},
+		&layers.IPv4{
+			Version:  4,
+			DstIP:    dstIP,
+			SrcIP:    net.IP{1, 1, 1, 1},
+			Protocol: layers.IPProtocolICMPv4,
+			IHL:      5,
+			TTL:      64,
+		},
+		&layers.ICMPv4{
+			TypeCode: layers.ICMPv4TypeEchoRequest,
+		},
+	); err != nil {
+		return 0, nil, fmt.Errorf("serializing input packet failed: %v", err)
+	}
+
+	return bpfObjects.UpfIpEntrypointFunc.Test(append(packet.Bytes(), 0x00, 0x00, 0x00, 0x00))
+}
+
+func TestDownlinkRebalance(t *testing.T) {
+
+	if err := IncreaseResourceLimits(); err != nil {
+		t.Fatalf("Can't increase resource limits: %s", err.Error())
+	}
+
+	bpfObjects := NewBpfObjects()
+
+	if err := bpfObjects.Load(); err != nil {
+		t.Fatalf("Loading bpf objects failed: %s", err.Error())
+	}
+
+	defer bpfObjects.Close()
+
+	entrypointConfig := IpEntrypointDataplaneConfig{
+		N3Ipv4Address:    binary.LittleEndian.Uint32(n3IP.To4()),
+		N9Ipv4Address:    binary.LittleEndian.Uint32(n9IP.To4()),
+		RebalanceIp:      binary.LittleEndian.Uint32(rebalanceIP.To4()),
+		RebalanceLocalIp: binary.LittleEndian.Uint32(rebalanceLocalIP.To4()),
+	}
+
+	if err := bpfObjects.GlobalConfig.Set(entrypointConfig); err != nil {
+		t.Fatalf("can't set dataplane global config: %s", err.Error())
+	}
+
+	prefixes := []string{"10.60.0.0/24", "10.61.0.0/16"}
+	for _, prefix := range prefixes {
+		if err := PutUeRebalancePrefix(bpfObjects.UeRebalanceLpmMap, prefix); err != nil {
+			t.Fatalf("can't put rebalance prefix %s: %s", prefix, err.Error())
+		}
+	}
+
+	inPrefix := []net.IP{
+		net.IP{10, 60, 0, 1},
+		net.IP{10, 60, 0, 254},
+		net.IP{10, 61, 128, 7},
+	}
+	for _, dstIP := range inPrefix {
+		bpfRet, bufOut, err := testDownlinkRebalancePacket(bpfObjects, dstIP)
+		if err != nil {
+			t.Fatalf("ebpf run failed: %v", err)
+		}
+		if bpfRet != 4 { // XDP_REDIRECT
+			t.Fatalf("unexpected return value for %s: %d", dstIP, bpfRet)
+		}
+
+		response := gopacket.NewPacket(bufOut, layers.LayerTypeEthernet, gopacket.Default)
+		ipLayer := response.Layer(layers.LayerTypeIPv4)
+		if ipLayer == nil {
+			t.Fatalf("no ipv4 layer in output packet for %s", dstIP)
+		}
+		outerIP, _ := ipLayer.(*layers.IPv4)
+		if !outerIP.DstIP.Equal(rebalanceIP) || !outerIP.SrcIP.Equal(rebalanceLocalIP) {
+			t.Fatalf("unexpected outer ipip header: %s -> %s", outerIP.SrcIP, outerIP.DstIP)
+		}
+		if outerIP.Protocol != layers.IPProtocolIPv4 {
+			t.Fatalf("unexpected outer protocol: %d", outerIP.Protocol)
+		}
+	}
+
+	outOfPrefix := []net.IP{
+		net.IP{10, 60, 1, 1},
+		net.IP{10, 62, 0, 1},
+		net.IP{8, 8, 8, 8},
+	}
+	for _, dstIP := range outOfPrefix {
+		bpfRet, _, err := testDownlinkRebalancePacket(bpfObjects, dstIP)
+		if err != nil {
+			t.Fatalf("ebpf run failed: %v", err)
+		}
+		if bpfRet != 2 { // XDP_PASS
+			t.Fatalf("unexpected return value for %s: %d", dstIP, bpfRet)
+		}
+	}
+}
+
 func TestEntrypointBenchmark(t *testing.T) {
 
 	if err := IncreaseResourceLimits(); err != nil {

cmd/ebpf/rebalance.go

@@ -0,0 +1,27 @@
+package ebpf
+
+import (
+	"net"
+	"unsafe"
+
+	"github.com/cilium/ebpf"
+)
+
+type UeRebalanceLpmKey struct {
+	PrefixLen uint32
+	Data      [4]byte
+}
+
+func PutUeRebalancePrefix(m *ebpf.Map, cidr string) error {
+	_, ipNet, err := net.ParseCIDR(cidr)
+	if err != nil {
+		return err
+	}
+	ones, _ := ipNet.Mask.Size()
+	key := UeRebalanceLpmKey{
+		PrefixLen: uint32(ones),
+	}
+	copy(key.Data[:], ipNet.IP.To4())
+	var value uint8 = 1
+	return m.Update(unsafe.Pointer(&key), unsafe.Pointer(&value), ebpf.UpdateAny)
+}

cmd/ebpf/xdp/n3n6_entrypoint.c

@@ -31,6 +31,7 @@
 #include "xdp/qer.h"
 #include "xdp/urr.h"
 #include "xdp/pdr.h"
+#include "xdp/rebalance.h"
 #include "xdp/sdf_filter.h"
 
 #include "xdp/utils/common.h"
@@ -67,6 +68,8 @@ struct dataplane_config {
     __u8   trace_in;
     __u8   trace_out;
     __u8   trace_blocked;
+    __u32  rebalance_local_ip;
+    __u32  rebalance_ip;
     __u8   ip6_ra_support;
     __u128 ip6_ra_prefix;
     __u16  ip6_ra_prefix_length;
@@ -202,11 +205,97 @@ static __always_inline const struct pdr* check_sdf_filters_gtp(struct packet_con
     return 0;
 }
 
+static __always_inline __u32 encap_ip_packet(struct packet_context *ctx, __u32 saddr, __u32 daddr) {
+
+    //static const size_t ip_encap_size = sizeof(struct iphdr);
+    size_t ip_packet_len = 0;
+    __u8 ip_proto = 0;
+    if (ctx->ip4) {
+        ip_packet_len = bpf_ntohs(ctx->ip4->tot_len);
+        ip_proto = 4;
+    }
+    // else if (ctx->ip6) {
+    //     ip_packet_len = bpf_ntohs(ctx->ip6->payload_len) + sizeof(struct ipv6hdr);
+    //     ip_proto = 41;
+    // }
+    else
+        return -1;
+
+    int result = bpf_xdp_adjust_head(ctx->xdp_ctx, (__s32) - sizeof(struct iphdr));
+    if (result)
+        return -1;
+
+    char *data = (char *)(long)ctx->xdp_ctx->data;
+    const char *data_end = (const char *)(long)ctx->xdp_ctx->data_end;
+
+    struct ethhdr *eth = (struct ethhdr *)data;
+    if ((const char *)(eth + 1) > data_end)
+        return -1;
+
+    struct ethhdr *orig_eth = (struct ethhdr *)(data + sizeof(struct iphdr));
+    if ((const char *)(orig_eth + 1) > data_end)
+        return -1;
+
+    __builtin_memcpy(eth, orig_eth, sizeof(*eth));
+    eth->h_proto = bpf_htons(ETH_P_IP);
+
+    struct iphdr *ip = (struct iphdr *)(eth + 1);
+    if ((const char *)(ip + 1) > data_end)
+        return -1;
+
+    /* Add the outer IP header */
+    ip->version = 4;
+    ip->ihl = 5; /* No options */
+    ip->tos = 0;
+    ip->tot_len = bpf_htons(ip_packet_len + sizeof(struct iphdr));
+    ip->id = 0;            /* No fragmentation */
+    ip->frag_off = 0x0040; /* Don't fragment; Fragment offset = 0 */
+    ip->ttl = 64;
+    ip->protocol = ip_proto;
+    ip->check = 0;
+    ip->saddr = saddr;
+    ip->daddr = daddr;
+
+    ip->check = ipv4_csum(ip, sizeof(*ip));
+
+    /* Update packet pointers */
+    context_set_ip4(ctx, (char *)(long)ctx->xdp_ctx->data, (const char *)(long)ctx->xdp_ctx->data_end, eth, ip, 0, 0);
+    return 0;
+}
+
+static __always_inline int try_rebalance(struct packet_context *ctx) {
+    if(!global_config.rebalance_ip || !ctx->ip4)
+        return -1;
+
+    struct ue_rebalance_lpm_key key = {
+        .prefix_len = 32,
+    };
+    __builtin_memcpy(key.data, &ctx->ip4->daddr, sizeof(key.data));
+    if(!bpf_map_lookup_elem(&ue_rebalance_lpm_map, &key))
+        return -1;
+
+    if(global_config.rebalance_local_ip) {
+        if(0 == encap_ip_packet(ctx, global_config.rebalance_local_ip, global_config.rebalance_ip))
+            return route_ipv4(ctx->xdp_ctx, ctx->eth, ctx->ip4);
+        return -1;
+    }
+    return route_ipv4_to(ctx->xdp_ctx, ctx->eth, ctx->ip4, global_config.rebalance_ip);
+}
+
 static __always_inline enum xdp_action handle_n6_packet_ipv4(struct packet_context *ctx) {
     const struct iphdr *ip4 = ctx->ip4;
     struct pdr_info *session = bpf_map_lookup_elem(&pdr_map_downlink_ip4, &ip4->daddr);
     if (!session) {
         upf_printk("upf: [n6] no downlink session for ip:%pI4", &ip4->daddr);
+
+        if(ctx->redirected) {
+            upf_printk("upf: [n6] drop redirected packet: no session for ip:%pI4", &ip4->daddr);
+            return XDP_DROP;
+        }
+
+        int rebalance_action = try_rebalance(ctx);
+        if(rebalance_action != -1)
+            return rebalance_action;
         return DEFAULT_XDP_ACTION;
     }
 
@@ -239,8 +328,16 @@ static __always_inline enum xdp_action handle_n6_packet_ipv4(struct packet_conte
     }
 
     // Only forwarding action is supported at the moment
-    if (!(far->action & FAR_FORW))
+    if (!(far->action & FAR_FORW)) {
+        if(ctx->redirected) {
+            upf_printk("upf: [n6] drop redirected: non-forward session ip:%pI4", &ip4->daddr);
+            return XDP_DROP;
+        }
+        int rebalance_action = try_rebalance(ctx);
+        if(rebalance_action != -1)
+            return rebalance_action;
         return XDP_DROP;
+    }
 
     // Only outer header GTP/UDP/IPv4 is supported at the moment
     if (!(far->outer_header_creation & OHC_GTP_U_UDP_IPv4))
@@ -532,6 +629,10 @@ static __always_inline enum xdp_action handle_gtpu(struct packet_context *ctx) {
 
 static __always_inline enum xdp_action handle_ip4(struct packet_context *ctx) {
     int l4_protocol = parse_ip4(ctx);
+    if(l4_protocol == IPPROTO_IPIP) {
+        l4_protocol = parse_ip4(ctx);
+        ctx->redirected = 1;
+    }
     switch (l4_protocol) {
         case IPPROTO_ICMP: {
             increment_counter(ctx->counters, rx_icmp);

cmd/ebpf/xdp/rebalance.h

@@ -0,0 +1,36 @@
+/**
+ * Copyright 2026 Edgecom LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <bpf/bpf_helpers.h>
+#include <linux/bpf.h>
+
+/* LPM trie key: prefix length in bits followed by the IPv4 prefix in network byte order */
+struct ue_rebalance_lpm_key {
+    __u32 prefix_len;
+    __u8  data[4];
+};
+
+/* UE prefixes eligible for downlink rebalance */
+struct
+{
+    __uint(type, BPF_MAP_TYPE_LPM_TRIE);
+    __type(key, struct ue_rebalance_lpm_key);
+    __type(value, __u8);
+    __uint(max_entries, 16);
+    __uint(map_flags, BPF_F_NO_PREALLOC);
+} ue_rebalance_lpm_map SEC(".maps");