From 79e3a90569e3f7ae6843ae4993fb921fc29e1c88 Mon Sep 17 00:00:00 2001
From: Porjo <porjo38@yahoo.com.au>
Date: Fri, 27 Jun 2014 17:29:55 +1000
Subject: [PATCH] Move per-container forward rules to DOCKER chain

Docker-DCO-1.1-Signed-off-by: Ian Bishop <ianbishop@pace7.com> (github: porjo)
---
 iptables/iptables.go | 77 +++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 69 insertions(+), 8 deletions(-)

diff --git a/iptables/iptables.go b/iptables/iptables.go
index b783347..a7d216a 100644
--- a/iptables/iptables.go
+++ b/iptables/iptables.go
@@ -15,8 +15,9 @@ import (
 type Action string
 
 const (
-	Add    Action = "-A"
+	Append Action = "-A"
 	Delete Action = "-D"
+	Insert Action = "-I"
 )
 
 var (
@@ -54,10 +55,10 @@ func NewChain(name, bridge string) (*Chain, error) {
 		Bridge: bridge,
 	}
 
-	if err := chain.Prerouting(Add, "-m", "addrtype", "--dst-type", "LOCAL"); err != nil {
+	if err := chain.Prerouting(Append, "-m", "addrtype", "--dst-type", "LOCAL"); err != nil {
 		return nil, fmt.Errorf("Failed to inject docker in PREROUTING chain: %s", err)
 	}
-	if err := chain.Output(Add, "-m", "addrtype", "--dst-type", "LOCAL", "!", "--dst", "127.0.0.0/8"); err != nil {
+	if err := chain.Output(Append, "-m", "addrtype", "--dst-type", "LOCAL", "!", "--dst", "127.0.0.0/8"); err != nil {
 		return nil, fmt.Errorf("Failed to inject docker in OUTPUT chain: %s", err)
 	}
 	return chain, nil
@@ -78,7 +79,7 @@ func (c *Chain) Forward(action Action, ip net.IP, port int, proto, dest_addr str
 		// value" by both iptables and ip6tables.
 		daddr = "0/0"
 	}
-	if output, err := Raw("-t", "nat", fmt.Sprint(action), c.Name,
+	if output, err := Raw("-t", "nat", string(action), c.Name,
 		"-p", proto,
 		"-d", daddr,
 		"--dport", strconv.Itoa(port),
@@ -90,11 +91,13 @@ func (c *Chain) Forward(action Action, ip net.IP, port int, proto, dest_addr str
 		return &ChainError{Chain: "FORWARD", Output: output}
 	}
 
-	fAction := action
-	if fAction == Add {
-		fAction = "-I"
+	if action != Delete {
+		if err := c.createForwardChain(); err != nil {
+			return err
+		}
 	}
-	if output, err := Raw(string(fAction), "FORWARD",
+
+	if output, err := Raw(string(action), c.Name,
 		"!", "-i", c.Bridge,
 		"-o", c.Bridge,
 		"-p", proto,
@@ -109,6 +112,39 @@ func (c *Chain) Forward(action Action, ip net.IP, port int, proto, dest_addr str
 	return nil
 }
 
+func (c *Chain) Link(action Action, ip1, ip2 net.IP, port int, proto string) error {
+	if action != Delete {
+		if err := c.createForwardChain(); err != nil {
+			return err
+		}
+	}
+	if output, err := Raw(string(action), c.Name,
+		"-i", c.Bridge, "-o", c.Bridge,
+		"-p", proto,
+		"-s", ip1.String(),
+		"--dport", strconv.Itoa(port),
+		"-d", ip2.String(),
+		"-j", "ACCEPT"); err != nil {
+		return err
+	} else if len(output) != 0 {
+		return fmt.Errorf("Error toggle iptables forward: %s", output)
+	}
+
+	if output, err := Raw(string(action), c.Name,
+		"-i", c.Bridge, "-o", c.Bridge,
+		"-p", proto,
+		"-s", ip2.String(),
+		"--dport", strconv.Itoa(port),
+		"-d", ip1.String(),
+		"-j", "ACCEPT"); err != nil {
+		return err
+	} else if len(output) != 0 {
+		return fmt.Errorf("Error toggle iptables forward: %s", output)
+	}
+
+	return nil
+}
+
 func (c *Chain) Prerouting(action Action, args ...string) error {
 	a := append(nat, fmt.Sprint(action), "PREROUTING")
 	if len(args) > 0 {
@@ -199,3 +235,28 @@ func Raw(args ...string) ([]byte, error) {
 
 	return output, err
 }
+
+func (c *Chain) createForwardChain() error {
+	// Add chain if doesn't exist
+	if _, err := Raw("-n", "-L", c.Name); err != nil {
+		output, err := Raw("-N", c.Name)
+		if err != nil {
+			return err
+		} else if len(output) != 0 {
+			return fmt.Errorf("Error iptables forward: %s", output)
+		}
+	}
+	// Add linking rule if it doesn't exist
+	if !Exists("FORWARD",
+		"-o", c.Bridge,
+		"-j", c.Name) {
+		if output2, err := Raw(string(Insert), "FORWARD",
+			"-o", c.Bridge,
+			"-j", c.Name); err != nil {
+			return err
+		} else if len(output2) != 0 {
+			return fmt.Errorf("Error iptables forward: %s", output2)
+		}
+	}
+	return nil
+}