Merge "veyron/tools/principal: Fix possible man in the middle attack by having the principal tool check the public key of the server."
diff --git a/lib/modules/registry.go b/lib/modules/registry.go
index f473d14..555c1a2 100644
--- a/lib/modules/registry.go
+++ b/lib/modules/registry.go
@@ -99,6 +99,27 @@
 	return append(newEnv, shellEntryPoint+"="+name)
 }
 
+// Dispatch will execute the requested subprocess command from a within a
+// a subprocess. It will return without an error if it is executed by a
+// process that does not specify an entry point in its environment.
+//
+// func main() {
+//     if modules.IsModulesProcess() {
+//         if err := modules.Dispatch(); err != nil {
+//             panic("error")
+//         }
+//         eturn
+//     }
+//     parent code...
+//
+func Dispatch() error {
+	if !IsModulesProcess() {
+		return nil
+	}
+	return registry.dispatch()
+}
+
+// TODO(cnicolaou): delete this in a subsequent CL.
 // DispatchInTest will execute the requested subproccess command from within
 // a unit test run as a subprocess.
 func DispatchInTest() {
@@ -111,30 +132,6 @@
 	os.Exit(0)
 }
 
-// Dispatch will execute the requested subprocess command from a within a
-// a subprocess that is not a unit test. It will return without an error
-// if it is executed by a process that does not specify an entry point
-// in its environment.
-//
-// func main() {
-//     if modules.IsModulesProcess() {
-//         if err := modules.Dispatch(); err != nil {
-//             panic("error")
-//          }
-//          return
-//      }
-//      parent code...
-//
-func Dispatch() error {
-	if !IsModulesProcess() {
-		return nil
-	}
-	if IsTestHelperProcess() {
-		return fmt.Errorf("use DispatchInTest in unittests")
-	}
-	return registry.dispatch()
-}
-
 // DispatchAndExit is like Dispatch except that it will call os.Exit(0)
 // when executed within a child process and the command succeeds, or panic
 // on encountering an error.
diff --git a/runtimes/google/ipc/server.go b/runtimes/google/ipc/server.go
index 3665d59..fc974f8 100644
--- a/runtimes/google/ipc/server.go
+++ b/runtimes/google/ipc/server.go
@@ -1037,7 +1037,7 @@
 	// on the server even if they will not be allowed to collect the
 	// results later.  This might be considered a DOS vector.
 	spanName := fmt.Sprintf("\"%s\".%s", fs.Name(), fs.Method())
-	fs.T, _ = ivtrace.SetContinuedSpan(fs.T, spanName, req.TraceRequest)
+	fs.T, _ = vtrace.SetContinuedTrace(fs.T, spanName, req.TraceRequest)
 
 	var cancel context.CancelFunc
 	if req.Timeout != ipc.NoTimeout {
diff --git a/runtimes/google/vtrace/vtrace.go b/runtimes/google/vtrace/vtrace.go
index 528d717..7ca647c 100644
--- a/runtimes/google/vtrace/vtrace.go
+++ b/runtimes/google/vtrace/vtrace.go
@@ -82,18 +82,6 @@
 	return vtrace.Response{}
 }
 
-// ContinuedSpan creates a span that represents a continuation of a trace from
-// a remote server.  name is a user readable string that describes the context
-// and req contains the parameters needed to connect this span with it's trace.
-func SetContinuedSpan(ctx *context.T, name string, req vtrace.Request) (*context.T, vtrace.Span) {
-	st := getStore(ctx)
-	if req.Method == vtrace.InMemory {
-		st.ForceCollect(req.TraceID)
-	}
-	newSpan := newSpan(req.SpanID, name, req.TraceID, st)
-	return context.WithValue(ctx, spanKey, newSpan), newSpan
-}
-
 type contextKey int
 
 const (
@@ -119,6 +107,19 @@
 	return context.WithValue(ctx, spanKey, s), s
 }
 
+// SetContinuedTrace creates a span that represents a continuation of
+// a trace from a remote server.  name is the name of the new span and
+// req contains the parameters needed to connect this span with it's
+// trace.
+func (m manager) SetContinuedTrace(ctx *context.T, name string, req vtrace.Request) (*context.T, vtrace.Span) {
+	st := getStore(ctx)
+	if req.Method == vtrace.InMemory {
+		st.ForceCollect(req.TraceID)
+	}
+	newSpan := newSpan(req.SpanID, name, req.TraceID, st)
+	return context.WithValue(ctx, spanKey, newSpan), newSpan
+}
+
 // SetNewSpan derives a context with a new Span that can be used to
 // trace and annotate operations across process boundaries.
 func (m manager) SetNewSpan(ctx *context.T, name string) (*context.T, vtrace.Span) {
diff --git a/services/identity/identityd/v23_test.go b/services/identity/identityd/v23_test.go
index 72afa15..b0ea5cb 100644
--- a/services/identity/identityd/v23_test.go
+++ b/services/identity/identityd/v23_test.go
@@ -13,7 +13,6 @@
 	cleanup := v23tests.UseSharedBinDir()
 	r := m.Run()
 	cleanup()
-	// TODO(cnicolaou): call modules.Dispatch and remove the need for TestHelperProcess
 	os.Exit(r)
 }
 
diff --git a/services/mgmt/application/applicationd/v23_test.go b/services/mgmt/application/applicationd/v23_test.go
index 9e43b89..f3e858c 100644
--- a/services/mgmt/application/applicationd/v23_test.go
+++ b/services/mgmt/application/applicationd/v23_test.go
@@ -13,7 +13,6 @@
 	cleanup := v23tests.UseSharedBinDir()
 	r := m.Run()
 	cleanup()
-	// TODO(cnicolaou): call modules.Dispatch and remove the need for TestHelperProcess
 	os.Exit(r)
 }
 
diff --git a/services/mgmt/binary/binaryd/v23_test.go b/services/mgmt/binary/binaryd/v23_test.go
index a74c88b..eec6864 100644
--- a/services/mgmt/binary/binaryd/v23_test.go
+++ b/services/mgmt/binary/binaryd/v23_test.go
@@ -13,7 +13,6 @@
 	cleanup := v23tests.UseSharedBinDir()
 	r := m.Run()
 	cleanup()
-	// TODO(cnicolaou): call modules.Dispatch and remove the need for TestHelperProcess
 	os.Exit(r)
 }
 
diff --git a/services/mgmt/build/buildd/v23_test.go b/services/mgmt/build/buildd/v23_test.go
index 6c1a4c8..8ed558e 100644
--- a/services/mgmt/build/buildd/v23_test.go
+++ b/services/mgmt/build/buildd/v23_test.go
@@ -13,7 +13,6 @@
 	cleanup := v23tests.UseSharedBinDir()
 	r := m.Run()
 	cleanup()
-	// TODO(cnicolaou): call modules.Dispatch and remove the need for TestHelperProcess
 	os.Exit(r)
 }
 
diff --git a/services/mgmt/device/deviced/server.go b/services/mgmt/device/deviced/server.go
index ac7e719..df62965 100644
--- a/services/mgmt/device/deviced/server.go
+++ b/services/mgmt/device/deviced/server.go
@@ -1,6 +1,8 @@
 package main
 
 import (
+	"crypto/rand"
+	"encoding/base64"
 	"flag"
 	"fmt"
 	"net"
@@ -30,6 +32,7 @@
 	restartExitCode = flag.Int("restart_exit_code", 0, "exit code to return when device manager should be restarted")
 	nhName          = flag.String("neighborhood_name", "", `if provided, it will enable sharing with the local neighborhood with the provided name. The address of the local mounttable will be published to the neighboorhood and everything in the neighborhood will be visible on the local mounttable.`)
 	dmPort          = flag.Int("deviced_port", 0, "the port number of assign to the device manager service. The hostname/IP address part of --veyron.tcp.address is used along with this port. By default, the port is assigned by the OS.")
+	usePairingToken = flag.Bool("use_pairing_token", false, "generate a pairing token for the device manager that will need to be provided when a device is claimed")
 )
 
 func runServer(*cmdline.Command, []string) error {
@@ -70,15 +73,25 @@
 		ns.ListenSpec = veyron2.GetListenSpec(ctx)
 		ns.Name = *publishAs
 	}
+	var pairingToken string
+	if *usePairingToken {
+		var token [8]byte
+		if _, err := rand.Read(token[:]); err != nil {
+			vlog.Errorf("unable to generate pairing token: %v", err)
+			return err
+		}
+		pairingToken = base64.URLEncoding.EncodeToString(token[:])
+		vlog.VI(0).Infof("Device manager pairing token: %v", pairingToken)
+	}
 	dev := starter.DeviceArgs{
 		ConfigState:     configState,
 		TestMode:        testMode,
 		RestartCallback: func() { exitErr = cmdline.ErrExitCode(*restartExitCode) },
+		PairingToken:    pairingToken,
 	}
 	if dev.ListenSpec, err = newDeviceListenSpec(ns.ListenSpec, *dmPort); err != nil {
 		return err
 	}
-
 	stop, err := starter.Start(ctx, starter.Args{Namespace: ns, Device: dev, MountGlobalNamespaceInLocalNamespace: true})
 	if err != nil {
 		return err
diff --git a/services/mgmt/device/impl/device_service.go b/services/mgmt/device/impl/device_service.go
index 533e3ea..678be1a 100644
--- a/services/mgmt/device/impl/device_service.go
+++ b/services/mgmt/device/impl/device_service.go
@@ -171,8 +171,8 @@
 	return args, nil
 }
 
