21324261259

Committed 25 Jan 2026 12:44AM UTC coverage: 36.989% (-0.7%) from 37.705%

Build # 21324261259

Build Type

push

github

Committed by

web-flow

Commit Message

feat(ebpf): Add backend balancing mechanism (#14)

Run Details

23 of 96 new or added lines in 4 files covered. (23.96%)

1 existing line in 1 file now uncovered.

371 of 1003 relevant lines covered (36.99%)

0.41 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/pkg/dataplane/proxy/forward.go

package proxy

import (
        "context"
        "fmt"
        "net"

        log "log/slog"

        "github.com/cilium/ebpf"
        "github.com/cilium/ebpf/link"
        "github.com/pkg/errors"

        wgebpf "github.com/stillya/wg-relay/ebpf"
        "github.com/stillya/wg-relay/pkg/bpf"
        "github.com/stillya/wg-relay/pkg/dataplane/config"
        "github.com/stillya/wg-relay/pkg/dataplane/maps"
        "github.com/stillya/wg-relay/pkg/utils"
)

// ForwardLoader manages XDP-based forward proxy
type ForwardLoader struct {
        cfg   config.ProxyConfig
        objs  *wgebpf.WgForwardProxyObjects
        links []link.Link
}

// NewForwardLoader creates a new forward proxy loader
func NewForwardLoader() (*ForwardLoader, error) {
        return &ForwardLoader{}, nil
}

// LoadAndAttach loads the eBPF program and attaches it to the configured interfaces.
func (fp *ForwardLoader) LoadAndAttach(ctx context.Context, cfg config.ProxyConfig) error {
        fp.cfg = cfg

        if err := fp.loadEBPF(); err != nil {
                return errors.Wrap(err, "failed to load eBPF objects")
        }

        if err := fp.attachToInterfaces(); err != nil {
                fp.Close()
                return errors.Wrap(err, "failed to attach to interfaces")
        }

        backendCount := len(cfg.Forward.Backends)

        log.Info("Forward proxy loaded and attached",
                "enabled", cfg.Enabled,
                "backend_count", backendCount,
        )

        return nil
}

func (fp *ForwardLoader) loadEBPF() error {
        spec, err := wgebpf.LoadWgForwardProxy()
        if err != nil {
                return errors.Wrap(err, "failed to load forward proxy spec")
        }

        if err := bpf.Configure(spec, &fp.cfg); err != nil {
                return errors.Wrap(err, "failed to configure static variables")
        }

        fp.objs = &wgebpf.WgForwardProxyObjects{}
        opts := &ebpf.CollectionOptions{
                Programs: ebpf.ProgramOptions{
                        LogLevel:     2,
                        LogSizeStart: 16777216,
                },
        }

        if err := spec.LoadAndAssign(fp.objs, opts); err != nil {
                return errors.Wrap(err, "failed to load forward proxy eBPF objects")
        }

        if err := fp.configureBackendMap(); err != nil {
                return errors.Wrap(err, "failed to configure backend map")
        }

        log.Info("Forward proxy eBPF program loaded")
        return nil
}

func (fp *ForwardLoader) configureBackendMap() error {
        backends := fp.cfg.GetBackends()

        if len(backends) == 0 {
                return errors.New("no backends configured")
        }

        for i, backend := range backends {
                key := uint32(i) //nolint:gosec // G304: it's fine
                ip, err := utils.IPToUint32(backend.IP)
                if err != nil {
                        return errors.Wrapf(err, "failed to convert IP address %s to uint32", backend.IP)
                }
                entry := &wgebpf.WgForwardProxyBackendEntry{
                        Ip:   ip,
                        Port: backend.Port,
                }
                if err := fp.objs.BackendMap.Put(&key, entry); err != nil {
                        return errors.Wrapf(err, "failed to add backend[%d] to map", i)
                }
        }

        countKey := uint32(0)
        count := uint32(len(backends)) //nolint:gosec // G304: it's fine
        if err := fp.objs.BackendCount.Put(&countKey, &count); err != nil {
                return errors.Wrap(err, "failed to set backend count")
        }

        backendAddrs := make([]string, len(backends))
        for i, b := range fp.cfg.Forward.Backends {
                if b.Port > 0 {
                        backendAddrs[i] = fmt.Sprintf("%s:%d", b.IP, b.Port)
                } else {
                        backendAddrs[i] = b.IP
                }
        }

        log.Info("Backend map configured",
                "count", len(backends),
                "backends", backendAddrs,
        )

        return nil
}

// attachToInterfaces attaches the XDP program to configured interfaces
func (fp *ForwardLoader) attachToInterfaces() error {
        for _, interfaceName := range fp.cfg.Interfaces {
                if err := fp.attachToInterface(interfaceName); err != nil {
                        fp.cleanupLinks()
                        return errors.Wrapf(err, "failed to attach to interface %s", interfaceName)
                }
        }
        return nil
}

// attachToInterface attaches XDP program to a single interface
func (fp *ForwardLoader) attachToInterface(interfaceName string) error {
        iface, err := net.InterfaceByName(interfaceName)
        if err != nil {
                return errors.Wrapf(err, "failed to get interface %s", interfaceName)
        }

        var flags link.XDPAttachFlags
        var mode string
        switch fp.cfg.DriverMode {
        case "generic":
                flags = link.XDPGenericMode
                mode = "generic"
        case "driver":
                flags = link.XDPDriverMode
                mode = "driver"
        case "offload":
                flags = link.XDPOffloadMode
                mode = "offload"
        default:
                // Default to driver mode
                flags = link.XDPDriverMode
                mode = "driver"
        }

        log.Info("Attaching XDP forward proxy", "interface", interfaceName, "index", iface.Index, "mode", mode)

        xdpLink, err := link.AttachXDP(link.XDPOptions{
                Program:   fp.objs.WgForwardProxy,
                Interface: iface.Index,
                Flags:     flags,
        })
        if err != nil {
                return errors.Wrapf(err, "failed to attach XDP program to interface %s in %s mode", interfaceName, mode)
        }

        log.Info("XDP forward proxy attached successfully", "interface", interfaceName, "mode", mode)
        fp.links = append(fp.links, xdpLink)
        return nil
}

// cleanupLinks cleans up all attached links
func (fp *ForwardLoader) cleanupLinks() {
        for _, l := range fp.links {
                if l != nil {
                        err := l.Close()
                        if err != nil {
                                log.Error("Failed to close XDP link", "error", err)
                        } else {
                                log.Info("XDP link closed successfully")
                        }
                }
        }
        fp.links = nil
}

// Close cleans up all resources
func (fp *ForwardLoader) Close() error {
        var errs []error

        for i, l := range fp.links {
                if l != nil {
                        if err := l.Close(); err != nil {
                                errs = append(errs, errors.Wrapf(err, "failed to close XDP l %d", i))
                        }
                }
        }

        if fp.objs != nil {
                if err := fp.objs.Close(); err != nil {
                        errs = append(errs, errors.Wrap(err, "failed to close forward proxy eBPF objects"))
                }
        }

        if len(errs) > 0 {
                return errs[0]
        }

        return nil
}

// Maps returns all eBPF maps used by the forward proxy
func (fp *ForwardLoader) Maps() *maps.Maps {
        var metricsMap *ebpf.Map
        if fp.objs != nil {
                metricsMap = fp.objs.MetricsMap
        }

        mapsCollection := maps.NewMaps(metricsMap)

        if fp.objs != nil {
                if fp.objs.ConnectionMap != nil {
                        mapsCollection.AddOtherMap("ConnectionMap", fp.objs.ConnectionMap)
                }

                if fp.objs.NatReverseMap != nil {
                        mapsCollection.AddOtherMap("NatReverseMap", fp.objs.NatReverseMap)
                }
        }

        return mapsCollection
}

1	package proxy
2
3	import (
4	"context"
5	"fmt"
6	"net"
7
8	log "log/slog"
9
10	"github.com/cilium/ebpf"
11	"github.com/cilium/ebpf/link"
12	"github.com/pkg/errors"
13
14	wgebpf "github.com/stillya/wg-relay/ebpf"
15	"github.com/stillya/wg-relay/pkg/bpf"
16	"github.com/stillya/wg-relay/pkg/dataplane/config"
17	"github.com/stillya/wg-relay/pkg/dataplane/maps"
18	"github.com/stillya/wg-relay/pkg/utils"
19	)
20
21	// ForwardLoader manages XDP-based forward proxy
22	type ForwardLoader struct {
23	cfg config.ProxyConfig
24	objs *wgebpf.WgForwardProxyObjects
25	links []link.Link
26	}
27
28	// NewForwardLoader creates a new forward proxy loader
29	func NewForwardLoader() (*ForwardLoader, error) {	×
30	return &ForwardLoader{}, nil	×
31	}	×
32
33	// LoadAndAttach loads the eBPF program and attaches it to the configured interfaces.
34	func (fp *ForwardLoader) LoadAndAttach(ctx context.Context, cfg config.ProxyConfig) error {	×
35	fp.cfg = cfg	×
36		×
37	if err := fp.loadEBPF(); err != nil {	×
38	return errors.Wrap(err, "failed to load eBPF objects")	×
39	}	×
40
41	if err := fp.attachToInterfaces(); err != nil {	×
42	fp.Close()	×
43	return errors.Wrap(err, "failed to attach to interfaces")	×
44	}	×
45
NEW 46	backendCount := len(cfg.Forward.Backends)	×
NEW 47		×
48	log.Info("Forward proxy loaded and attached",	×
49	"enabled", cfg.Enabled,	×
NEW 50	"backend_count", backendCount,	×
NEW 51	)	×
52		×
53	return nil	×
54	}
55
56	func (fp *ForwardLoader) loadEBPF() error {	×
57	spec, err := wgebpf.LoadWgForwardProxy()	×
58	if err != nil {	×
59	return errors.Wrap(err, "failed to load forward proxy spec")	×
60	}	×
61
62	if err := bpf.Configure(spec, &fp.cfg); err != nil {	×
63	return errors.Wrap(err, "failed to configure static variables")	×
64	}	×
65
66	fp.objs = &wgebpf.WgForwardProxyObjects{}	×
67	opts := &ebpf.CollectionOptions{	×
68	Programs: ebpf.ProgramOptions{	×
69	LogLevel: 2,	×
70	LogSizeStart: 16777216,	×
71	},	×
72	}	×
73		×
74	if err := spec.LoadAndAssign(fp.objs, opts); err != nil {	×
75	return errors.Wrap(err, "failed to load forward proxy eBPF objects")	×
76	}	×
77
78	if err := fp.configureBackendMap(); err != nil {	×
79	return errors.Wrap(err, "failed to configure backend map")	×
80	}	×
81
82	log.Info("Forward proxy eBPF program loaded")	×
83	return nil	×
84	}
85
86	func (fp *ForwardLoader) configureBackendMap() error {	×
NEW 87	backends := fp.cfg.GetBackends()	×
NEW 88		×
NEW 89	if len(backends) == 0 {	×
NEW 90	return errors.New("no backends configured")	×
NEW 91	}	×
92
NEW 93	for i, backend := range backends {	×
NEW 94	key := uint32(i) //nolint:gosec // G304: it's fine	×
NEW 95	ip, err := utils.IPToUint32(backend.IP)	×
NEW 96	if err != nil {	×
NEW 97	return errors.Wrapf(err, "failed to convert IP address %s to uint32", backend.IP)	×
NEW 98	}	×
NEW 99	entry := &wgebpf.WgForwardProxyBackendEntry{	×
NEW 100	Ip: ip,	×
NEW 101	Port: backend.Port,	×
NEW 102	}	×
NEW 103	if err := fp.objs.BackendMap.Put(&key, entry); err != nil {	×
NEW 104	return errors.Wrapf(err, "failed to add backend[%d] to map", i)	×
NEW 105	}	×
106	}
107
NEW 108	countKey := uint32(0)	×
NEW 109	count := uint32(len(backends)) //nolint:gosec // G304: it's fine	×
NEW 110	if err := fp.objs.BackendCount.Put(&countKey, &count); err != nil {	×
NEW 111	return errors.Wrap(err, "failed to set backend count")	×
NEW 112	}	×
113
NEW 114	backendAddrs := make([]string, len(backends))	×
NEW 115	for i, b := range fp.cfg.Forward.Backends {	×
NEW 116	if b.Port > 0 {	×
NEW 117	backendAddrs[i] = fmt.Sprintf("%s:%d", b.IP, b.Port)	×
NEW 118	} else {	×
NEW 119	backendAddrs[i] = b.IP	×
NEW 120	}	×
121	}
122
NEW 123	log.Info("Backend map configured",	×
NEW 124	"count", len(backends),	×
NEW 125	"backends", backendAddrs,	×
NEW 126	)	×
127		×
128	return nil	×
129	}
130
131	// attachToInterfaces attaches the XDP program to configured interfaces
132	func (fp *ForwardLoader) attachToInterfaces() error {	×
133	for _, interfaceName := range fp.cfg.Interfaces {	×
134	if err := fp.attachToInterface(interfaceName); err != nil {	×
135	fp.cleanupLinks()	×
136	return errors.Wrapf(err, "failed to attach to interface %s", interfaceName)	×
137	}	×
138	}
139	return nil	×
140	}
141
142	// attachToInterface attaches XDP program to a single interface
143	func (fp *ForwardLoader) attachToInterface(interfaceName string) error {	×
144	iface, err := net.InterfaceByName(interfaceName)	×
145	if err != nil {	×
146	return errors.Wrapf(err, "failed to get interface %s", interfaceName)	×
147	}	×
148
149	var flags link.XDPAttachFlags	×
150	var mode string	×
151	switch fp.cfg.DriverMode {	×
152	case "generic":	×
153	flags = link.XDPGenericMode	×
154	mode = "generic"	×
155	case "driver":	×
156	flags = link.XDPDriverMode	×
157	mode = "driver"	×
158	case "offload":	×
159	flags = link.XDPOffloadMode	×
160	mode = "offload"	×
161	default:	×
162	// Default to driver mode	×
163	flags = link.XDPDriverMode	×
164	mode = "driver"	×
165	}
166
167	log.Info("Attaching XDP forward proxy", "interface", interfaceName, "index", iface.Index, "mode", mode)	×
168		×
169	xdpLink, err := link.AttachXDP(link.XDPOptions{	×
170	Program: fp.objs.WgForwardProxy,	×
171	Interface: iface.Index,	×
172	Flags: flags,	×
173	})	×
174	if err != nil {	×
175	return errors.Wrapf(err, "failed to attach XDP program to interface %s in %s mode", interfaceName, mode)	×
176	}	×
177
178	log.Info("XDP forward proxy attached successfully", "interface", interfaceName, "mode", mode)	×
179	fp.links = append(fp.links, xdpLink)	×
180	return nil	×
181	}
182
183	// cleanupLinks cleans up all attached links
184	func (fp *ForwardLoader) cleanupLinks() {	×
185	for _, l := range fp.links {	×
186	if l != nil {	×
187	err := l.Close()	×
188	if err != nil {	×
189	log.Error("Failed to close XDP link", "error", err)	×
190	} else {	×
191	log.Info("XDP link closed successfully")	×
192	}	×
193	}
194	}
195	fp.links = nil	×
196	}
197
198	// Close cleans up all resources
199	func (fp *ForwardLoader) Close() error {	×
200	var errs []error	×
201		×
202	for i, l := range fp.links {	×
203	if l != nil {	×
204	if err := l.Close(); err != nil {	×
205	errs = append(errs, errors.Wrapf(err, "failed to close XDP l %d", i))	×
206	}	×
207	}
208	}
209
210	if fp.objs != nil {	×
211	if err := fp.objs.Close(); err != nil {	×
212	errs = append(errs, errors.Wrap(err, "failed to close forward proxy eBPF objects"))	×
213	}	×
214	}
215
216	if len(errs) > 0 {	×
217	return errs[0]	×
218	}	×
219
220	return nil	×
221	}
222
223	// Maps returns all eBPF maps used by the forward proxy
224	func (fp ForwardLoader) Maps() maps.Maps {	×
225	var metricsMap *ebpf.Map	×
226	if fp.objs != nil {	×
227	metricsMap = fp.objs.MetricsMap	×
228	}	×
229
230	mapsCollection := maps.NewMaps(metricsMap)	×
231		×
232	if fp.objs != nil {	×
233	if fp.objs.ConnectionMap != nil {	×
234	mapsCollection.AddOtherMap("ConnectionMap", fp.objs.ConnectionMap)	×
235	}	×
236
237	if fp.objs.NatReverseMap != nil {	×
238	mapsCollection.AddOtherMap("NatReverseMap", fp.objs.NatReverseMap)	×
239	}	×
240	}
241
242	return mapsCollection	×
243	}

stillya / wg-relay / 21324261259

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous