Merge lp:~themue/juju-core/go-cmd-firewall into lp:~juju/juju-core/trunk

[Frank Mueller]

cmd: added port of state/firewall.py

The port of the current Python firewall has been moved to
the cmd package because it's only used by the provisioning
agent. The implementation is split into several subtypes
and done, the tests are done for services.

THIS PROPOSAL IS STILL EVOLVING.

https://codereview.appspot.com/6374047/

[Roger Peppe (rogpeppe)]

some initial superficial comments.

i'm still mulling it over.

https://codereview.appspot.com/6374047/diff/1/cmd/firewall.go
File cmd/firewall.go (right):

https://codereview.appspot.com/6374047/diff/1/cmd/firewall.go#newcode1
cmd/firewall.go:1: package cmd
this is the wrong place for this code.
given that it's really just part of the provisioning agent, it should be
alongside the PA code, in cmd/jujud, or the provisioning package,
if/when that happens.

https://codereview.appspot.com/6374047/diff/1/cmd/firewall.go#newcode14
cmd/firewall.go:14: type Ports []state.Port
i'm not sure that this type or its methods deserve export.
not exporting the type means that environs implementors don't
necessarily need to depend on this package.

since we're talking about a set of ports, i wonder if
map[state.Port]bool might be a nicer representation.

https://codereview.appspot.com/6374047/diff/1/cmd/firewall.go#newcode55
cmd/firewall.go:55: OpenPorts(instanceId string, machineId int) (Ports,
error)
i wonder if just "Ports" might be better here.
the name OpenPorts seems to suggest it's just a version of OpenPort that
opens more than one port.

https://codereview.appspot.com/6374047/diff/1/cmd/firewall.go#newcode106
cmd/firewall.go:106: // startd the watchers of the services exposed
flag.
s/startd/starts/