-func (s *deviceService) Claim(ctx ipc.ServerContext) error {
-	return s.disp.claimDeviceManager(ctx)
+func (s *deviceService) Claim(ctx ipc.ServerContext, pairingToken string) error {
+	return s.disp.claimDeviceManager(ctx, pairingToken)
 }
 
 func (*deviceService) Describe(ipc.ServerContext) (device.Description, error) {
diff --git a/services/mgmt/device/impl/dispatcher.go b/services/mgmt/device/impl/dispatcher.go
index 6674f3e..fec7b3a 100644
--- a/services/mgmt/device/impl/dispatcher.go
+++ b/services/mgmt/device/impl/dispatcher.go
@@ -1,6 +1,7 @@
 package impl
 
 import (
+	"crypto/subtle"
 	"fmt"
 	"os"
 	"path"
@@ -55,6 +56,9 @@
 	mtAddress string // The address of the local mounttable.
 	// reap is the app process monitoring subsystem.
 	reap reaper
+	// pairingToken (if present) is a token that the device manager expects
+	// to be replayed by the principal who claims the device.
+	pairingToken string
 }
 
 var _ ipc.Dispatcher = (*dispatcher)(nil)
@@ -75,14 +79,15 @@
 	ErrUpdateNoOp          = verror.Register(pkgPath+".UpdateNoOp", verror.NoRetry, "{1:}{2:} update is no op{:_}")
 	ErrInvalidOperation    = verror.Register(pkgPath+".InvalidOperation", verror.NoRetry, "{1:}{2:} invalid operation{:_}")
 	ErrInvalidBlessing     = verror.Register(pkgPath+".InvalidBlessing", verror.NoRetry, "{1:}{2:} invalid blessing{:_}")
+	ErrInvalidPairingToken = verror.Register(pkgPath+".InvalidPairingToken", verror.NoRetry, "{1:}{2:} pairing token mismatch{:_}")
 )
 
 // NewDispatcher is the device manager dispatcher factory.
-func NewDispatcher(ctx *context.T, config *config.State, mtAddress string, testMode bool, restartHandler func()) (ipc.Dispatcher, error) {
+func NewDispatcher(ctx *context.T, config *config.State, mtAddress string, pairingToken string, testMode bool, restartHandler func()) (ipc.Dispatcher, error) {
 	if err := config.Validate(); err != nil {
 		return nil, fmt.Errorf("invalid config %v: %v", config, err)
 	}
-	// TODO(caprita): use some mechansim (a file lock or presence of entry
+	// TODO(caprita): use some mechanism (a file lock or presence of entry
 	// in mounttable) to ensure only one device manager is running in an
 	// installation?
 	mi := &managerInfo{
@@ -107,12 +112,13 @@
 			restartHandler: restartHandler,
 			testMode:       testMode,
 		},
-		config:    config,
-		uat:       uat,
-		locks:     acls.NewLocks(),
-		principal: veyron2.GetPrincipal(ctx),
-		mtAddress: mtAddress,
-		reap:      reap,
+		config:       config,
+		uat:          uat,
+		locks:        acls.NewLocks(),
+		principal:    veyron2.GetPrincipal(ctx),
+		mtAddress:    mtAddress,
+		reap:         reap,
+		pairingToken: pairingToken,
 	}
 
 	// If we're in 'security agent mode', set up the key manager agent.
@@ -147,9 +153,13 @@
 	return filepath.Join(d.config.Root, "device-manager", "device-data", "acls")
 }
 
-func (d *dispatcher) claimDeviceManager(ctx ipc.ServerContext) error {
+func (d *dispatcher) claimDeviceManager(ctx ipc.ServerContext, pairingToken string) error {
 	// TODO(rjkroege): Scrub the state tree of installation and instance ACL files.
 
+	// Verify that the claimer pairing tokens match that of the device manager.
+	if subtle.ConstantTimeCompare([]byte(pairingToken), []byte(d.pairingToken)) != 1 {
+		return verror.New(ErrInvalidPairingToken, ctx.Context())
+	}
 	// Get the blessings to be used by the claimant.
 	blessings := ctx.Blessings()
 	if blessings == nil {
diff --git a/services/mgmt/device/impl/impl_test.go b/services/mgmt/device/impl/impl_test.go
index a26fbf4..9e12a48 100644
--- a/services/mgmt/device/impl/impl_test.go
+++ b/services/mgmt/device/impl/impl_test.go
@@ -160,11 +160,15 @@
 	// This exemplifies how to override or set specific config fields, if,
 	// for example, the device manager is invoked 'by hand' instead of via a
 	// script prepared by a previous version of the device manager.
+	var pairingToken string
 	if len(args) > 0 {
-		if want, got := 4, len(args); want != got {
-			vlog.Fatalf("expected %d additional arguments, got %d instead: %q", want, got, args)
+		if want, got := 4, len(args); want > got {
+			vlog.Fatalf("expected atleast %d additional arguments, got %d instead: %q", want, got, args)
 		}
 		configState.Root, configState.Helper, configState.Origin, configState.CurrentLink = args[0], args[1], args[2], args[3]
+		if len(args) > 4 {
+			pairingToken = args[4]
+		}
 	}
 
 	stop, err := starter.Start(ctx, starter.Args{
@@ -177,6 +181,7 @@
 			ConfigState:     configState,
 			TestMode:        strings.HasSuffix(fmt.Sprint(veyron2.GetPrincipal(ctx).BlessingStore().Default()), "/testdm"),
 			RestartCallback: func() { fmt.Println("restart handler") },
+			PairingToken:    pairingToken,
 		},
 		// TODO(rthellend): Wire up the local mounttable like the real device
 		// manager, i.e. mount the device manager and the apps on it, and mount
@@ -824,7 +829,8 @@
 
 	// Set up the device manager.  Since we won't do device manager updates,
 	// don't worry about its application envelope and current link.
-	_, dms := mgmttest.RunShellCommand(t, sh, nil, deviceManagerCmd, "dm", root, helperPath, "unused_app_repo_name", "unused_curr_link")
+	pairingToken := "abcxyz"
+	_, dms := mgmttest.RunShellCommand(t, sh, nil, deviceManagerCmd, "dm", root, helperPath, "unused_app_repo_name", "unused_curr_link", pairingToken)
 	pid := mgmttest.ReadPID(t, dms)
 	defer syscall.Kill(pid, syscall.SIGINT)
 
@@ -844,8 +850,14 @@
 	// Devicemanager should have open ACLs before we claim it and so an
 	// Install from octx should succeed.
 	installApp(t, octx)
+
+	// Claim the devicemanager with an incorrect pairing token should fail
+	if err := deviceStub.Claim(claimantCtx, "abcxy", &granter{p: veyron2.GetPrincipal(claimantCtx), extension: "mydevice"}); err == nil || !verror.Is(err, impl.ErrInvalidPairingToken.ID) {
+		t.Fatalf("Claim with an incorrect pairing token should have failed: %v", err)
+	}
+
 	// Claim the devicemanager with claimantRT as <defaultblessing>/mydevice
-	if err := deviceStub.Claim(claimantCtx, &granter{p: veyron2.GetPrincipal(claimantCtx), extension: "mydevice"}); err != nil {
+	if err := deviceStub.Claim(claimantCtx, pairingToken, &granter{p: veyron2.GetPrincipal(claimantCtx), extension: "mydevice"}); err != nil {
 		t.Fatal(err)
 	}
 
@@ -926,7 +938,8 @@
 	}
 
 	// Claim the devicemanager as "root/self/mydevice"
-	if err := deviceStub.Claim(selfCtx, &granter{p: veyron2.GetPrincipal(selfCtx), extension: "mydevice"}); err != nil {
+	var pairingToken string
+	if err := deviceStub.Claim(selfCtx, pairingToken, &granter{p: veyron2.GetPrincipal(selfCtx), extension: "mydevice"}); err != nil {
 		t.Fatal(err)
 	}
 	expectedACL := make(access.TaggedACLMap)
@@ -1386,7 +1399,8 @@
 	}
 
 	// self claims the device manager.
-	if err := deviceStub.Claim(selfCtx, &granter{p: veyron2.GetPrincipal(selfCtx), extension: "alice"}); err != nil {
+	var pairingToken string
+	if err := deviceStub.Claim(selfCtx, pairingToken, &granter{p: veyron2.GetPrincipal(selfCtx), extension: "alice"}); err != nil {
 		t.Fatalf("Claim failed: %v", err)
 	}
 
@@ -1501,7 +1515,8 @@
 	appID := installApp(t, selfCtx)
 
 	// Claim the devicemanager with selfCtx as root/self/alice
-	if err := deviceStub.Claim(selfCtx, &granter{p: veyron2.GetPrincipal(selfCtx), extension: "alice"}); err != nil {
+	var pairingToken string
+	if err := deviceStub.Claim(selfCtx, pairingToken, &granter{p: veyron2.GetPrincipal(selfCtx), extension: "alice"}); err != nil {
 		t.Fatal(err)
 	}
 
diff --git a/services/mgmt/device/impl/util.go b/services/mgmt/device/impl/util.go
index 6ce2afa..096f100 100644
--- a/services/mgmt/device/impl/util.go
+++ b/services/mgmt/device/impl/util.go
@@ -49,7 +49,7 @@
 		vlog.Errorf("Publisher binary(%v) signature verification failed", bin.File)
 		return err
 	}
-	path, perm := filepath.Join(workspace, fileName), os.FileMode(755)
+	path, perm := filepath.Join(workspace, fileName), os.FileMode(0755)
 	if err := ioutil.WriteFile(path, data, perm); err != nil {
 		vlog.Errorf("WriteFile(%v, %v) failed: %v", path, perm, err)
 		return verror.New(ErrOperationFailed, nil)
diff --git a/services/mgmt/device/starter/starter.go b/services/mgmt/device/starter/starter.go
index 62c84ad..c7c21f1 100644
--- a/services/mgmt/device/starter/starter.go
+++ b/services/mgmt/device/starter/starter.go
@@ -32,6 +32,7 @@
 	ConfigState     *config.State  // Configuration for the device.
 	TestMode        bool           // Whether the device is running in test mode or not.
 	RestartCallback func()         // Callback invoked when the device service is restarted.
+	PairingToken    string         // PairingToken that a claimer needs to provide.
 }
 
 type Args struct {
@@ -101,7 +102,7 @@
 	args.ConfigState.Name = endpoints[0].Name()
 	vlog.Infof("Device manager object name: %v", args.ConfigState.Name)
 
-	dispatcher, err := impl.NewDispatcher(ctx, args.ConfigState, mt, args.TestMode, args.RestartCallback)
+	dispatcher, err := impl.NewDispatcher(ctx, args.ConfigState, mt, args.PairingToken, args.TestMode, args.RestartCallback)
 	if err != nil {
 		shutdown()
 		return nil, err
diff --git a/services/mgmt/profile/profiled/v23_test.go b/services/mgmt/profile/profiled/v23_test.go
index 2917072..cb048e1 100644
--- a/services/mgmt/profile/profiled/v23_test.go
+++ b/services/mgmt/profile/profiled/v23_test.go
@@ -13,7 +13,6 @@
 	cleanup := v23tests.UseSharedBinDir()
 	r := m.Run()
 	cleanup()
-	// TODO(cnicolaou): call modules.Dispatch and remove the need for TestHelperProcess
 	os.Exit(r)
 }
 
diff --git a/services/mounttable/mounttabled/v23_test.go b/services/mounttable/mounttabled/v23_test.go
index 7cfad44..a464bad 100644
--- a/services/mounttable/mounttabled/v23_test.go
+++ b/services/mounttable/mounttabled/v23_test.go
@@ -13,7 +13,6 @@
 	cleanup := v23tests.UseSharedBinDir()
 	r := m.Run()
 	cleanup()
-	// TODO(cnicolaou): call modules.Dispatch and remove the need for TestHelperProcess
 	os.Exit(r)
 }
 
diff --git a/tools/debug/v23_test.go b/tools/debug/v23_test.go
index 0dc059b..451bd67 100644
--- a/tools/debug/v23_test.go
+++ b/tools/debug/v23_test.go
@@ -13,7 +13,6 @@
 	cleanup := v23tests.UseSharedBinDir()
 	r := m.Run()
 	cleanup()
-	// TODO(cnicolaou): call modules.Dispatch and remove the need for TestHelperProcess
 	os.Exit(r)
 }
 
diff --git a/tools/mgmt/device/devicex b/tools/mgmt/device/devicex
index 17fc8ea..4650a66 100755
--- a/tools/mgmt/device/devicex
+++ b/tools/mgmt/device/devicex
@@ -117,12 +117,12 @@
 
 ###############################################################################
 # Installs device manager: fetches binaries, configures suidhelper, calls the
-# self-install command on deviced.
+# install command on deviced.
 # Globals:
 #   VANADIUM_DEVICE_DIR
 # Arguments:
 #   source of binaries (optional)
-#   args for self-install command and for device manager (optional)
+#   args for install command and for device manager (optional)
 # Returns:
 #   None
 ###############################################################################
@@ -167,11 +167,29 @@
   done
   local -r SETUID_SCRIPT="${BIN_INSTALL}/suidhelper"
   if [[ ${SINGLE_USER} == false ]]; then
-    sudo bash -c "chown root:root \"${SETUID_SCRIPT}\"; chmod 4551 \"${SETUID_SCRIPT}\""
+    case "$(uname)" in
+    "Darwin")
+      # Group root not available on Darwin and chown
+      # not in default path.
+      sudo /usr/sbin/chown root:wheel "${SETUID_SCRIPT}"
+      ;;
+    "Linux")
+      sudo chown root:root "${SETUID_SCRIPT}"
+      ;;
+     esac
+     sudo chmod 4551 "${SETUID_SCRIPT}"
   fi
   local -r INIT_SCRIPT="${BIN_INSTALL}/inithelper"
   if [[ ${INIT_MODE} == true ]]; then
-    sudo bash -c "chown root:root \"${INIT_SCRIPT}\"; chmod 4551 \"${INIT_SCRIPT}\""
+    case "$(uname)" in
+    "Darwin")
+      sudo /usr/sbin/chown root:wheel "${INIT_SCRIPT}"
+      ;;
+    "Linux")
+      sudo chown root:root "${INIT_SCRIPT}"
+      ;;
+    esac
+    sudo chmod 4551 "${INIT_SCRIPT}"
   fi
   echo "Helpers configured."
 
diff --git a/tools/mgmt/device/doc.go b/tools/mgmt/device/doc.go
index ff4c731..7501841 100644
--- a/tools/mgmt/device/doc.go
+++ b/tools/mgmt/device/doc.go
@@ -193,7 +193,7 @@
 Claim the device.
 
 Usage:
-   device claim <device> <grant extension>
+   device claim <device> <grant extension> <pairing token>
 
 <device> is the veyron object name of the device manager's device service.
 
diff --git a/tools/mgmt/device/impl/devicemanager_mock_test.go b/tools/mgmt/device/impl/devicemanager_mock_test.go
index 5b652d1..7c039e9 100644
--- a/tools/mgmt/device/impl/devicemanager_mock_test.go
+++ b/tools/mgmt/device/impl/devicemanager_mock_test.go
@@ -68,7 +68,7 @@
 	return mni.simpleCore(AddAssociationStimulus{"AssociateAccount", identityNames, accountName}, "AssociateAccount")
 }
 
-func (mni *mockDeviceInvoker) Claim(call ipc.ServerContext) error {
+func (mni *mockDeviceInvoker) Claim(call ipc.ServerContext, pairingToken string) error {
 	return mni.simpleCore("Claim", "Claim")
 }
 
diff --git a/tools/mgmt/device/impl/impl.go b/tools/mgmt/device/impl/impl.go
index 4f75fd4..e3b789e 100644
--- a/tools/mgmt/device/impl/impl.go
+++ b/tools/mgmt/device/impl/impl.go
@@ -148,7 +148,7 @@
 	Name:     "claim",
 	Short:    "Claim the device.",
 	Long:     "Claim the device.",
-	ArgsName: "<device> <grant extension>",
+	ArgsName: "<device> <grant extension> <pairing token>",
 	ArgsLong: `
 <device> is the veyron object name of the device manager's device service.
 
@@ -157,12 +157,16 @@
 }
 
 func runClaim(cmd *cmdline.Command, args []string) error {
-	if expected, got := 2, len(args); expected != got {
-		return cmd.UsageErrorf("claim: incorrect number of arguments, expected %d, got %d", expected, got)
+	if expected, max, got := 2, 3, len(args); expected > got || got > max {
+		return cmd.UsageErrorf("claim: incorrect number of arguments, expected atleast %d (max: %d), got %d", expected, max, got)
 	}
 	deviceName, grant := args[0], args[1]
+	var pairingToken string
+	if len(args) > 2 {
+		pairingToken = args[2]
+	}
 	principal := veyron2.GetPrincipal(gctx)
-	if err := device.DeviceClient(deviceName).Claim(gctx, &granter{p: principal, extension: grant}); err != nil {
+	if err := device.DeviceClient(deviceName).Claim(gctx, pairingToken, &granter{p: principal, extension: grant}); err != nil {
 		return fmt.Errorf("Claim failed: %v", err)
 	}
 	fmt.Fprintln(cmd.Stdout(), "Successfully claimed.")
diff --git a/tools/mgmt/device/impl/impl_test.go b/tools/mgmt/device/impl/impl_test.go
index d4118dd..d37bd3d 100644
--- a/tools/mgmt/device/impl/impl_test.go
+++ b/tools/mgmt/device/impl/impl_test.go
@@ -293,17 +293,17 @@
 	if err := cmd.Execute([]string{"claim", "nope"}); err == nil {
 		t.Fatalf("wrongly failed to receive a non-nil error.")
 	}
-	if expected, got := "ERROR: claim: incorrect number of arguments, expected 2, got 1", strings.TrimSpace(stderr.String()); !strings.HasPrefix(got, expected) {
+	if expected, got := "ERROR: claim: incorrect number of arguments, expected atleast 2 (max: 3), got 1", strings.TrimSpace(stderr.String()); !strings.HasPrefix(got, expected) {
 		t.Fatalf("Unexpected output from claim. Got %q, expected prefix %q", got, expected)
 	}
 	stdout.Reset()
 	stderr.Reset()
 	tape.Rewind()
 
-	if err := cmd.Execute([]string{"claim", "nope", "nope", "nope"}); err == nil {
+	if err := cmd.Execute([]string{"claim", "nope", "nope", "nope", "nope"}); err == nil {
 		t.Fatalf("wrongly failed to receive a non-nil error.")
 	}
-	if expected, got := "ERROR: claim: incorrect number of arguments, expected 2, got 3", strings.TrimSpace(stderr.String()); !strings.HasPrefix(got, expected) {
+	if expected, got := "ERROR: claim: incorrect number of arguments, expected atleast 2 (max: 3), got 4", strings.TrimSpace(stderr.String()); !strings.HasPrefix(got, expected) {
 		t.Fatalf("Unexpected output from claim. Got %q, expected prefix %q", got, expected)
 	}
 	stdout.Reset()
@@ -314,8 +314,9 @@
 	tape.SetResponses([]interface{}{
 		nil,
 	})
-	if err := cmd.Execute([]string{"claim", deviceName, "grant"}); err != nil {
-		t.Fatalf("Claim(%s, %s) failed: %v", deviceName, "grant", err)
+	var pairingToken string
+	if err := cmd.Execute([]string{"claim", deviceName, "grant", pairingToken}); err != nil {
+		t.Fatalf("Claim(%s, %s, %s) failed: %v", deviceName, "grant", pairingToken, err)
 	}
 	if got, expected := len(tape.Play()), 1; got != expected {
 		t.Errorf("invalid call sequence. Got %v, want %v", got, expected)
@@ -337,7 +338,7 @@
 	tape.SetResponses([]interface{}{
 		verror.New(errOops, nil),
 	})
-	if err := cmd.Execute([]string{"claim", deviceName, "grant"}); err == nil {
+	if err := cmd.Execute([]string{"claim", deviceName, "grant", pairingToken}); err == nil {
 		t.Fatalf("claim() failed to detect error", err)
 	}
 	expected = []interface{}{
diff --git a/tools/naming/simulator/v23_test.go b/tools/naming/simulator/v23_test.go
index 9c42af4..043c638 100644
--- a/tools/naming/simulator/v23_test.go
+++ b/tools/naming/simulator/v23_test.go
@@ -13,7 +13,6 @@
 	cleanup := v23tests.UseSharedBinDir()
 	r := m.Run()
 	cleanup()
-	// TODO(cnicolaou): call modules.Dispatch and remove the need for TestHelperProcess
 	os.Exit(r)
 }
 
diff --git a/tools/principal/v23_test.go b/tools/principal/v23_test.go
index 806ef4d..665f5c3 100644
--- a/tools/principal/v23_test.go
+++ b/tools/principal/v23_test.go
@@ -13,7 +13,6 @@
 	cleanup := v23tests.UseSharedBinDir()
 	r := m.Run()
 	cleanup()
-	// TODO(cnicolaou): call modules.Dispatch and remove the need for TestHelperProcess
 	os.Exit(r)
 }
 
diff --git a/tools/vrun/internal/v23_test.go b/tools/vrun/internal/v23_test.go
index ebbc180..b214a3e 100644
--- a/tools/vrun/internal/v23_test.go
+++ b/tools/vrun/internal/v23_test.go
@@ -13,7 +13,6 @@
 	cleanup := v23tests.UseSharedBinDir()
 	r := m.Run()
 	cleanup()
-	// TODO(cnicolaou): call modules.Dispatch and remove the need for TestHelperProcess
 	os.Exit(r)
 }