tstest/natlab: add a stateful firewall.

The firewall provides a ProcessPacket handler, and implements an
address-and-port endpoint dependent firewall that allows all
traffic to egress from the trusted interface, and only allows
inbound traffic if corresponding outbound traffic was previously
seen.

Signed-off-by: David Anderson <danderson@tailscale.com>

tstest/clock.go

@@ -30,16 +30,27 @@ type Clock struct {
 func (c *Clock) Now() time.Time {
 	c.Lock()
 	defer c.Unlock()
+	c.initLocked()
+	step := c.Step
+	ret := c.Present
+	c.Present = c.Present.Add(step)
+	return ret
+}
+
+func (c *Clock) Advance(d time.Duration) {
+	c.Lock()
+	defer c.Unlock()
+	c.initLocked()
+	c.Present = c.Present.Add(d)
+}
+
+func (c *Clock) initLocked() {
 	if c.Start.IsZero() {
 		c.Start = time.Now()
 	}
 	if c.Present.Before(c.Start) {
 		c.Present = c.Start
 	}
-	step := c.Step
-	ret := c.Present
-	c.Present = c.Present.Add(step)
-	return ret
 }
 
 // Reset rewinds the virtual clock to its start time.

tstest/natlab/firewall.go

@@ -0,0 +1,76 @@
+// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package natlab
+
+import (
+	"sync"
+	"time"
+
+	"inet.af/netaddr"
+)
+
+type session struct {
+	src netaddr.IPPort
+	dst netaddr.IPPort
+}
+
+type Firewall struct {
+	// TrustedInterface is the interface that's allowed to send
+	// anywhere. All other interfaces can only respond to traffic from
+	// TrustedInterface.
+	TrustedInterface *Interface
+	// SessionTimeout is the lifetime of idle sessions in the firewall
+	// state. Packets transiting from the TrustedInterface reset the
+	// session lifetime to SessionTimeout.
+	SessionTimeout time.Duration
+	// TimeNow is a function returning the current time. If nil,
+	// time.Now is used.
+	TimeNow func() time.Time
+
+	// TODO: tuple-ness pickiness: EIF, ADF, APDF
+	// TODO: refresh directionality: outbound-only, both
+
+	mu   sync.Mutex
+	seen map[session]time.Time // session -> deadline
+}
+
+func (f *Firewall) timeNow() time.Time {
+	if f.TimeNow != nil {
+		return f.TimeNow()
+	}
+	return time.Now()
+}
+
+func (f *Firewall) HandlePacket(p []byte, inIf *Interface, dst, src netaddr.IPPort) PacketVerdict {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if f.seen == nil {
+		f.seen = map[session]time.Time{}
+	}
+
+	if inIf == f.TrustedInterface {
+		sess := session{
+			src: src,
+			dst: dst,
+		}
+		f.seen[sess] = f.timeNow().Add(f.SessionTimeout)
+		trace(p, "mach=%s iface=%s src=%s dst=%s firewall out ok", inIf.Machine().Name, inIf.name, src, dst)
+		return Continue
+	} else {
+		// reverse src and dst because the session table is from the
+		// POV of outbound packets.
+		sess := session{
+			src: dst,
+			dst: src,
+		}
+		now := f.timeNow()
+		if now.After(f.seen[sess]) {
+			trace(p, "mach=%s iface=%s src=%s dst=%s firewall drop", inIf.Machine().Name, inIf.name, src, dst)
+			return Drop
+		}
+		trace(p, "mach=%s iface=%s src=%s dst=%s firewall in ok", inIf.Machine().Name, inIf.name, src, dst)
+		return Continue
+	}
+}

tstest/natlab/natlab.go

@@ -166,6 +166,14 @@ type Interface struct {
 	ips     []netaddr.IP // static; not mutated once created
 }
 
+func (f *Interface) Machine() *Machine {
+	return f.machine
+}
+
+func (f *Interface) Network() *Network {
+	return f.net
+}
+
 // V4 returns the machine's first IPv4 address, or the zero value if none.
 func (f *Interface) V4() netaddr.IP { return f.pickIP(netaddr.IP.Is4) }
 

tstest/natlab/natlab_test.go

@@ -8,8 +8,10 @@
 	"context"
 	"fmt"
 	"testing"
+	"time"
 
 	"inet.af/netaddr"
+	"tailscale.com/tstest"
 )
 
 func TestAllocIPs(t *testing.T) {
@@ -217,5 +219,55 @@ func TestPacketHandler(t *testing.T) {
 	if addr.String() != mappedAddr.String() {
 		t.Errorf("addr = %q; want %q", addr, mappedAddr)
 	}
-
+}
+
+func TestFirewall(t *testing.T) {
+	clock := &tstest.Clock{}
+
+	wan := NewInternet()
+	lan := &Network{
+		Name:    "lan",
+		Prefix4: mustPrefix("10.0.0.0/8"),
+	}
+	m := &Machine{Name: "test"}
+	trust := m.Attach("trust", lan)
+	untrust := m.Attach("untrust", wan)
+
+	f := &Firewall{
+		TrustedInterface: trust,
+		SessionTimeout:   30 * time.Second,
+		TimeNow:          clock.Now,
+	}
+
+	client := ipp("192.168.0.2:1234")
+	serverA := ipp("2.2.2.2:5678")
+	serverB := ipp("7.7.7.7:9012")
+	tests := []struct {
+		iface    *Interface
+		src, dst netaddr.IPPort
+		want     PacketVerdict
+	}{
+		{trust, client, serverA, Continue},
+		{untrust, serverA, client, Continue},
+		{untrust, serverA, client, Continue},
+		{untrust, serverB, client, Drop},
+		{trust, client, serverB, Continue},
+		{untrust, serverB, client, Continue},
+	}
+
+	for _, test := range tests {
+		clock.Advance(time.Second)
+		got := f.HandlePacket(nil, test.iface, test.dst, test.src)
+		if got != test.want {
+			t.Errorf("iface=%s src=%s dst=%s got %v, want %v", test.iface.name, test.src, test.dst, got, test.want)
+		}
+	}
+}
+
+func ipp(str string) netaddr.IPPort {
+	ipp, err := netaddr.ParseIPPort(str)
+	if err != nil {
+		panic(err)
+	}
+	return ipp
 }