Merge pull request #16 from Shopify/add-udp-balancing

Adding load balancing between ranges for UDP
Shopify · Jul 4, 2024 · eafd49e · eafd49e
2 parents f4ddc6f + 4a3dbd3
commit eafd49e
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 11 deletions.
diff --git a/cmd/ip-masq-agent/ip-masq-agent.go b/cmd/ip-masq-agent/ip-masq-agent.go
@@ -388,16 +388,22 @@ func writeNonMasqRule(lines *bytes.Buffer, cidr string) {
 const masqRuleComment = `-m comment --comment "ip-masq-agent: outbound traffic is subject to MASQUERADE (must be last in chain)"`
 
 func writeTcpMasqRule(lines *bytes.Buffer) {
-    args1 := []string{masqRuleComment, "-j", "MASQUERADE", "-p", "tcp", "-m", "statistic", "--mode", "random", "--probability", "0.5", "--to-ports", "1024-29999"}
-    args2 := []string{masqRuleComment, "-j", "MASQUERADE", "-p", "tcp", "--to-ports", "32768-65535"}
+	args1 := []string{masqRuleComment, "-j", "MASQUERADE", "-p", "tcp", "-m", "statistic", "--mode", "random", "--probability", "0.5", "--to-ports", "1024-29999"}
+	args2 := []string{masqRuleComment, "-j", "MASQUERADE", "-p", "tcp", "--to-ports", "32768-65535"}
+	args3 := []string{masqRuleComment, "-j", "MASQUERADE", "-p", "udp", "-m", "statistic", "--mode", "random", "--probability", "0.5", "--to-ports", "1024-29999"}
+	args4 := []string{masqRuleComment, "-j", "MASQUERADE", "-p", "udp", "--to-ports", "32768-65535"}
 
-    if *randomFully {
-        args1 = append(args1, "--random-fully")
-        args2 = append(args2, "--random-fully")
-    }
+	if *randomFully {
+		args1 = append(args1, "--random-fully")
+		args2 = append(args2, "--random-fully")
+		args3 = append(args3, "--random-fully")
+		args4 = append(args4, "--random-fully")
+	}
 
-    writeRule(lines, utiliptables.Append, masqChain, args1...)
-    writeRule(lines, utiliptables.Append, masqChain, args2...)
+	writeRule(lines, utiliptables.Append, masqChain, args1...)
+	writeRule(lines, utiliptables.Append, masqChain, args2...)
+	writeRule(lines, utiliptables.Append, masqChain, args3...)
+	writeRule(lines, utiliptables.Append, masqChain, args4...)
 }
 
 func writeMasqRule(lines *bytes.Buffer) {

diff --git a/cmd/ip-masq-agent/ip-masq-agent_test.go b/cmd/ip-masq-agent/ip-masq-agent_test.go
@@ -41,8 +41,8 @@ func TestMain(m *testing.M) {
 	ec := 0
 	randomFully := " --random-fully"
 
-	for _, tc := range []struct{
-		arg string
+	for _, tc := range []struct {
+		arg  string
 		want string
 	}{
 		{
@@ -52,7 +52,7 @@ func TestMain(m *testing.M) {
 			arg: "false",
 		},
 		{
-			arg: "true",
+			arg:  "true",
 			want: randomFully,
 		},
 	} {
@@ -312,6 +312,8 @@ func TestSyncMasqRules(t *testing.T) {
 -A ` + string(masqChain) + ` ` + nonMasqRuleComment + ` -d 169.254.0.0/16 -j RETURN
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp --to-ports 32768-65535` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp --to-ports 32768-65535` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE` + wantRandomFully + `
 COMMIT
 `,
@@ -330,6 +332,8 @@ COMMIT
 -A ` + string(masqChain) + ` ` + nonMasqRuleComment + ` -d 192.168.0.0/16 -j RETURN
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp --to-ports 32768-65535` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp --to-ports 32768-65535` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE` + wantRandomFully + `
 COMMIT
 `,
@@ -356,6 +360,8 @@ COMMIT
 -A ` + string(masqChain) + ` ` + nonMasqRuleComment + ` -d 240.0.0.0/4 -j RETURN
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp --to-ports 32768-65535` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp --to-ports 32768-65535` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE` + wantRandomFully + `
 COMMIT
 `,
@@ -377,6 +383,8 @@ COMMIT
 -A ` + string(masqChain) + ` ` + nonMasqRuleComment + ` -d 10.244.0.0/16 -j RETURN
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp --to-ports 32768-65535` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp --to-ports 32768-65535` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE` + wantRandomFully + `
 COMMIT
 `,
@@ -421,6 +429,8 @@ func TestSyncMasqRulesIPv6(t *testing.T) {
 -A ` + string(masqChain) + ` ` + nonMasqRuleComment + ` -d fe80::/10 -j RETURN
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp --to-ports 32768-65535` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp --to-ports 32768-65535` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE` + wantRandomFully + `
 COMMIT
 `,
@@ -442,6 +452,8 @@ COMMIT
 -A ` + string(masqChain) + ` ` + nonMasqRuleComment + ` -d fc00::/7 -j RETURN
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp --to-ports 32768-65535` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp --to-ports 32768-65535` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE` + wantRandomFully + `
 COMMIT
 `,
@@ -456,6 +468,8 @@ COMMIT
 				fmt.Sprintf(postRoutingMasqChainCommentFormat, masqChain) + ` -m addrtype ! --dst-type LOCAL -j ` + string(masqChain) + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p tcp --to-ports 32768-65535` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp -m statistic --mode random --probability 0.5 --to-ports 1024-29999` + wantRandomFully + `
+-A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE -p udp --to-ports 32768-65535` + wantRandomFully + `
 -A ` + string(masqChain) + ` ` + masqRuleComment + ` -j MASQUERADE` + wantRandomFully + `
 COMMIT
 `,