From d793677cd424fef10ac0b080871d181db0bcdec0 Mon Sep 17 00:00:00 2001
From: Kevin Krakauer <krakauer@google.com>
Date: Fri, 10 Jan 2020 18:07:15 -0800
Subject: [PATCH] I think INPUT works with protocol

---
 pkg/sentry/socket/netfilter/netfilter.go |  3 ++-
 pkg/tcpip/iptables/BUILD                 |  1 +
 pkg/tcpip/iptables/iptables.go           |  4 +++-
 pkg/tcpip/iptables/types.go              |  2 +-
 pkg/tcpip/network/ipv4/ipv4.go           |  3 ++-
 pkg/tcpip/packet_buffer.go               | 25 +++++++++++++++++++++++-
 6 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/pkg/sentry/socket/netfilter/netfilter.go b/pkg/sentry/socket/netfilter/netfilter.go
index f30461936..175466f19 100644
--- a/pkg/sentry/socket/netfilter/netfilter.go
+++ b/pkg/sentry/socket/netfilter/netfilter.go
@@ -25,6 +25,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
 	"gvisor.dev/gvisor/pkg/sentry/usermem"
 	"gvisor.dev/gvisor/pkg/syserr"
+	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/iptables"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
@@ -455,7 +456,7 @@ func filterFromIPTIP(iptip linux.IPTIP) (iptables.IPHeaderFilter, *syserr.Error)
 		return iptables.IPHeaderFilter{}, syserr.ErrInvalidArgument
 	}
 	return iptables.IPHeaderFilter{
-		Protocol: iptip.Protocol,
+		Protocol: tcpip.TransportProtocolNumber(iptip.Protocol),
 	}, nil
 }
 
diff --git a/pkg/tcpip/iptables/BUILD b/pkg/tcpip/iptables/BUILD
index 2893c80cd..297eaccaf 100644
--- a/pkg/tcpip/iptables/BUILD
+++ b/pkg/tcpip/iptables/BUILD
@@ -14,5 +14,6 @@ go_library(
     deps = [
         "//pkg/log",
         "//pkg/tcpip",
+        "//pkg/tcpip/header",
     ],
 )
diff --git a/pkg/tcpip/iptables/iptables.go b/pkg/tcpip/iptables/iptables.go
index 4e1700fdb..3cff879a2 100644
--- a/pkg/tcpip/iptables/iptables.go
+++ b/pkg/tcpip/iptables/iptables.go
@@ -21,6 +21,7 @@ import (
 
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/header"
 )
 
 const (
@@ -183,12 +184,13 @@ func (it *IPTables) checkTable(hook Hook, pkt tcpip.PacketBuffer, tablename stri
 	panic(fmt.Sprintf("Traversed past the entire list of iptables rules in table %q.", tablename))
 }
 
+// Precondition: pk.NetworkHeader is set.
 func (it *IPTables) checkRule(hook Hook, pkt tcpip.PacketBuffer, table Table, ruleIdx int) Verdict {
 	rule := table.Rules[ruleIdx]
 
 	// First check whether the packet matches the IP header filter.
 	// TODO(gvisor.dev/issue/170): Support other fields of the filter.
-	if rule.Filter.Protocol != pkt.Protocol {
+	if rule.Filter.Protocol != header.IPv4(pkt.NetworkHeader).TransportProtocol() {
 		return Continue
 	}
 
diff --git a/pkg/tcpip/iptables/types.go b/pkg/tcpip/iptables/types.go
index 4bedd9bc8..4f2a4d65e 100644
--- a/pkg/tcpip/iptables/types.go
+++ b/pkg/tcpip/iptables/types.go
@@ -173,7 +173,7 @@ type IPHeaderFilter struct {
 	InputInterface      string
 	OutputInterfaceMask string
 	InputInterfaceMask  string
-	Protocol            uint16
+	Protocol            tcpip.TransportProtocolNumber
 	Flags               uint8
 	InverseFlags        uint8
 }
diff --git a/pkg/tcpip/network/ipv4/ipv4.go b/pkg/tcpip/network/ipv4/ipv4.go
index f856081e6..5388d2549 100644
--- a/pkg/tcpip/network/ipv4/ipv4.go
+++ b/pkg/tcpip/network/ipv4/ipv4.go
@@ -353,7 +353,8 @@ func (e *endpoint) HandlePacket(r *stack.Route, pkt tcpip.PacketBuffer) {
 	}
 	pkt.NetworkHeader = headerView[:h.HeaderLength()]
 
-	// iptables filtering.
+	// iptables filtering. All packets that reach here are intended for
+	// this machine and will not be forwarded.
 	ipt := e.stack.IPTables()
 	if ok := ipt.Check(iptables.Input, pkt); !ok {
 		// iptables is telling us to drop the packet.
diff --git a/pkg/tcpip/packet_buffer.go b/pkg/tcpip/packet_buffer.go
index ab24372e7..7a036b93c 100644
--- a/pkg/tcpip/packet_buffer.go
+++ b/pkg/tcpip/packet_buffer.go
@@ -13,7 +13,9 @@
 
 package tcpip
 
-import "gvisor.dev/gvisor/pkg/tcpip/buffer"
+import (
+	"gvisor.dev/gvisor/pkg/tcpip/buffer"
+)
 
 // A PacketBuffer contains all the data of a network packet.
 //
@@ -65,3 +67,24 @@ func (pk PacketBuffer) Clone() PacketBuffer {
 	pk.Data = pk.Data.Clone(nil)
 	return pk
 }
+
+//// TransportProtocol returns the transport protocol of pk.
+////
+//// Precondition: pk.NetworkHeader is set.
+//func (pk PacketBuffer) TransportProtocolIPv4() uint16 {
+//	if pk.NetworkHeader == nil {
+//		panic("This should only be called when pk.NetworkHeader is set.")
+//	}
+//	return header.IPv4(pk.NetworkHeader).TransportProtocol()
+//}
+
+// func (pk Packet) findNetHeader() header.IPv4 {
+// 	// Inbound:
+// 	// Data holds everything, but may have had some headers shaved off.
+// 	// Figure out whether it's set or still somewhere in data and return
+// 	// appropriately.
+
+// 	// Outbound:
+// 	// NetworkHeader will be set if we've added one. Otherwise there's no
+// 	// header.
+// }