added refactored GetConsistencyProof

3 files changed, 158 insertions, 13 deletions

diff --git a/trillian/client.go b/trillian/client.go
index c22e9cc..7b8a00a 100644
--- a/trillian/client.go
+++ b/trillian/client.go
@@ -75,19 +75,36 @@ func (c *Client) GetTreeHead(ctx context.Context) (*types.TreeHead, error) {
 	if len(r.RootHash) != types.HashSize {
 		return nil, fmt.Errorf("unexpected hash length: %d", len(r.RootHash))
 	}
-
-	var hash [types.HashSize]byte
-	th := types.TreeHead{
-		Timestamp: uint64(r.TimestampNanos / 1000 / 1000 / 1000),
-		TreeSize:  uint64(r.TreeSize),
-		RootHash:  &hash,
-	}
-	copy(th.RootHash[:], r.RootHash)
-	return &th, nil
+	return treeHeadFromLogRoot(&r), nil
 }
 
 func (c *Client) GetConsistencyProof(ctx context.Context, req *types.ConsistencyProofRequest) (*types.ConsistencyProof, error) {
-	return nil, fmt.Errorf("TODO")
+	rsp, err := c.GRPC.GetConsistencyProof(ctx, &trillian.GetConsistencyProofRequest{
+		LogId:          c.TreeID,
+		FirstTreeSize:  int64(req.OldSize),
+		SecondTreeSize: int64(req.NewSize),
+	})
+	if err != nil {
+		return nil, fmt.Errorf("backend failure: %v", err)
+	}
+	if rsp == nil {
+		return nil, fmt.Errorf("no response")
+	}
+	if rsp.Proof == nil {
+		return nil, fmt.Errorf("no consistency proof")
+	}
+	if len(rsp.Proof.Hashes) == 0 {
+		return nil, fmt.Errorf("not a consistency proof: empty")
+	}
+	path, err := nodePathFromHashes(rsp.Proof.Hashes)
+	if err != nil {
+		return nil, fmt.Errorf("not a consistency proof: %v", err)
+	}
+	return &types.ConsistencyProof{
+		OldSize: req.OldSize,
+		NewSize: req.NewSize,
+		Path:    path,
+	}, nil
 }
 
 func (c *Client) GetInclusionProof(ctx context.Context, req *types.InclusionProofRequest) (*types.InclusionProof, error) {
diff --git a/trillian/client_test.go b/trillian/client_test.go
index 1807615..1d1c16f 100644
--- a/trillian/client_test.go
+++ b/trillian/client_test.go
@@ -191,6 +191,101 @@ func TestGetTreeHead(t *testing.T) {
 	}
 }
 
-func TestGetConsistencyProof(t *testing.T) {}
-func TestGetInclusionProof(t *testing.T)   {}
-func TestGetLeaves(t *testing.T)           {}
+func TestGetConsistencyProof(t *testing.T) {
+	req := &types.ConsistencyProofRequest{
+		OldSize: 1,
+		NewSize: 3,
+	}
+	for _, table := range []struct {
+		description string
+		req         *types.ConsistencyProofRequest
+		rsp         *trillian.GetConsistencyProofResponse
+		err         error
+		wantErr     bool
+		wantProof   *types.ConsistencyProof
+	}{
+		{
+			description: "invalid: backend failure",
+			req:         req,
+			err:         fmt.Errorf("something went wrong"),
+			wantErr:     true,
+		},
+		{
+			description: "invalid: no response",
+			req:         req,
+			wantErr:     true,
+		},
+		{
+			description: "invalid: no consistency proof",
+			req:         req,
+			rsp:         &trillian.GetConsistencyProofResponse{},
+			wantErr:     true,
+		},
+		{
+			description: "invalid: not a consistency proof (1/2)",
+			req:         req,
+			rsp: &trillian.GetConsistencyProofResponse{
+				Proof: &trillian.Proof{
+					Hashes: [][]byte{},
+				},
+			},
+			wantErr: true,
+		},
+		{
+			description: "invalid: not a consistency proof (2/2)",
+			req:         req,
+			rsp: &trillian.GetConsistencyProofResponse{
+				Proof: &trillian.Proof{
+					Hashes: [][]byte{
+						make([]byte, types.HashSize),
+						make([]byte, types.HashSize+1),
+					},
+				},
+			},
+			wantErr: true,
+		},
+		{
+			description: "valid",
+			req:         req,
+			rsp: &trillian.GetConsistencyProofResponse{
+				Proof: &trillian.Proof{
+					Hashes: [][]byte{
+						make([]byte, types.HashSize),
+						make([]byte, types.HashSize),
+					},
+				},
+			},
+			wantProof: &types.ConsistencyProof{
+				OldSize: 1,
+				NewSize: 3,
+				Path: []*[types.HashSize]byte{
+					&[types.HashSize]byte{},
+					&[types.HashSize]byte{},
+				},
+			},
+		},
+	} {
+		// Run deferred functions at the end of each iteration
+		func() {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+			grpc := mockclient.NewMockTrillianLogClient(ctrl)
+			grpc.EXPECT().GetConsistencyProof(gomock.Any(), gomock.Any()).Return(table.rsp, table.err)
+			client := Client{GRPC: grpc}
+
+			proof, err := client.GetConsistencyProof(context.Background(), table.req)
+			if got, want := err != nil, table.wantErr; got != want {
+				t.Errorf("got error %v but wanted %v in test %q: %v", got, want, table.description, err)
+			}
+			if err != nil {
+				return
+			}
+			if got, want := proof, table.wantProof; !reflect.DeepEqual(got, want) {
+				t.Errorf("got proof\n\t%v\nbut wanted\n\t%v\nin test %q", got, want, table.description)
+			}
+		}()
+	}
+}
+
+func TestGetInclusionProof(t *testing.T) {}
+func TestGetLeaves(t *testing.T)         {}
diff --git a/trillian/util.go b/trillian/util.go
new file mode 100644
index 0000000..87e64b6
--- /dev/null
+++ b/trillian/util.go
@@ -0,0 +1,33 @@
+package trillian
+
+import (
+	"fmt"
+
+	trillian "github.com/google/trillian/types"
+	siglog "github.com/system-transparency/stfe/types"
+)
+
+func treeHeadFromLogRoot(lr *trillian.LogRootV1) *siglog.TreeHead {
+	var hash [siglog.HashSize]byte
+	th := siglog.TreeHead{
+		Timestamp: uint64(lr.TimestampNanos / 1000 / 1000 / 1000),
+		TreeSize:  uint64(lr.TreeSize),
+		RootHash:  &hash,
+	}
+	copy(th.RootHash[:], lr.RootHash)
+	return &th
+}
+
+func nodePathFromHashes(hashes [][]byte) ([]*[siglog.HashSize]byte, error) {
+	var path []*[siglog.HashSize]byte
+	for _, hash := range hashes {
+		if len(hash) != siglog.HashSize {
+			return nil, fmt.Errorf("unexpected hash length: %v", len(hash))
+		}
+
+		var h [siglog.HashSize]byte
+		copy(h[:], hash)
+		path = append(path, &h)
+	}
+	return path, nil
+}