Merge pull request #18 from google/fix

Polices can now be directly unlocked with Protectors

3 files changed, 96 insertions, 2 deletions

diff --git a/actions/context.go b/actions/context.go
index 8bf0287..fb25b54 100644
--- a/actions/context.go
+++ b/actions/context.go
@@ -44,7 +44,7 @@ var (
 	ErrBadConfigFile    = errors.New("global config file has invalid data")
 	ErrConfigFileExists = errors.New("global config file already exists")
 	ErrBadConfig        = errors.New("invalid Config structure provided")
-	ErrLocked           = errors.New("method needs a call to Unlock() first")
+	ErrLocked           = errors.New("key needs to be unlocked first")
 )
 
 // Context contains the necessary global state to perform most of fscrypt's
diff --git a/actions/policy.go b/actions/policy.go
index 0d0ed02..1291e6b 100644
--- a/actions/policy.go
+++ b/actions/policy.go
@@ -248,6 +248,27 @@ func (policy *Policy) Unlock(optionFn OptionFunc, keyFn KeyFunc) error {
 	return err
 }
 
+// UnlockWithProtector uses an unlocked Protector to unlock a policy. An error
+// is returned if the Protector is not yet unlocked or does not protect the
+// policy. Does nothing if policy is already unlocked.
+func (policy *Policy) UnlockWithProtector(protector *Protector) error {
+	if policy.key != nil {
+		return nil
+	}
+	if protector.key == nil {
+		return ErrLocked
+	}
+	idx, ok := policy.findWrappedKeyIndex(protector.Descriptor())
+	if !ok {
+		return ErrNotProtected
+	}
+
+	var err error
+	wrappedPolicyKey := policy.data.WrappedPolicyKeys[idx].WrappedKey
+	policy.key, err = crypto.Unwrap(protector.key, wrappedPolicyKey)
+	return err
+}
+
 // Lock wipes a Policy's internal Key. It should always be called after using a
 // Policy. This is often done with a defer statement. There is no effect if
 // called multiple times.
diff --git a/actions/policy_test.go b/actions/policy_test.go
index 96b9bb0..11c9c3e 100644
--- a/actions/policy_test.go
+++ b/actions/policy_test.go
@@ -19,7 +19,11 @@
 
 package actions
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/pkg/errors"
+)
 
 // Makes a protector and policy
 func makeBoth() (*Protector, *Policy, error) {
@@ -139,3 +143,72 @@ func TestPolicyBadRemoveProtector(t *testing.T) {
 		t.Error("we should not be able to remove all the protectors from a policy")
 	}
 }
+
+// Tests that policy can be unlocked with a callback.
+func TestPolicyUnlockWithCallback(t *testing.T) {
+	// Our optionFunc just selects the first protector
+	optionFn := func(policyDescriptor string, options []*ProtectorOption) (int, error) {
+		return 0, nil
+	}
+
+	pro1, pol, err := makeBoth()
+	defer cleanupProtector(pro1)
+	defer cleanupPolicy(pol)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := pol.Lock(); err != nil {
+		t.Fatal(err)
+	}
+	if err := pol.Unlock(optionFn, goodCallback); err != nil {
+		t.Error(err)
+	}
+	if err := pol.Lock(); err != nil {
+		t.Error(err)
+	}
+}
+
+// Tests that policy can be unlock with an unlocked protector.
+func TestPolicyUnlockWithProtector(t *testing.T) {
+	pro1, pol, err := makeBoth()
+	defer cleanupProtector(pro1)
+	defer cleanupPolicy(pol)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := pol.Lock(); err != nil {
+		t.Fatal(err)
+	}
+	if err := pol.UnlockWithProtector(pro1); err != nil {
+		t.Error(err)
+	}
+	if err := pol.Lock(); err != nil {
+		t.Error(err)
+	}
+}
+
+// Tests that locked protectors cannot unlock a policy.
+func TestPolicyUnlockWithLockedProtector(t *testing.T) {
+	pro1, pol, err := makeBoth()
+	defer cleanupProtector(pro1)
+	defer cleanupPolicy(pol)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := pol.Lock(); err != nil {
+		t.Fatal(err)
+	}
+	if err := pro1.Lock(); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := pol.UnlockWithProtector(pro1); errors.Cause(err) != ErrLocked {
+		t.Errorf("Expected a cause of %v got %v", ErrLocked, err)
+		if err == nil {
+			pol.Lock()
+		}
+	}
+}