Refactored ebtables to remove redundant code

Author: ofiliz

ebtables/ebtables.go

@@ -6,16 +6,18 @@ package ebtables
 import (
 	"fmt"
 	"io/ioutil"
+	"net"
 	"os/exec"
 	"strings"
 
 	"github.com/Azure/azure-container-networking/log"
 )
 
-// Init initializes the ebtables module.
-func init() {
-	installEbtables()
-}
+const (
+	// Ebtables actions.
+	Append = "-A"
+	Delete = "-D"
+)
 
 // InstallEbtables installs the ebtables package.
 func installEbtables() {
@@ -31,64 +33,31 @@ func installEbtables() {
 	}
 }
 
-// SetupSnatForOutgoingPackets sets up snat
-func SetupSnatForOutgoingPackets(interfaceName string, snatAddress string) error {
-	command := fmt.Sprintf("ebtables -t nat -A POSTROUTING -o %s -j snat --to-source %s --snat-arp", interfaceName, snatAddress)
-	err := executeShellCommand(command)
-	if err != nil {
-		return err
-	}
-	return nil
-}
+// SetSnatForInterface sets a MAC SNAT rule for an interface.
+func SetSnatForInterface(interfaceName string, macAddress net.HardwareAddr, action string) error {
+	command := fmt.Sprintf(
+		"ebtables -t nat %s POSTROUTING -o %s -j snat --to-src %s --snat-arp",
+		action, interfaceName, macAddress.String())
 
-// CleanupSnatForOutgoingPackets cleans up snat
-func CleanupSnatForOutgoingPackets(interfaceName string, snatAddress string) error {
-	command := fmt.Sprintf("ebtables -t nat -D POSTROUTING -o %s -j snat --to-source %s --snat-arp", interfaceName, snatAddress)
-	err := executeShellCommand(command)
-	if err != nil {
-		return err
-	}
-	return nil
+	return executeShellCommand(command)
 }
 
-// SetupDnatForArpReplies sets up dnat
-func SetupDnatForArpReplies(interfaceName string) error {
-	command := fmt.Sprintf("ebtables -t nat -A PREROUTING -i %s -p arp -j dnat --to-destination ff:ff:ff:ff:ff:ff", interfaceName)
-	err := executeShellCommand(command)
-	if err != nil {
-		return err
-	}
-	return nil
-}
+// SetDnatForArpReplies sets a MAC DNAT rule for ARP replies received on an interface.
+func SetDnatForArpReplies(interfaceName string, action string) error {
+	command := fmt.Sprintf(
+		"ebtables -t nat %s PREROUTING -p ARP -i %s -j dnat --to-dst ff:ff:ff:ff:ff:ff",
+		action, interfaceName)
 
-// CleanupDnatForArpReplies cleans up dnat
-func CleanupDnatForArpReplies(interfaceName string) error {
-	command := fmt.Sprintf("ebtables -t nat -D PREROUTING -i %s -p arp -j dnat --to-destination ff:ff:ff:ff:ff:ff", interfaceName)
-	err := executeShellCommand(command)
-	if err != nil {
-		return err
-	}
-	return nil
+	return executeShellCommand(command)
 }
 
-// SetupDnatBasedOnIPV4Address sets up dnat
-func SetupDnatBasedOnIPV4Address(ipv4Address string, macAddress string) error {
-	command := fmt.Sprintf("ebtables -t nat -A PREROUTING -p IPv4  --ip-dst %s -j dnat --to-dst %s --dnat-target ACCEPT", ipv4Address, macAddress)
-	err := executeShellCommand(command)
-	if err != nil {
-		return err
-	}
-	return nil
-}
+// SetDnatForIPAddress sets a MAC DNAT rule for an IP address.
+func SetDnatForIPAddress(ipAddress net.IP, macAddress net.HardwareAddr, action string) error {
+	command := fmt.Sprintf(
+		"ebtables -t nat %s PREROUTING -p IPv4 --ip-dst %s -j dnat --to-dst %s",
+		action, ipAddress.String(), macAddress.String())
 
-// RemoveDnatBasedOnIPV4Address cleans up dnat
-func RemoveDnatBasedOnIPV4Address(ipv4Address string, macAddress string) error {
-	command := fmt.Sprintf("ebtables -t nat -D PREROUTING -p IPv4  --ip-dst %s -j dnat --to-dst %s --dnat-target ACCEPT", ipv4Address, macAddress)
-	err := executeShellCommand(command)
-	if err != nil {
-		return err
-	}
-	return nil
+	return executeShellCommand(command)
 }
 
 func executeShellCommand(command string) error {

network/endpoint.go

@@ -101,7 +101,7 @@ func (nw *network) newEndpoint(epInfo *EndpointInfo) (*endpoint, error) {
 	// Setup MAC address translation rules for container interface.
 	log.Printf("[net] Setting up MAC address translation rules for endpoint %v.", contIfName)
 	for _, ipAddr := range epInfo.IPAddresses {
-		err = ebtables.SetupDnatBasedOnIPV4Address(ipAddr.IP.String(), containerIf.HardwareAddr.String())
+		err = ebtables.SetDnatForIPAddress(ipAddr.IP, containerIf.HardwareAddr, ebtables.Append)
 		if err != nil {
 			goto cleanup
 		}
@@ -240,7 +240,7 @@ func (nw *network) deleteEndpoint(endpointId string) error {
 	// Delete MAC address translation rule.
 	log.Printf("[net] Deleting MAC address translation rules for endpoint %v.", endpointId)
 	for _, ipAddr := range ep.IPAddresses {
-		err = ebtables.RemoveDnatBasedOnIPV4Address(ipAddr.IP.String(), ep.MacAddress.String())
+		err = ebtables.SetDnatForIPAddress(ipAddr.IP, ep.MacAddress, ebtables.Delete)
 		if err != nil {
 			goto cleanup
 		}

network/network.go

@@ -172,12 +172,12 @@ func (nm *networkManager) connectExternalInterface(extIf *externalInterface, bri
 
 	// Setup MAC address translation rules for external interface.
 	log.Printf("[net] Setting up MAC address translation rules for %v.", hostIf.Name)
-	err = ebtables.SetupSnatForOutgoingPackets(hostIf.Name, hostIf.HardwareAddr.String())
+	err = ebtables.SetSnatForInterface(hostIf.Name, hostIf.HardwareAddr, ebtables.Append)
 	if err != nil {
 		goto cleanup
 	}
 
-	err = ebtables.SetupDnatForArpReplies(hostIf.Name)
+	err = ebtables.SetDnatForArpReplies(hostIf.Name, ebtables.Append)
 	if err != nil {
 		goto cleanup
 	}
@@ -240,8 +240,8 @@ cleanup:
 	log.Printf("[net] Connecting interface %v failed, err:%v.", extIf.Name, err)
 
 	// Roll back the changes for the network.
-	ebtables.CleanupDnatForArpReplies(extIf.Name)
-	ebtables.CleanupSnatForOutgoingPackets(extIf.Name, extIf.MacAddress.String())
+	ebtables.SetDnatForArpReplies(extIf.Name, ebtables.Delete)
+	ebtables.SetSnatForInterface(extIf.Name, extIf.MacAddress, ebtables.Delete)
 
 	netlink.DeleteLink(bridgeName)
 
@@ -253,8 +253,8 @@ func (nm *networkManager) disconnectExternalInterface(extIf *externalInterface)
 	log.Printf("[net] Disconnecting interface %v.", extIf.Name)
 
 	// Cleanup MAC address translation rules.
-	ebtables.CleanupDnatForArpReplies(extIf.Name)
-	ebtables.CleanupSnatForOutgoingPackets(extIf.Name, extIf.MacAddress.String())
+	ebtables.SetDnatForArpReplies(extIf.Name, ebtables.Delete)
+	ebtables.SetSnatForInterface(extIf.Name, extIf.MacAddress, ebtables.Delete)
 
 	// Disconnect external interface from its bridge.
 	err := netlink.SetLinkMaster(extIf.Name, "")