Add shim client to simulate allowed IP removal

Direct allowed IP removal is a new feature, so its availability is
limited for now. Clients who want to take advantage of this capability
would need to know ahead of time if WireGuard on their system supports
it or probe to see if they're running on a system that supports it. To
ease the transition, add the WithShim option for clients:

	c, err := wgctrl.New(wgctrl.WithShim)

WithShim wraps internal clients with a shim client that probes to see
if direct IP removal is supported on their system. If not, the shim
client emulates the effect by assigning IPs to a dummy peer then
removing that peer. At some point in the future, this option should no
longer be necessary once the feature becomes commonplace.

In my case, I plan to use WithShim to simplify Cilium's WireGuard
orchestration logic.

Signed-off-by: Jordan Rife <[email protected]>

Author: jrife

client.go

@@ -5,11 +5,22 @@ import (
 	"os"
 
 	"golang.zx2c4.com/wireguard/wgctrl/internal/wginternal"
+	"golang.zx2c4.com/wireguard/wgctrl/internal/wgshim"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 )
 
-// Expose an identical interface to the underlying packages.
-var _ wginternal.Client = &Client{}
+// An Option configures a client in some way. Options may be provided when
+// calling New().
+type Option func(wginternal.Client) wginternal.Client
+
+// WithShim wraps the client in a shim that probes for the capabilities
+// supported by the underlying WireGuard implementation and emulates missing
+// capabilities.
+//
+// This option ensures backwards and forwards compatibility.
+func WithShim(c wginternal.Client) wginternal.Client {
+	return wgshim.New(c)
+}
 
 // A Client provides access to WireGuard device information.
 type Client struct {
@@ -18,13 +29,20 @@ type Client struct {
 	cs []wginternal.Client
 }
 
-// New creates a new Client.
-func New() (*Client, error) {
+// New creates a new Client. Callers may provide a list of Options that modify
+// client behavior.
+func New(opts ...Option) (*Client, error) {
 	cs, err := newClients()
 	if err != nil {
 		return nil, err
 	}
 
+	for _, opt := range opts {
+		for i := range cs {
+			cs[i] = opt(cs[i])
+		}
+	}
+
 	return &Client{
 		cs: cs,
 	}, nil

client_test.go

@@ -2,11 +2,13 @@ package wgctrl
 
 import (
 	"errors"
+	"net"
 	"os"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"golang.zx2c4.com/wireguard/wgctrl/internal/wginternal"
+	"golang.zx2c4.com/wireguard/wgctrl/internal/wgtest"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 )
 
@@ -224,11 +226,263 @@ func TestClientConfigureDevice(t *testing.T) {
 	}
 }
 
+func TestClientConfigureDeviceWithShim(t *testing.T) {
+	type devicesFunc func() ([]*wgtypes.Device, error)
+	type supportsFunc func(name string) (bool, error)
+
+	var (
+		ip           = wgtest.MustCIDR("192.0.2.0/32")
+		peerKey      = wgtest.MustPublicKey()
+		dummyPeerKey = wgtypes.Key{}
+		device       = "wg0"
+
+		notSupported = func(_ string) (bool, error) {
+			return false, nil
+		}
+
+		supported = func(_ string) (bool, error) {
+			return true, nil
+		}
+
+		returnsError = func(_ string) (bool, error) {
+			return false, errFoo
+		}
+
+		peerHasIP = func() ([]*wgtypes.Device, error) {
+			return []*wgtypes.Device{
+				{
+					Name: device,
+					Peers: []wgtypes.Peer{
+						{
+							PublicKey: peerKey,
+							AllowedIPs: []net.IPNet{
+								ip,
+							},
+						},
+					},
+				},
+			}, nil
+		}
+
+		peerDoesNotHaveIP = func() ([]*wgtypes.Device, error) {
+			return []*wgtypes.Device{
+				{
+					Name: device,
+					Peers: []wgtypes.Peer{
+						{
+							PublicKey:  peerKey,
+							AllowedIPs: []net.IPNet{},
+						},
+					},
+				},
+			}, nil
+		}
+
+		otherPeerHasIP = func() ([]*wgtypes.Device, error) {
+			return []*wgtypes.Device{
+				{
+					Name: device,
+					Peers: []wgtypes.Peer{
+						{
+							PublicKey:  peerKey,
+							AllowedIPs: []net.IPNet{},
+						},
+						{
+							PublicKey: wgtest.MustPublicKey(),
+							AllowedIPs: []net.IPNet{
+								ip,
+							},
+						},
+					},
+				},
+			}, nil
+		}
+
+		removeAllowedIP = wgtypes.Config{
+			Peers: []wgtypes.PeerConfig{
+				{
+					PublicKey: peerKey,
+					AllowedIPs: []wgtypes.AllowedIPConfig{
+						{
+							IPNet:  ip,
+							Remove: true,
+						},
+					},
+				},
+			},
+		}
+
+		removeAllowedIPUndone = wgtypes.Config{
+			Peers: []wgtypes.PeerConfig{
+				{
+					PublicKey: peerKey,
+					AllowedIPs: []wgtypes.AllowedIPConfig{
+						{
+							IPNet:  ip,
+							Remove: true,
+						},
+						{
+							IPNet: ip,
+						},
+					},
+				},
+			},
+		}
+
+		simulateRemoveAllowedIP = wgtypes.Config{
+			Peers: []wgtypes.PeerConfig{
+				{
+					PublicKey: dummyPeerKey,
+					AllowedIPs: []wgtypes.AllowedIPConfig{
+						{
+							IPNet: ip,
+						},
+					},
+				},
+				{
+					PublicKey: peerKey,
+				},
+				{
+					PublicKey: dummyPeerKey,
+					Remove:    true,
+				},
+			},
+		}
+
+		addAllowedIP = wgtypes.Config{
+			Peers: []wgtypes.PeerConfig{
+				{
+					PublicKey: peerKey,
+					AllowedIPs: []wgtypes.AllowedIPConfig{
+						{
+							IPNet: ip,
+						},
+					},
+				},
+			},
+		}
+
+		dontRemoveIP = wgtypes.Config{
+			Peers: []wgtypes.PeerConfig{
+				{
+					PublicKey: peerKey,
+				},
+			},
+		}
+	)
+
+	tests := []struct {
+		name       string
+		supportsFn supportsFunc
+		devicesFn  devicesFunc
+		cfg        wgtypes.Config
+		expectCfg  wgtypes.Config
+		err        error
+	}{
+		{
+			name:       "not supported + remove IP + peer has IP",
+			supportsFn: notSupported,
+			devicesFn:  peerHasIP,
+			cfg:        removeAllowedIP,
+			expectCfg:  simulateRemoveAllowedIP,
+			err:        nil,
+		},
+		{
+			name:       "not supported + remove IP + peer does not have IP",
+			supportsFn: notSupported,
+			devicesFn:  peerDoesNotHaveIP,
+			cfg:        removeAllowedIP,
+			expectCfg:  dontRemoveIP,
+			err:        nil,
+		},
+		{
+			name:       "not supported + remove IP + other peer has IP",
+			supportsFn: notSupported,
+			devicesFn:  otherPeerHasIP,
+			cfg:        removeAllowedIP,
+			expectCfg:  dontRemoveIP,
+			err:        nil,
+		},
+		{
+			name:       "not supported + remove IP undone",
+			supportsFn: notSupported,
+			devicesFn:  peerHasIP,
+			cfg:        removeAllowedIPUndone,
+			expectCfg:  addAllowedIP,
+			err:        nil,
+		},
+		{
+			name:       "not supported + don't remove IP",
+			supportsFn: notSupported,
+			devicesFn:  peerHasIP,
+			cfg:        addAllowedIP,
+			expectCfg:  addAllowedIP,
+			err:        nil,
+		},
+		{
+			name:       "supported + remove IP",
+			supportsFn: supported,
+			cfg:        removeAllowedIP,
+			expectCfg:  removeAllowedIP,
+			err:        nil,
+		},
+		{
+			name:       "supported + don't remove IP",
+			supportsFn: supported,
+			cfg:        addAllowedIP,
+			expectCfg:  addAllowedIP,
+			err:        nil,
+		},
+		{
+			name:       "probe error + remove IP",
+			supportsFn: returnsError,
+			cfg:        removeAllowedIP,
+			expectCfg:  wgtypes.Config{},
+			err:        errFoo,
+		},
+		{
+			name:       "probe error + don't remove IP",
+			supportsFn: returnsError,
+			cfg:        addAllowedIP,
+			expectCfg:  wgtypes.Config{},
+			err:        errFoo,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var finalCfg wgtypes.Config
+
+			cs := WithShim(&testClient{
+				ConfigureDeviceFunc: func(name string, cfg wgtypes.Config) error {
+					finalCfg = cfg
+
+					return nil
+				},
+				DevicesFunc:                 tt.devicesFn,
+				SupportsAllowedIPRemoveFunc: tt.supportsFn,
+			})
+
+			c := &Client{cs: []wginternal.Client{cs}}
+
+			err := c.ConfigureDevice(device, tt.cfg)
+			if !errors.Is(err, tt.err) {
+				t.Fatalf("unexpected error: got %s, want %s", err, tt.err)
+			}
+
+			if diff := cmp.Diff(tt.expectCfg, finalCfg); diff != "" {
+				t.Fatalf("unexpected config (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
 type testClient struct {
-	CloseFunc           func() error
-	DevicesFunc         func() ([]*wgtypes.Device, error)
-	DeviceFunc          func(name string) (*wgtypes.Device, error)
-	ConfigureDeviceFunc func(name string, cfg wgtypes.Config) error
+	CloseFunc                   func() error
+	DevicesFunc                 func() ([]*wgtypes.Device, error)
+	DeviceFunc                  func(name string) (*wgtypes.Device, error)
+	ConfigureDeviceFunc         func(name string, cfg wgtypes.Config) error
+	SupportsAllowedIPRemoveFunc func(name string) (bool, error)
 }
 
 func (c *testClient) Close() error                        { return c.CloseFunc() }
@@ -240,3 +494,7 @@ func (c *testClient) Device(name string) (*wgtypes.Device, error) {
 func (c *testClient) ConfigureDevice(name string, cfg wgtypes.Config) error {
 	return c.ConfigureDeviceFunc(name, cfg)
 }
+
+func (c *testClient) SupportsAllowedIPRemove(name string) (bool, error) {
+	return c.SupportsAllowedIPRemoveFunc(name)
+}

internal/wgfreebsd/client_freebsd.go

@@ -217,6 +217,10 @@ func (c *Client) ConfigureDevice(name string, cfg wgtypes.Config) error {
 	return nil
 }
 
+func (c *Client) SupportsAllowedIPRemove(name string) (bool, error) {
+	return false, nil
+}
+
 // deviceName converts an interface name string to the format required to pass
 // with wgh.WGGetServ.
 func deviceName(name string) ([16]byte, error) {

internal/wginternal/client.go

@@ -18,4 +18,5 @@ type Client interface {
 	Devices() ([]*wgtypes.Device, error)
 	Device(name string) (*wgtypes.Device, error)
 	ConfigureDevice(name string, cfg wgtypes.Config) error
+	SupportsAllowedIPRemove(name string) (bool, error)
 }