Client with (partially) verified AddLeaf

Client AddLeaf:
- Fetches current Signed Tree Root
- Queues the leaf
- Waits for an STR update
- TODO: Verify STR
- Fetch and verify inclusion proof

Integration test improvements:
- LogEnv also creates a signer

client/client.go

@@ -18,40 +18,152 @@ package client
 import (
 	"context"
 	"crypto/sha256"
+	"time"
 
 	"github.com/google/trillian"
+	"github.com/google/trillian/merkle"
+	"github.com/jpillora/Backoff"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 )
 
 // LogClient represents a client for a given Trillian log instance.
 type LogClient struct {
-	LogID  int64
-	client trillian.TrillianLogClient
+	LogID    int64
+	client   trillian.TrillianLogClient
+	hasher   merkle.TreeHasher
+	STR      trillian.SignedLogRoot
+	MaxTries int
 }
 
 // New returns a new LogClient.
-func New(logID int64, cc *grpc.ClientConn) *LogClient {
+func New(logID int64, cc *grpc.ClientConn, hasher merkle.TreeHasher) *LogClient {
 	return &LogClient{
-		LogID:  logID,
-		client: trillian.NewTrillianLogClient(cc),
+		LogID:    logID,
+		client:   trillian.NewTrillianLogClient(cc),
+		hasher:   hasher,
+		MaxTries: 3,
 	}
 }
 
-// AddLeaf adds leaf to the append only log. It blocks until a verifiable response is received.
+// AddLeaf adds leaf to the append only log.
+// Blocks until it gets a verifiable response.
 func (c *LogClient) AddLeaf(ctx context.Context, data []byte) error {
+	// Fetch the current STR so we can detect when we update.
+	if err := c.UpdateSTR(ctx); err != nil {
+		return err
+	}
+
+	leaf := buildLeaf(data)
+	err := c.queueLeaf(ctx, leaf)
+	switch {
+	case grpc.Code(err) == codes.AlreadyExists:
+		// If the leaf already exists, don't wait for an update.
+		return c.getInclusionProof(ctx, leaf.MerkleLeafHash, c.STR.TreeSize)
+	case err != nil:
+		return err
+	default:
+		err := grpc.Errorf(codes.NotFound, "Pre-loop condition")
+		for i := 0; grpc.Code(err) == codes.NotFound && i < c.MaxTries; i++ {
+			// Wait for TreeSize to update.
+			if err := c.waitForSTRUpdate(ctx, c.MaxTries); err != nil {
+				return err
+			}
+
+			// Get proof by hash.
+			err = c.getInclusionProof(ctx, leaf.MerkleLeafHash, c.STR.TreeSize)
+		}
+		return err
+	}
+}
+
+// waitForSTRUpdate repeatedly fetches the STR until the TreeSize changes
+// or until a maximum number of attempts has been tried.
+func (c *LogClient) waitForSTRUpdate(ctx context.Context, attempts int) error {
+	b := &backoff.Backoff{
+		Min:    100 * time.Millisecond,
+		Max:    10 * time.Second,
+		Factor: 2,
+		Jitter: true,
+	}
+	startTreeSize := c.STR.TreeSize
+	for i := 0; ; i++ {
+		if err := c.UpdateSTR(ctx); err != nil {
+			return err
+		}
+		if c.STR.TreeSize > startTreeSize {
+			return nil
+		}
+		if i >= (attempts - 1) {
+			return grpc.Errorf(codes.DeadlineExceeded,
+				"UpdateSTR().TreeSize: %v, want > %v. Tried %v times.",
+				c.STR.TreeSize, startTreeSize, i+1)
+		}
+		time.Sleep(b.Duration())
+	}
+}
+
+// UpdateSTR retrieves the current SignedLogRoot and verifies it.
+func (c *LogClient) UpdateSTR(ctx context.Context) error {
+	req := &trillian.GetLatestSignedLogRootRequest{
+		LogId: c.LogID,
+	}
+	resp, err := c.client.GetLatestSignedLogRoot(ctx, req)
+	if err != nil {
+		return err
+	}
+	// TODO(gdbelvin): Verify SignedLogRoot
+
+	c.STR = *resp.SignedLogRoot
+	return nil
+}
+
+func (c *LogClient) getInclusionProof(ctx context.Context, leafHash []byte, treeSize int64) error {
+	req := &trillian.GetInclusionProofByHashRequest{
+		LogId:    c.LogID,
+		LeafHash: leafHash,
+		TreeSize: treeSize,
+	}
+	resp, err := c.client.GetInclusionProofByHash(ctx, req)
+	if err != nil {
+		return err
+	}
+	for _, proof := range resp.Proof {
+		neighbors := convertProof(proof)
+		v := merkle.NewLogVerifier(c.hasher)
+		if err := v.VerifyInclusionProof(proof.LeafIndex, treeSize, neighbors, c.STR.RootHash, leafHash); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// convertProof returns a slice of neighbor nodes from a trillian Proof.
+// TODO(martin): adjust the public API to do this in the server before returing a proof.
+func convertProof(proof *trillian.Proof) [][]byte {
+	neighbors := make([][]byte, len(proof.ProofNode))
+	for i, node := range proof.ProofNode {
+		neighbors[i] = node.NodeHash
+	}
+	return neighbors
+}
+
+func buildLeaf(data []byte) *trillian.LogLeaf {
 	hash := sha256.Sum256(data)
 	leaf := &trillian.LogLeaf{
 		LeafValue:        data,
 		MerkleLeafHash:   hash[:],
 		LeafIdentityHash: hash[:],
 	}
+	return leaf
+}
+
+func (c *LogClient) queueLeaf(ctx context.Context, leaf *trillian.LogLeaf) error {
+	// Queue Leaf
 	req := trillian.QueueLeafRequest{
 		LogId: c.LogID,
 		Leaf:  leaf,
 	}
 	_, err := c.client.QueueLeaf(ctx, &req)
-	// TODO(gdbelvin): Get proof by hash
-	// TODO(gdbelvin): backoff with jitter
-	// TODO(gdbelvin): verify proof
 	return err
 }

client/client_test.go

@@ -18,47 +18,65 @@ import (
 	"context"
 	"testing"
 
-	"github.com/google/trillian/storage/mysql"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+
+	"github.com/golang/glog"
+	"github.com/google/trillian/testonly"
 	"github.com/google/trillian/testonly/integration"
 )
 
+const logID = int64(1234)
+
+func TestAddGetLeaf(t *testing.T) {
+	// TODO: Build a GetLeaf method and test a full get/set cycle.
+}
+
 func TestAddLeaf(t *testing.T) {
 	ctx := context.Background()
-	logID := int64(1234)
-	env, err := integration.NewLogEnv("client")
+	env, err := integration.NewLogEnv("TestAddLeaf")
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer env.Close()
-	if err := mysql.CreateTree(logID, env.DB); err != nil {
+	if err := env.CreateLog(logID); err != nil {
 		t.Errorf("Failed to create log: %v", err)
 	}
 
-	client := New(logID, env.ClientConn)
+	client := New(logID, env.ClientConn, testonly.Hasher)
+	client.MaxTries = 1
 
+	if err, want := client.AddLeaf(ctx, []byte("foo")), codes.DeadlineExceeded; grpc.Code(err) != want {
+		t.Errorf("AddLeaf(): %v, want, %v", err, want)
+	}
+	env.Sequencer.OperationLoop() // Sequence the new node.
+	glog.Infof("try AddLeaf again")
 	if err := client.AddLeaf(ctx, []byte("foo")); err != nil {
 		t.Errorf("Failed to add Leaf: %v", err)
 	}
 }
 
-func TestAddSameLeaf(t *testing.T) {
+func TestUpdateSTR(t *testing.T) {
 	ctx := context.Background()
-	logID := int64(1234)
-	t.Skip("Submitting two leaves currently breaks")
-	env, err := integration.NewLogEnv("client")
+	env, err := integration.NewLogEnv("TestUpdateSTR")
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer env.Close()
-	client := New(logID, env.ClientConn)
-
-	if err := mysql.CreateTree(logID, env.DB); err != nil {
+	if err := env.CreateLog(logID); err != nil {
 		t.Errorf("Failed to create log: %v", err)
 	}
-	if err := client.AddLeaf(ctx, []byte("foo")); err != nil {
-		t.Error(err)
+	client := New(logID, env.ClientConn, testonly.Hasher)
+
+	before := client.STR.TreeSize
+	if err, want := client.AddLeaf(ctx, []byte("foo")), codes.DeadlineExceeded; grpc.Code(err) != want {
+		t.Errorf("AddLeaf(): %v, want, %v", err, want)
 	}
-	if err := client.AddLeaf(ctx, []byte("foo")); err != nil {
+	env.Sequencer.OperationLoop() // Sequence the new node.
+	if err := client.UpdateSTR(ctx); err != nil {
 		t.Error(err)
 	}
+	if got, want := client.STR.TreeSize, before; got <= want {
+		t.Errorf("Tree size after add Leaf: %v, want > %v", got, want)
+	}
 }

log/sequencer.go

@@ -205,6 +205,7 @@ func (s Sequencer) SequenceBatch(ctx context.Context, logID int64, limit int) (i
 	// current one is too old. If there's work to be done then we'll be creating a root anyway.
 	if len(leaves) == 0 {
 		// We have nothing to integrate into the tree
+		glog.Infof("No leaves sequenced in this signing operation.")
 		return 0, tx.Commit()
 	}
 

merkle/log_verifier.go

@@ -43,8 +43,8 @@ func NewLogVerifier(hasher TreeHasher) LogVerifier {
 }
 
 // VerifyInclusionProof verifies the correctness of the proof given the passed in information about the tree and leaf.
-func (v LogVerifier) VerifyInclusionProof(leafIndex, treeSize int64, proof [][]byte, root []byte, leaf []byte) error {
-	calcRoot, err := v.RootFromInclusionProof(leafIndex, treeSize, proof, leaf)
+func (v LogVerifier) VerifyInclusionProof(leafIndex, treeSize int64, proof [][]byte, root []byte, leafHash []byte) error {
+	calcRoot, err := v.RootFromInclusionProof(leafIndex, treeSize, proof, leafHash)
 	if err != nil {
 		return err
 	}
@@ -60,7 +60,7 @@ func (v LogVerifier) VerifyInclusionProof(leafIndex, treeSize int64, proof [][]b
 // RootFromInclusionProof calculates the expected tree root given the proof and leaf.
 // leafIndex starts at 0.  treeSize is the number of nodes in the tree.
 // proof is an array of neighbor nodes from the bottom to the root.
-func (v LogVerifier) RootFromInclusionProof(leafIndex, treeSize int64, proof [][]byte, leaf []byte) ([]byte, error) {
+func (v LogVerifier) RootFromInclusionProof(leafIndex, treeSize int64, proof [][]byte, leafHash []byte) ([]byte, error) {
 	if leafIndex < 0 {
 		return nil, errors.New("invalid leafIndex < 0")
 	}
@@ -73,7 +73,7 @@ func (v LogVerifier) RootFromInclusionProof(leafIndex, treeSize int64, proof [][
 	}
 
 	cntIndex := leafIndex
-	cntHash := v.hasher.HashLeaf(leaf)
+	cntHash := leafHash
 	proofIndex := 0
 
 	// Tree is numbered as follows, where nodes at each level are counted from left to right.