Implement ECMP for unsafe_routes (#1332)

routing/balance.go

@@ -0,0 +1,39 @@
+package routing
+
+import (
+	"net/netip"
+
+	"github.com/slackhq/nebula/firewall"
+)
+
+// Hashes the packet source and destination port and always returns a positive integer
+// Based on 'Prospecting for Hash Functions'
+//   - https://nullprogram.com/blog/2018/07/31/
+//   - https://github.com/skeeto/hash-prospector
+//     [16 21f0aaad 15 d35a2d97 15] = 0.10760229515479501
+func hashPacket(p *firewall.Packet) int {
+	x := (uint32(p.LocalPort) << 16) | uint32(p.RemotePort)
+	x ^= x >> 16
+	x *= 0x21f0aaad
+	x ^= x >> 15
+	x *= 0xd35a2d97
+	x ^= x >> 15
+
+	return int(x) & 0x7FFFFFFF
+}
+
+// For this function to work correctly it requires that the buckets for the gateways have been calculated
+// If the contract is violated balancing will not work properly and the second return value will return false
+func BalancePacket(fwPacket *firewall.Packet, gateways []Gateway) (netip.Addr, bool) {
+	hash := hashPacket(fwPacket)
+
+	for i := range gateways {
+		if hash <= gateways[i].BucketUpperBound() {
+			return gateways[i].Addr(), true
+		}
+	}
+
+	// If you land here then the buckets for the gateways are not properly calculated
+	// Fallback to random routing and let the caller know
+	return gateways[hash%len(gateways)].Addr(), false
+}

routing/balance_test.go

@@ -0,0 +1,144 @@
+package routing
+
+import (
+	"net/netip"
+	"testing"
+
+	"github.com/slackhq/nebula/firewall"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPacketsAreBalancedEqually(t *testing.T) {
+
+	gateways := []Gateway{}
+
+	gw1Addr := netip.MustParseAddr("1.0.0.1")
+	gw2Addr := netip.MustParseAddr("1.0.0.2")
+	gw3Addr := netip.MustParseAddr("1.0.0.3")
+
+	gateways = append(gateways, NewGateway(gw1Addr, 1))
+	gateways = append(gateways, NewGateway(gw2Addr, 1))
+	gateways = append(gateways, NewGateway(gw3Addr, 1))
+
+	CalculateBucketsForGateways(gateways)
+
+	gw1count := 0
+	gw2count := 0
+	gw3count := 0
+
+	iterationCount := uint16(65535)
+	for i := uint16(0); i < iterationCount; i++ {
+		packet := firewall.Packet{
+			LocalAddr:  netip.MustParseAddr("192.168.1.1"),
+			RemoteAddr: netip.MustParseAddr("10.0.0.1"),
+			LocalPort:  i,
+			RemotePort: 65535 - i,
+			Protocol:   6, // TCP
+			Fragment:   false,
+		}
+
+		selectedGw, ok := BalancePacket(&packet, gateways)
+		assert.True(t, ok)
+
+		switch selectedGw {
+		case gw1Addr:
+			gw1count += 1
+		case gw2Addr:
+			gw2count += 1
+		case gw3Addr:
+			gw3count += 1
+		}
+
+	}
+
+	// Assert packets are balanced, allow variation of up to 100 packets per gateway
+	assert.InDeltaf(t, iterationCount/3, gw1count, 100, "Expected %d +/- 100, but got %d", iterationCount/3, gw1count)
+	assert.InDeltaf(t, iterationCount/3, gw2count, 100, "Expected %d +/- 100, but got %d", iterationCount/3, gw1count)
+	assert.InDeltaf(t, iterationCount/3, gw3count, 100, "Expected %d +/- 100, but got %d", iterationCount/3, gw1count)
+
+}
+
+func TestPacketsAreBalancedByPriority(t *testing.T) {
+
+	gateways := []Gateway{}
+
+	gw1Addr := netip.MustParseAddr("1.0.0.1")
+	gw2Addr := netip.MustParseAddr("1.0.0.2")
+
+	gateways = append(gateways, NewGateway(gw1Addr, 10))
+	gateways = append(gateways, NewGateway(gw2Addr, 5))
+
+	CalculateBucketsForGateways(gateways)
+
+	gw1count := 0
+	gw2count := 0
+
+	iterationCount := uint16(65535)
+	for i := uint16(0); i < iterationCount; i++ {
+		packet := firewall.Packet{
+			LocalAddr:  netip.MustParseAddr("192.168.1.1"),
+			RemoteAddr: netip.MustParseAddr("10.0.0.1"),
+			LocalPort:  i,
+			RemotePort: 65535 - i,
+			Protocol:   6, // TCP
+			Fragment:   false,
+		}
+
+		selectedGw, ok := BalancePacket(&packet, gateways)
+		assert.True(t, ok)
+
+		switch selectedGw {
+		case gw1Addr:
+			gw1count += 1
+		case gw2Addr:
+			gw2count += 1
+		}
+
+	}
+
+	iterationCountAsFloat := float32(iterationCount)
+
+	assert.InDeltaf(t, iterationCountAsFloat*(2.0/3.0), gw1count, 100, "Expected %d +/- 100, but got %d", iterationCountAsFloat*(2.0/3.0), gw1count)
+	assert.InDeltaf(t, iterationCountAsFloat*(1.0/3.0), gw2count, 100, "Expected %d +/- 100, but got %d", iterationCountAsFloat*(1.0/3.0), gw2count)
+}
+
+func TestBalancePacketDistributsRandomlyAndReturnsFalseIfBucketsNotCalculated(t *testing.T) {
+	gateways := []Gateway{}
+
+	gw1Addr := netip.MustParseAddr("1.0.0.1")
+	gw2Addr := netip.MustParseAddr("1.0.0.2")
+
+	gateways = append(gateways, NewGateway(gw1Addr, 10))
+	gateways = append(gateways, NewGateway(gw2Addr, 5))
+
+	iterationCount := uint16(65535)
+	gw1count := 0
+	gw2count := 0
+
+	for i := uint16(0); i < iterationCount; i++ {
+		packet := firewall.Packet{
+			LocalAddr:  netip.MustParseAddr("192.168.1.1"),
+			RemoteAddr: netip.MustParseAddr("10.0.0.1"),
+			LocalPort:  i,
+			RemotePort: 65535 - i,
+			Protocol:   6, // TCP
+			Fragment:   false,
+		}
+
+		selectedGw, ok := BalancePacket(&packet, gateways)
+		assert.False(t, ok)
+
+		switch selectedGw {
+		case gw1Addr:
+			gw1count += 1
+		case gw2Addr:
+			gw2count += 1
+		}
+
+	}
+
+	assert.Equal(t, int(iterationCount), (gw1count + gw2count))
+	assert.NotEqual(t, 0, gw1count)
+	assert.NotEqual(t, 0, gw2count)
+
+}

routing/gateway.go

@@ -0,0 +1,70 @@
+package routing
+
+import (
+	"fmt"
+	"net/netip"
+)
+
+const (
+	// Sentinal value
+	BucketNotCalculated = -1
+)
+
+type Gateways []Gateway
+
+func (g Gateways) String() string {
+	str := ""
+	for i, gw := range g {
+		str += gw.String()
+		if i < len(g)-1 {
+			str += ", "
+		}
+	}
+	return str
+}
+
+type Gateway struct {
+	addr             netip.Addr
+	weight           int
+	bucketUpperBound int
+}
+
+func NewGateway(addr netip.Addr, weight int) Gateway {
+	return Gateway{addr: addr, weight: weight, bucketUpperBound: BucketNotCalculated}
+}
+
+func (g *Gateway) BucketUpperBound() int {
+	return g.bucketUpperBound
+}
+
+func (g *Gateway) Addr() netip.Addr {
+	return g.addr
+}
+
+func (g *Gateway) String() string {
+	return fmt.Sprintf("{addr: %s, weight: %d}", g.addr, g.weight)
+}
+
+// Divide and round to nearest integer
+func divideAndRound(v uint64, d uint64) uint64 {
+	var tmp uint64 = v + d/2
+	return tmp / d
+}
+
+// Implements Hash-Threshold mapping, equivalent to the implementation in the linux kernel.
+// After this function returns each gateway will have a
+// positive bucketUpperBound with a maximum value of 2147483647 (INT_MAX)
+func CalculateBucketsForGateways(gateways []Gateway) {
+
+	var totalWeight int = 0
+	for i := range gateways {
+		totalWeight += gateways[i].weight
+	}
+
+	var loopWeight int = 0
+	for i := range gateways {
+		loopWeight += gateways[i].weight
+		gateways[i].bucketUpperBound = int(divideAndRound(uint64(loopWeight)<<31, uint64(totalWeight))) - 1
+	}
+
+}

routing/gateway_test.go

@@ -0,0 +1,34 @@
+package routing
+
+import (
+	"net/netip"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRebalance3_2Split(t *testing.T) {
+	gateways := []Gateway{}
+
+	gateways = append(gateways, Gateway{addr: netip.Addr{}, weight: 10})
+	gateways = append(gateways, Gateway{addr: netip.Addr{}, weight: 5})
+
+	CalculateBucketsForGateways(gateways)
+
+	assert.Equal(t, 1431655764, gateways[0].bucketUpperBound) // INT_MAX/3*2
+	assert.Equal(t, 2147483647, gateways[1].bucketUpperBound) // INT_MAX
+}
+
+func TestRebalanceEqualSplit(t *testing.T) {
+	gateways := []Gateway{}
+
+	gateways = append(gateways, Gateway{addr: netip.Addr{}, weight: 1})
+	gateways = append(gateways, Gateway{addr: netip.Addr{}, weight: 1})
+	gateways = append(gateways, Gateway{addr: netip.Addr{}, weight: 1})
+
+	CalculateBucketsForGateways(gateways)
+
+	assert.Equal(t, 715827882, gateways[0].bucketUpperBound)  // INT_MAX/3
+	assert.Equal(t, 1431655764, gateways[1].bucketUpperBound) // INT_MAX/3*2
+	assert.Equal(t, 2147483647, gateways[2].bucketUpperBound) // INT_MAX
+}