https://codereview.appspot.com/6374047/diff/1/cmd/firewall.go#newcode159
cmd/firewall.go:159: func (fw *Firewall)
AddOpenClosePortsOnMachineObserver(observer func(int)) {
i'm not sure i understand the use-case for this function. isn't
observing ports exactly what all this code is about? why do we need
another hook? if it's about testing, i'd have thought that the
PortsManager interface should give us a sufficient interface.

https://codereview.appspot.com/6374047/diff/1/cmd/firewall.go#newcode280
cmd/firewall.go:280: type serviceExposedWatcher struct {
see comment on newServiceUnitsWatcher below.

serviceExposedPortManager, perhaps?

https://codereview.appspot.com/6374047/diff/1/cmd/firewall.go#newcode332
cmd/firewall.go:332: }
i think this may read better:

sew.exposed = exposed
if exposed {
     sew.suw = newServiceUnitsWatcher(sew.fw, sew)
}

https://codereview.appspot.com/6374047/diff/1/cmd/firewall.go#newcode360
cmd/firewall.go:360: // the passed service.
this isn't really a watcher in the sense of the other watchers in state.
more of a serviceUnitPortManager, perhaps?

https://codereview.appspot.com/6374047/diff/1/cmd/firewall.go#newcode414
cmd/firewall.go:414: type unitPortsWatcher struct {
unitPortManager?

https://codereview.appspot.com/6374047/

[Gustavo Niemeyer (niemeyer)]

Frank, I'm sure a lot of what is in this branch will be useful, but this
is not the way we should be pushing this kind of feature. 700 lines of
code, apparently in a location that, agreeing with roger, is not the
right one, with 5 lines of tests.

I understand you say it's an evolving proposal, but we've covered this
kind of thing before, and it's going to be painful again but I'll have
to ask you to chunk this down into smaller portions that we can more
easily agree on.

I'm happy to have a run down with you tomorrow to see how to best chunk
this up, and what is the first piece for the puzzle.

https://codereview.appspot.com/6374047/

Unmerged revisions

266. By Frank Mueller on 2012-07-10

cmd: added port of state/firewall.py

265. By Frank Mueller on 2012-07-10

state: preparing for trunk merging

264. By Frank Mueller on 2012-07-09

cmd: prepared firewall testing

263. By Frank Mueller on 2012-07-06

cmd: added initial version of the firewall

262. By Dave Cheney on 2012-06-25

state: machine.InstanceId error proposal

machine.InstanceId returns NoInstanceIdError when
there is no instance id set in the machine *ConfigNode.

R=niemeyer
CC=
https://codereview.appspot.com/6297105

261. By Gustavo Niemeyer on 2012-06-25

mstate: update to the mgo API version r2012.06.22

This was already reviewed in:

https://code.launchpad.net/~niemeyer/juju-core/store-update-mgo/+merge/111546

But Launchpad simply can't serve that branch without errors anymore. :-(

R=
CC=
https://codereview.appspot.com/6336052

260. By Gustavo Niemeyer on 2012-06-25

store: update mgo API to r2012.06.22

R=rog
CC=
https://codereview.appspot.com/6327049

259. By Frank Mueller on 2012-06-22

state: adding of FlagWatcher and integration into Service

The port of Firewall (in a different CL) needs the implementation
of the FlagWatcher and its integration into Service to watch the
exposed flag. Both is done in this CL.

R=niemeyer
CC=
https://codereview.appspot.com/6332048

258. By Roger Peppe on 2012-06-22

cmd/jujud: make provisioner use StateInfo from environment.

This means we can pass in a localhost address to the provisioning
agent and it will nonetheless use the correct address for zookeeper
when firing up clients.

Some rearrangement of testing code was necessary too.

R=niemeyer
CC=
https://codereview.appspot.com/6336047

257. By Dave Cheney on 2012-06-22

environs/dummy: enhance start/stop operations

This is an alternate proposal to enhance dummy to return
the actual environ.Instance objects created during
Start/StopInstance operations.

This would be used in the PA tests to double check that
m, err := s.st.AddMachine(n); m.InstanceId matches the
Instance.Id() provided by dummy.

R=niemeyer
CC=
https://codereview.appspot.com/6304104

=== modified file 'cmd/cmd_test.go'
--- cmd/cmd_test.go	2012-06-21 20:40:39 +0000
+++ cmd/cmd_test.go	2012-07-10 15:08:35 +0000
@@ -3,11 +3,14 @@
 import (
 	. "launchpad.net/gocheck"
 	"launchpad.net/juju-core/cmd"
+	"launchpad.net/juju-core/testing"
 	"path/filepath"
-	"testing"
+	stdtesting "testing"
 )
 
-func Test(t *testing.T) { TestingT(t) }
+func Test(t *stdtesting.T) {
+	testing.ZkTestPackage(t)
+}
 
 type CmdSuite struct{}
 

=== added file 'cmd/firewall.go'
--- cmd/firewall.go	1970-01-01 00:00:00 +0000
+++ cmd/firewall.go	2012-07-10 15:08:35 +0000
@@ -0,0 +1,472 @@
+package cmd
+
+import (
+	"errors"
+	"fmt"
+	"launchpad.net/juju-core/environs"
+	"launchpad.net/juju-core/log"
+	"launchpad.net/juju-core/state"
+	"launchpad.net/juju-core/state/watcher"
+	"launchpad.net/tomb"
+)
+
+// Ports is a list of ports.
+type Ports []state.Port
+
+// Missing returns the ports that exist in p but not in other.
+func (p Ports) Missing(other Ports) (missing Ports) {
+next:
+	for _, portA := range p {
+		for _, portB := range other {
+			if portA.Number == portB.Number && portA.Protocol == portB.Protocol {
+				continue next
+			}
+		}
+		missing = append(missing, portA)
+	}
+	return
+}
+
+// Len is the number of ports.
+func (p Ports) Len() int {
+	return len(p)
+}
+
+// Less returns whether the port with index i should sort
+// before the port with index j.
+func (p Ports) Less(i, j int) bool {
+	if p[i].Protocol == p[j].Protocol {
+		return p[i].Number < p[j].Number
+	}
+	return p[i].Protocol < p[j].Protocol
+}
+
+// Swap swaps the ports with indexes i and j.
+func (p Ports) Swap(i, j int) {
+	p[i], p[j] = p[j], p[i]
+}
+
+// PortManager defines the interface for the component
+// responsible for the management of the ports on the
+// instances and machines.
+type PortManager interface {
+	// OpenPorts returns the current open ports of the 
+	// machine on the instance.
+	OpenPorts(instanceId string, machineId int) (Ports, error)
+	// OpenPort opens the passed port on the machine and instance.
+	OpenPort(instanceId string, machineId int, port state.Port) error
+	// ClosePort closes the passed port on the machine and instance.
+	ClosePort(instanceId string, machineId int, port state.Port) error
+}
+
+// Firewall manages the opening and closing of ports in the firewall.
+type Firewall struct {
+	state                            *state.State
+	isRunning                        func() bool
+	portManager                      PortManager
+	watchedMachines                  map[int]*machineUnitsWatcher
+	watchedServices                  map[string]*serviceExposedWatcher
+	retryMachinesOnPortError         map[int]bool
+	openClosePortsObservers          []func(*state.Unit)
+	openClosePortsOnMachineObservers []func(int)
+}
+
+// NewFirewall creates a new firewall. is Running is a function
+// that returns whether the associated agent is running or not.
+func NewFirewall(st *state.State, isRunning func() bool, portManager PortManager) *Firewall {
+	return &Firewall{
+		state:                            st,
+		isRunning:                        isRunning,
+		portManager:                      portManager,
+		watchedMachines:                  make(map[int]*machineUnitsWatcher),
+		watchedServices:                  make(map[string]*serviceExposedWatcher),
+		retryMachinesOnPortError:         make(map[int]bool),
+		openClosePortsObservers:          make([]func(*state.Unit), 0),
+		openClosePortsOnMachineObservers: make([]func(int), 0),
+	}
+}
+
+// ProcessMachine ensures that a watch is setup per machine and
+// performs any necessary retry.
+func (fw *Firewall) ProcessMachine(machine *state.Machine) error {
+	machineId := machine.Id()
+	if fw.retryMachinesOnPortError[machineId] {
+		delete(fw.retryMachinesOnPortError, machineId)
+		if err := fw.OpenClosePortsOnMachine(machineId); err != nil {
+			return err
+		}
+	}
+	if fw.watchedMachines[machineId] == nil {
+		fw.watchedMachines[machineId] = newMachineUnitsWatcher(fw, machineId)
+	}
+	return nil
+}
+
+// WatchServiceChanges is called when services are added or removed and
+// startd the watchers of the services exposed flag.
+func (fw *Firewall) WatchServiceChanges(change *state.ServicesChange) error {
+	for _, service := range change.Removed {
+		sew, ok := fw.watchedServices[service.Name()]
+		if ok {
+			if err := sew.stop(); err != nil {
+				return err
+			}
+			delete(fw.watchedServices, service.Name())
+		}
+	}
+	for _, service := range change.Added {
+		if _, ok := fw.watchedServices[service.Name()]; ok {
+			// Already watched.
+			continue
+		}
+		fw.watchedServices[service.Name()] = newServiceExposedWatcher(fw, service)
+	}
+	return nil
+}
+
+// AddOpenClosePortsObserver adds an observer for the opening
+// or closing of ports for a service unit.
+func (fw *Firewall) AddOpenClosePortsObserver(observer func(*state.Unit)) {
+	fw.openClosePortsObservers = append(fw.openClosePortsObservers, observer)
+}
+
+// OpenClosePorts is called upon changes that *may* open/close 
+// ports for a service unit.
+func (fw *Firewall) OpenClosePorts(unit *state.Unit) (err error) {
+	defer errorContextf(&err, "can't open/close ports for unit %q", unit)
+	if !fw.isRunning() {
+		return errors.New("agent associated to firewall is not running")
+	}
+	defer func() {
+		// Ensure that the observations runs after the
+		// corresponding action completes. In particular, tests
+		// that use observation depend on this ordering to ensure
+		// that the action has in fact happened before they can
+		// proceed.
+		for _, observer := range fw.openClosePortsObservers {
+			observer(unit)
+		}
+	}()
+	machineId, err := unit.AssignedMachineId()
+	if err != nil {
+		return err
+	}
+	return fw.OpenClosePortsOnMachine(machineId)
+}
+
+// AddOpenClosePortsOnMachineObserver adds an observer for the opening
+// or closing of ports on a machine.
+func (fw *Firewall) AddOpenClosePortsOnMachineObserver(observer func(int)) {
+	fw.openClosePortsOnMachineObservers = append(fw.openClosePortsOnMachineObservers, observer)
+}
+
+// OpenClosePortsOnMachine is called upon changes that *may* open/close
+// ports for a machine.
+func (fw *Firewall) OpenClosePortsOnMachine(machineId int) (err error) {
+	defer errorContextf(&err, "can't open/close ports on machine %d", machineId)
+	if !fw.isRunning() {
+		return errors.New("agent associated to firewall is not running")
+	}
+	defer func() {
+		if err == environs.ErrNoInstances {
+			log.Printf("no provisioned instance for machine %d", machineId)
+		} else if err != nil {
+			log.Printf("got error in opening/closing ports, will retry")
+			fw.retryMachinesOnPortError[machineId] = true
+		}
+		// Ensure that the observation runs after the corresponding
+		// action completes. In particular, tests that use
+		// observation depend on this ordering to ensure that this
+		// action has happened before they can proceed.
+		for _, observer := range fw.openClosePortsOnMachineObservers {
+			observer(machineId)
+		}
+
+	}()
+	machine, err := fw.state.Machine(machineId)
+	if err != nil {
+		return err
+	}
+	instanceId, err := machine.InstanceId()
+	if err != nil {
+		return err
+	}
+	units, err := machine.Units()
+	if err != nil {
+		return err
+	}
+	policyPorts := Ports{}
+	for _, unit := range units {
+		service, err := fw.state.Service(unit.ServiceName())
+		if err != nil {
+			return err
+		}
+		exposed, err := service.IsExposed()
+		if err != nil {
+			return err
+		}
+		if exposed {
+			ports, err := unit.OpenPorts()
+			if err != nil {
+				return err
+			}
+			policyPorts = append(policyPorts, ports...)
+		}
+	}
+	currentPorts, err := fw.portManager.OpenPorts(instanceId, machineId)
+	toOpen := policyPorts.Missing(currentPorts)
+	toClose := currentPorts.Missing(policyPorts)
+	for _, port := range toOpen {
+		if err := fw.portManager.OpenPort(instanceId, machineId, port); err != nil {
+			return err
+		}
+	}
+	for _, port := range toClose {
+		if err := fw.portManager.ClosePort(instanceId, machineId, port); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// machineUnitsWatcher manages the opening and closing of ports in the firewall
+// for one machine.
+type machineUnitsWatcher struct {
+	fw        *Firewall
+	machineId int
+	tomb      tomb.Tomb
+}
+
+// newMachineUnitsWatcher creates a new watcher for the passed machine.
+func newMachineUnitsWatcher(fw *Firewall, machineId int) *machineUnitsWatcher {
+	muw := &machineUnitsWatcher{
+		fw:        fw,
+		machineId: machineId,
+	}
+	go muw.loop()
+	return muw
+}
+
+// loop watches assigned units for changes possible require port management.
+func (muw *machineUnitsWatcher) loop() {
+	defer muw.tomb.Done()
+	machine, err := muw.fw.state.Machine(muw.machineId)
+	if err != nil {
+		muw.tomb.Kill(err)
+		return
+	}
+	unitsWatcher := machine.WatchUnits()
+	defer watcher.Stop(unitsWatcher, &muw.tomb)
+
+	for {
+		select {
+		case <-muw.tomb.Dying():
+			return
+		case _, ok := <-unitsWatcher.Changes():
+			if !ok {
+				muw.tomb.Kill(watcher.MustErr(unitsWatcher))
+				return
+			}
+			if err := muw.fw.OpenClosePortsOnMachine(muw.machineId); err != nil {
+				muw.tomb.Kill(err)
+				return
+			}
+		}
+	}
+}
+
+// serviceExposedWatcher manages the opening and closing of ports in the firewall
+// for a service.
+type serviceExposedWatcher struct {
+	fw      *Firewall
+	service *state.Service
+	suw     *serviceUnitsWatcher
+	exposed bool
+	tomb    tomb.Tomb
+}
+
+// newServiceExposedWatcher creates a new watcher for the passed service.
+func newServiceExposedWatcher(fw *Firewall, service *state.Service) *serviceExposedWatcher {
+	sew := &serviceExposedWatcher{
+		fw:      fw,
+		service: service,
+	}
+	go sew.loop()
+	return sew
+}
+
+// loop watches if the exposed flag of the service is set.
+func (sew *serviceExposedWatcher) loop() {
+	log.Debugf("start watching of %q on changes to being exposed", sew.service)
+	defer sew.tomb.Done()
+	exposedWatcher := sew.service.WatchExposed()
+	defer watcher.Stop(exposedWatcher, &sew.tomb)
+
+	for {
+		select {
+		case <-sew.tomb.Dying():
+			return
+		case exposed, ok := <-exposedWatcher.Changes():
+			if !ok {
+				sew.tomb.Kill(watcher.MustErr(exposedWatcher))
+				return
+			}
+			if !sew.fw.isRunning() {
+				sew.tomb.Kill(fmt.Errorf("agent associated to firewall is not running"))
+				return
+
+			}
+			if exposed {
+				log.Debugf("service %q exposed flag is set", sew.service)
+			} else {
+				log.Debugf("service %q exposed flag is unset", sew.service)
+			}
+			units, err := sew.service.AllUnits()
+			if err != nil {
+				log.Debugf("stopping watch on %q: %v", sew.service, err)
+				sew.tomb.Kill(err)
+				return
+			}
+			for _, unit := range units {
+				sew.fw.OpenClosePorts(unit)
+			}
+			if !exposed {
+				sew.exposed = false
+			} else {
+				sew.exposed = true
+				sew.suw = newServiceUnitsWatcher(sew.fw, sew)
+			}
+		}
+	}
+}
+
+// stop stops the service exposed watcher and returns any error 
+// encountered while watching.
+func (sew *serviceExposedWatcher) stop() error {
+	sew.tomb.Kill(nil)
+	return sew.tomb.Wait()
+}
+
+// serviceUnitsWatcher manages the opening and closing of ports in the firewall
+// for the units of one service.
+type serviceUnitsWatcher struct {
+	fw   *Firewall
+	sew  *serviceExposedWatcher
+	upws map[string]*unitPortsWatcher
+	tomb tomb.Tomb
+}
+
+// newServiceUnitsWatcher creates a new watcher for the units of 
+// the passed service.
+func newServiceUnitsWatcher(fw *Firewall, sew *serviceExposedWatcher) *serviceUnitsWatcher {
+	suw := &serviceUnitsWatcher{
+		fw:   fw,
+		sew:  sew,
+		upws: make(map[string]*unitPortsWatcher),
+	}
+	go suw.loop()
+	return suw
+}
+
+// loop watches on the units of its service.
+func (suw *serviceUnitsWatcher) loop() {
+	log.Debugf("start watching of service units for service %q", suw.sew.service)
+	defer suw.tomb.Done()
+	defer func() {
+		for _, upw := range suw.upws {
+			upw.stop()
+		}
+	}()
+	unitsWatcher := suw.sew.service.WatchUnits()
+	defer watcher.Stop(unitsWatcher, &suw.tomb)
+
+	for {
+		select {
+		case <-suw.tomb.Dying():
+			return
+		case <-suw.sew.tomb.Dying():
+			return
+		case changes, ok := <-unitsWatcher.Changes():
+			if !ok {
+				suw.tomb.Kill(watcher.MustErr(unitsWatcher))
+				return
+			}
+			if !suw.fw.isRunning() || !suw.sew.exposed {
+				suw.tomb.Kill(fmt.Errorf("agent associated to firewall is not running or service is not exposed"))
+				return
+			}
+			for _, unit := range changes.Removed {
+				if !suw.fw.isRunning() {
+					suw.tomb.Kill(fmt.Errorf("agent associated to firewall is not running"))
+				}
+				delete(suw.upws, unit.Name())
+
+				suw.fw.OpenClosePorts(unit)
+			}
+			for _, unit := range changes.Added {
+				suw.upws[unit.Name()] = newUnitPortsWatcher(suw.fw, suw, unit)
+			}
+		}
+	}
+}
+
+// unitPortsWatcher watches the opening and closing of the ports for a unit.
+type unitPortsWatcher struct {
+	fw   *Firewall
+	suw  *serviceUnitsWatcher
+	unit *state.Unit
+	tomb tomb.Tomb
+}
+
+// newUnitPortsWatch creates the new watcher for the ports of the passed unit.
+func newUnitPortsWatcher(fw *Firewall, suw *serviceUnitsWatcher, unit *state.Unit) *unitPortsWatcher {
+	upw := &unitPortsWatcher{
+		fw:   fw,
+		suw:  suw,
+		unit: unit,
+	}
+	go upw.loop()
+	return upw
+}
+
+// loop watches on the ports of its unit.
+func (upw *unitPortsWatcher) loop() {
+	log.Debugf("start watching of ports for unit %q", upw.unit)
+	defer upw.tomb.Done()
+	portsWatcher := upw.unit.WatchPorts()
+	defer watcher.Stop(portsWatcher, &upw.tomb)
+
+	for {
+		select {
+		case <-upw.tomb.Dying():
+			return
+		case _, ok := <-portsWatcher.Changes():
+			if !ok {
+				upw.tomb.Kill(watcher.MustErr(portsWatcher))
+				return
+			}
+			_, unitIsWatched := upw.suw.upws[upw.unit.Name()]
+			if !upw.fw.isRunning() || !upw.suw.sew.exposed || !unitIsWatched {
+				upw.tomb.Kill(fmt.Errorf("agent associated to firewall is not running or service is not exposed"))
+				return
+			}
+			upw.fw.OpenClosePorts(upw.unit)
+		}
+	}
+}
+
+// stop stops the unit ports watcher and returns any error 
+// encountered while watching.
+func (upw *unitPortsWatcher) stop() error {
+	upw.tomb.Kill(nil)
+	return upw.tomb.Wait()
+}
+
+// errorContextf prefixes any error stored in err with text formatted
+// according to the format specifier. If err does not contain an error,
+// errorContextf does nothing.
+func errorContextf(err *error, format string, args ...interface{}) {
+	if *err != nil {
+		*err = errors.New(fmt.Sprintf(format, args...) + ": " + (*err).Error())
+	}
+}

=== added file 'cmd/firewall_test.go'
--- cmd/firewall_test.go	1970-01-01 00:00:00 +0000
+++ cmd/firewall_test.go	2012-07-10 15:08:35 +0000
@@ -0,0 +1,348 @@
+package cmd_test
+
+import (
+	"fmt"
+	. "launchpad.net/gocheck"
+	"launchpad.net/juju-core/cmd"
+	"launchpad.net/juju-core/state"
+	"launchpad.net/juju-core/state/testing"
+	"sort"
+	"time"
+)
+
+type FirewallSuite struct {
+	testing.StateSuite
+	charm           *state.Charm
+	running         bool
+	portManager     cmd.PortManager
+	fw              *cmd.Firewall
+	servicesWatcher *state.ServicesWatcher
+	seenUnitChan    chan bool
+	seenMachineChan chan bool
+}
+
+func (s *FirewallSuite) SetUpTest(c *C) {
+	s.StateSuite.SetUpTest(c)
+	s.charm = s.AddTestingCharm(c, "dummy")
+}
+
+func (s *FirewallSuite) TearDownTest(c *C) {
+	s.StateSuite.TearDownTest(c)
+}
+
+func (s *FirewallSuite) startFirewall() {
+	s.portManager = newPortManagerMock()
+	s.fw = cmd.NewFirewall(s.State, func() bool { return s.running }, s.portManager)
+	s.servicesWatcher = s.State.WatchServices()
+	go func() {
+		for change := range s.servicesWatcher.Changes() {
+			s.fw.WatchServiceChanges(change)
+		}
+	}()
+	s.seenUnitChan = make(chan bool)
+	s.seenMachineChan = make(chan bool)
+	s.running = true
+}
+
+func (s *FirewallSuite) stopFirewall() {
+	s.servicesWatcher.Stop()
+	s.fw = nil
+}
+
+func (s *FirewallSuite) observeUnits(unitNames ...string) func() error {
+	seen := map[string]bool{}
+	s.fw.AddOpenClosePortsObserver(func(unit *state.Unit) {
+		seen[unit.Name()] = true
+		for _, unitName := range unitNames {
+			if seen[unitName] {
+				select {
+				case s.seenUnitChan <- true:
+				default:
+				}
+			}
+		}
+	})
+	return func() error {
+		select {
+		case <-s.seenUnitChan:
+			return nil
+		case <-time.After(2 * time.Second):
+			return fmt.Errorf("timeout during wait for observed units")
+		}
+		return nil
+	}
+}
+
+func (s *FirewallSuite) observeMachines(machineIds ...int) func() error {
+	seen := map[int]bool{}
+	s.fw.AddOpenClosePortsOnMachineObserver(func(machineId int) {
+		seen[machineId] = true
+		for _, machineId := range machineIds {
+			if seen[machineId] {
+				select {
+				case s.seenMachineChan <- true:
+				default:
+				}
+			}
+		}
+	})
+	return func() error {
+		select {
+		case <-s.seenMachineChan:
+			return nil
+		case <-time.After(2 * time.Second):
+			return fmt.Errorf("timeout during wait for observed machines")
+		}
+		return nil
+	}
+}
+
+func (s *FirewallSuite) providerPorts(c *C, machine *state.Machine) cmd.Ports {
+	instanceId, err := machine.InstanceId()
+	c.Assert(err, IsNil)
+	ports, err := s.portManager.OpenPorts(instanceId, machine.Id())
+	c.Assert(err, IsNil)
+	sort.Sort(ports)
+	return ports
+}
+
+var _ = Suite(&FirewallSuite{})
+
+func (s *FirewallSuite) TestWatchServiceChanges(c *C) {
+	// Verify that a service unit is checked whenever a change
+	// occurs such that ports may need to be opened and/or closed
+	// for the machine corresponding to a given service unit.
+	s.startFirewall()
+
+	wait := s.observeUnits("wordpress/0")
+	wordpress, err := s.State.AddService("wordpress", s.charm)
+	c.Assert(err, IsNil)
+	_, err = wordpress.AddUnit()
+	c.Assert(err, IsNil)
+	err = wordpress.SetExposed()
+	c.Assert(err, IsNil)
+	c.Assert(wait(), IsNil)
+
+	err = wordpress.ClearExposed()
+	c.Assert(err, IsNil)
+	c.Assert(wait(), IsNil)
+
+	err = wordpress.SetExposed()
+	c.Assert(err, IsNil)
+	c.Assert(wait(), IsNil)
+
+	s.stopFirewall()
+}
+
+func (s *FirewallSuite) TestAddRemoveUnitsForExposedService(c *C) {
+	// Verify that adding/removing service units for an exposed
+	// service triggers the appropriate firewall management of
+	// opening/closing ports on the machines for the corresponding
+	//service units.
+	s.startFirewall()
+
+	wordpress, err := s.State.AddService("wordpress", s.charm)
+	c.Assert(err, IsNil)
+	err = wordpress.SetExposed()
+	c.Assert(err, IsNil)
+
+	wait := s.observeUnits("wordpress/0", "wordpress/1")
+	wordpress0, err := wordpress.AddUnit()
+	c.Assert(err, IsNil)
+	c.Assert(wait(), IsNil)
+
+	wait = s.observeUnits("wordpress/2")
+	err = wordpress.RemoveUnit(wordpress0)
+	c.Assert(err, IsNil)
+	_, err = wordpress.AddUnit()
+	c.Assert(err, IsNil)
+	c.Assert(wait(), IsNil)
+
+	s.stopFirewall()
+}
+
+func (s *FirewallSuite) TestOpenClosePorts(c *C) {
+	// Verify that opening/closing ports triggers the appropriate
+	// firewall management for the corresponding service units.
+	s.startFirewall()
+
+	wait := s.observeUnits("wordpress/0")
+	wordpress, err := s.State.AddService("wordpress", s.charm)
+	c.Assert(err, IsNil)
+	err = wordpress.SetExposed()
+	c.Assert(err, IsNil)
+	wordpress0, err := wordpress.AddUnit()
+	_, err = wordpress.AddUnit()
+	wordpress0.OpenPort("tcp", 443)
+	wordpress0.OpenPort("tcp", 80)
+	wordpress0.ClosePort("tcp", 443)
+	c.Assert(wait(), IsNil)
+
+	wait = s.observeUnits("wordpress/1", "wordpress/3")
+	wordpress3, err := wordpress.AddUnit()
+	wordpress0.OpenPort("tcp", 53)
+	wordpress3.OpenPort("tcp", 80)
+	c.Assert(wait(), IsNil)
+
+	wait = s.observeUnits("wordpress/0", "wordpress/1", "wordpress/3")
+	wordpress.ClearExposed()
+	c.Assert(wait(), IsNil)
+
+	s.stopFirewall()
+}
+
+func (s *FirewallSuite) TestRemoveService(c *C) {
+	// Verify that firewall for corresponding service units
+	// is triggered upon the service's removal.
+	s.startFirewall()
+
+	wait := s.observeUnits("wordpress/0", "wordpress/1")
+	wordpress, err := s.State.AddService("wordpress", s.charm)
+	c.Assert(err, IsNil)
+	wordpress.AddUnit()
+	wordpress.AddUnit()
+	wordpress.SetExposed()
+	c.Assert(wait(), IsNil)
+
+	// Do not clear the exposed flag prior to removal, triggering
+	// should still occur as expected
+	s.State.RemoveService(wordpress)
+
+	s.stopFirewall()
+}
+
+func (s *FirewallSuite) TestPortManagementForUnexposedServiceIsANop(c *C) {
+	// Verify that activity on an unexposed service does NOT
+	// trigger firewall for the corresponding service unit.
+	s.startFirewall()
+
+	wait := s.observeUnits("not-called")
+	wordpress, err := s.State.AddService("wordpress", s.charm)
+	c.Assert(err, IsNil)
+	wordpress0, err := wordpress.AddUnit()
+	c.Assert(err, IsNil)
+	wordpress0.OpenPort("tcp", 53)
+	c.Assert(wait(), ErrorMatches, `timeout during wait for observed units`)
+
+	s.stopFirewall()
+}
+
+func (s *FirewallSuite) TestProvisioningAgentRestart(c *C) {
+	// Verify that firewall management is correct if the agent restarts.
+	// In particular, this test verifies that all state relevant for
+	// firewall management is stored in ZK and not in the agent
+	// itself.
+	wordpress, err := s.State.AddService("wordpress", s.charm)
+	c.Assert(err, IsNil)
+	wordpress0, err := wordpress.AddUnit()
+	wordpress1, err := wordpress.AddUnit()
+	wordpress1.OpenPort("tcp", 443)
+	wordpress1.OpenPort("tcp", 80)
+	wordpress.SetExposed()
+
+	// Simulate agent start.
+	s.startFirewall()
+
+	wait := s.observeUnits("wordpress/0", "wordpress/1")
+	wordpress0.OpenPort("tcp", 53)
+	wordpress1.ClosePort("tcp", 443)
+	c.Assert(wait(), IsNil)
+
+	wait = s.observeUnits("wordpress/1")
+	wordpress1.ClosePort("tcp", 80)
+
+	wait = s.observeUnits("wordpress/0", "wordpress/1")
+	wordpress.ClearExposed()
+	c.Assert(wait(), IsNil)
+
+	s.stopFirewall()
+}
+
+func (s *FirewallSuite) TestOpenClosePortsOnMachine(c *C) {
+	// Verify opening/closing ports on a machine works properly.
+	s.startFirewall()
+
+	machine, err := s.State.AddMachine()
+	c.Assert(err, IsNil)
+	machine.SetInstanceId("m-0")
+	s.fw.ProcessMachine(machine)
+
+	// Expose a service.
+	wordpress, err := s.State.AddService("wordpress", s.charm)
+	c.Assert(err, IsNil)
+	wordpress.SetExposed()
+	wordpress0, err := wordpress.AddUnit()
+	wordpress0.OpenPort("tcp", 80)
+	wordpress0.OpenPort("tcp", 443)
+	wordpress0.AssignToMachine(machine)
+	err = s.fw.OpenClosePortsOnMachine(machine.Id())
+	c.Assert(err, IsNil)
+	ports := s.providerPorts(c, machine)
+	c.Assert(ports, DeepEquals, cmd.Ports{{"tcp", 80}, {"tcp", 443}})
+
+	// Change port setup.
+	wordpress0.OpenPort("tcp", 8080)
+	wordpress0.ClosePort("tcp", 443)
+	err = s.fw.OpenClosePortsOnMachine(machine.Id())
+	c.Assert(err, IsNil)
+	ports = s.providerPorts(c, machine)
+	c.Assert(ports, DeepEquals, cmd.Ports{{"tcp", 80}, {"tcp", 8080}})
+
+	s.stopFirewall()
+}
+
+type portManagerMock struct {
+	ports map[string]map[int]map[state.Port]bool
+}
+
+func newPortManagerMock() cmd.PortManager {
+	return &portManagerMock{make(map[string]map[int]map[state.Port]bool)}
+}
+
+func (m *portManagerMock) OpenPorts(instanceId string, machineId int) (cmd.Ports, error) {
+	instance, ok := m.ports[instanceId]
+	if !ok {
+		instance = make(map[int]map[state.Port]bool)
+		m.ports[instanceId] = instance
+	}
+	machine, ok := instance[machineId]
+	if !ok {
+		machine = make(map[state.Port]bool)
+		instance[machineId] = machine
+	}
+	ports := cmd.Ports{}
+	for port, open := range machine {
+		if open {
+			ports = append(ports, port)
+		}
+	}
+	return ports, nil
+}
+
+func (m *portManagerMock) OpenPort(instanceId string, machineId int, port state.Port) error {
+	instance, ok := m.ports[instanceId]
+	if !ok {
+		instance = make(map[int]map[state.Port]bool)
+		m.ports[instanceId] = instance
+	}
+	machine, ok := instance[machineId]
+	if !ok {
+		machine = make(map[state.Port]bool)
+		instance[machineId] = machine
+	}
+	machine[port] = true
+	return nil
+}
+
+func (m *portManagerMock) ClosePort(instanceId string, machineId int, port state.Port) error {
+	instance, ok := m.ports[instanceId]
+	if !ok {
+		return fmt.Errorf("illegal instance %q", instanceId)
+	}
+	machine, ok := instance[machineId]
+	if !ok {
+		return fmt.Errorf("illegal machine %d", machineId)
+	}
+	delete(machine, port)
+	return nil
+}