Permalink
Browse files

Add GetLeavesByRevision to fake.TrillianMapClient (#872)

* Add GetLeavesByRevision to fake.TrillianMapClient

Also moves the fake TrillianMapClient into core/fake
and cleans up the test to use a proper client interface.

* pass options in grpc calls
  • Loading branch information...
gdbelvin committed Dec 7, 2017
1 parent 0cbe7d5 commit 0b84de20391fbb312469bbdaaf4ed183be74f90a
Showing with 100 additions and 52 deletions.
  1. +89 −0 core/fake/trillian_map_client.go
  2. +11 −52 core/mutationserver/mutation_test.go
@@ -0,0 +1,89 @@
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package fake
import (
"context"
"google.golang.org/grpc"
tpb "github.com/google/trillian"
)
type mapServer struct {
// roots by revision number
roots map[int64]*tpb.SignedMapRoot
revision int64
}
// NewTrillianMapClient returns a fake tpb.TrillianMapClient
// This client only stores tpb.MapperMetadata in roots. It does not store
// leaves compute inclusion proofs, or sign roots. This client is not
// threadsafe.
func NewTrillianMapClient() tpb.TrillianMapClient {
m := &mapServer{
roots: make(map[int64]*tpb.SignedMapRoot),
}
m.roots[0] = &tpb.SignedMapRoot{} // Set the initial root
return m
}
// GetLeaves just returns the indexes requested. No leaf data, no inclusion proofs.
func (m *mapServer) GetLeaves(ctx context.Context, in *tpb.GetMapLeavesRequest, opts ...grpc.CallOption) (*tpb.GetMapLeavesResponse, error) {
return m.GetLeavesByRevision(ctx, &tpb.GetMapLeavesByRevisionRequest{
MapId: in.MapId,
Index: in.Index,
Revision: m.revision,
}, opts...)
}
// GetLeavesByRevision just returns the indexes requested. No leaf data, no inclusion proofs.
func (m *mapServer) GetLeavesByRevision(ctx context.Context, in *tpb.GetMapLeavesByRevisionRequest, opts ...grpc.CallOption) (*tpb.GetMapLeavesResponse, error) {
leaves := make([]*tpb.MapLeafInclusion, 0, len(in.Index))
for _, index := range in.Index {
leaves = append(leaves, &tpb.MapLeafInclusion{
Leaf: &tpb.MapLeaf{
Index: index,
},
})
}
return &tpb.GetMapLeavesResponse{
MapLeafInclusion: leaves,
}, nil
}
// SetLeaves is not thread safe. It will store the root metadata.
func (m *mapServer) SetLeaves(ctx context.Context, in *tpb.SetMapLeavesRequest, opts ...grpc.CallOption) (*tpb.SetMapLeavesResponse, error) {
m.revision++
m.roots[m.revision] = &tpb.SignedMapRoot{
Metadata: in.Metadata,
MapRevision: m.revision,
}
return nil, nil
}
// GetSignedMapRootByRevision returns the current map root.
func (m *mapServer) GetSignedMapRoot(ctx context.Context, in *tpb.GetSignedMapRootRequest, opts ...grpc.CallOption) (*tpb.GetSignedMapRootResponse, error) {
return m.GetSignedMapRootByRevision(ctx, &tpb.GetSignedMapRootByRevisionRequest{
Revision: m.revision,
}, opts...)
}
// GetSignedMapRootByRevision returns the saved map root.
func (m *mapServer) GetSignedMapRootByRevision(ctx context.Context, in *tpb.GetSignedMapRootByRevisionRequest, opts ...grpc.CallOption) (*tpb.GetSignedMapRootResponse, error) {
return &tpb.GetSignedMapRootResponse{
MapRoot: m.roots[in.Revision],
}, nil
}
@@ -34,7 +34,7 @@ import (
"google.golang.org/grpc/codes"
pb "github.com/google/keytransparency/core/proto/keytransparency_v1_proto"
"github.com/google/trillian"
tpb "github.com/google/trillian"
)
const (
@@ -57,24 +57,23 @@ func updates(t *testing.T, start, end int) []*pb.EntryUpdate {
return kvs
}
func prepare(t *testing.T, mapID int64, mutations mutator.MutationStorage, fakeMap *fakeTrillianMapClient) {
createEpoch(t, mapID, mutations, fakeMap, 1, 1, 6)
createEpoch(t, mapID, mutations, fakeMap, 2, 7, 10)
func prepare(ctx context.Context, t *testing.T, mapID int64, mutations mutator.MutationStorage, tmap tpb.TrillianMapClient) {
createEpoch(ctx, t, mapID, mutations, tmap, 1, 1, 6)
createEpoch(ctx, t, mapID, mutations, tmap, 2, 7, 10)
}
func createEpoch(t *testing.T, mapID int64, mutations mutator.MutationStorage, fakeMap *fakeTrillianMapClient, epoch int64, start, end int) {
func createEpoch(ctx context.Context, t *testing.T, mapID int64, mutations mutator.MutationStorage, tmap tpb.TrillianMapClient, epoch int64, start, end int) {
kvs := updates(t, start, end)
for _, kv := range kvs {
if _, err := mutations.Write(nil, mapID, kv); err != nil {
t.Fatalf("mutations.Write failed: %v", err)
}
}
fakeMap.tmap[epoch] = &trillian.SignedMapRoot{
tmap.SetLeaves(ctx, &tpb.SetMapLeavesRequest{
Metadata: mustMetadataAsAny(t, &pb.MapperMetadata{
HighestFullyCompletedSeq: int64(end),
}),
MapRevision: epoch,
}
})
}
func mustMetadataAsAny(t *testing.T, meta *pb.MapperMetadata) *any.Any {
@@ -98,11 +97,11 @@ func TestGetMutations(t *testing.T) {
mapID := int64(2)
fakeMutations := fake.NewMutationStorage()
fakeAdmin := fake.NewAdminStorage()
fakeMap := newFakeTrillianMapClient()
fakeMap := fake.NewTrillianMapClient()
if err := fakeAdmin.Write(ctx, domainID, mapID, 0, nil, nil, 1*time.Second, 5*time.Second); err != nil {
t.Fatalf("admin.Write(): %v", err)
}
prepare(t, mapID, fakeMutations, fakeMap)
prepare(ctx, t, mapID, fakeMutations, fakeMap)
for _, tc := range []struct {
description string
@@ -159,10 +158,10 @@ func TestGetMutations(t *testing.T) {
func TestLowestSequenceNumber(t *testing.T) {
ctx := context.Background()
fakeMutations := fake.NewMutationStorage()
fakeMap := newFakeTrillianMapClient()
fakeMap := fake.NewTrillianMapClient()
fakeAdmin := &fake.AdminStorage{}
mapID := int64(1)
prepare(t, mapID, fakeMutations, fakeMap)
prepare(ctx, t, mapID, fakeMutations, fakeMap)
for _, tc := range []struct {
token string
@@ -187,46 +186,6 @@ func TestLowestSequenceNumber(t *testing.T) {
}
}
// trillian.TrillianMapClient fake.
type fakeTrillianMapClient struct {
tmap map[int64]*trillian.SignedMapRoot
}
func newFakeTrillianMapClient() *fakeTrillianMapClient {
return &fakeTrillianMapClient{
tmap: make(map[int64]*trillian.SignedMapRoot),
}
}
func (*fakeTrillianMapClient) GetLeaves(ctx context.Context, in *trillian.GetMapLeavesRequest, opts ...grpc.CallOption) (*trillian.GetMapLeavesResponse, error) {
leaves := make([]*trillian.MapLeafInclusion, 0, len(in.Index))
for _, index := range in.Index {
leaves = append(leaves, &trillian.MapLeafInclusion{
Leaf: &trillian.MapLeaf{
Index: index,
},
})
}
return &trillian.GetMapLeavesResponse{
MapLeafInclusion: leaves,
}, nil
}
func (*fakeTrillianMapClient) SetLeaves(ctx context.Context, in *trillian.SetMapLeavesRequest, opts ...grpc.CallOption) (*trillian.SetMapLeavesResponse, error) {
return nil, nil
}
func (*fakeTrillianMapClient) GetSignedMapRoot(ctx context.Context, in *trillian.GetSignedMapRootRequest, opts ...grpc.CallOption) (*trillian.GetSignedMapRootResponse, error) {
return nil, nil
}
// GetSignedMapRootByRevision echos in.MapId in HighestFullyCompletedSeq.
func (m *fakeTrillianMapClient) GetSignedMapRootByRevision(ctx context.Context, in *trillian.GetSignedMapRootByRevisionRequest, opts ...grpc.CallOption) (*trillian.GetSignedMapRootResponse, error) {
return &trillian.GetSignedMapRootResponse{
MapRoot: m.tmap[in.Revision],
}, nil
}
// transaction.Txn fake.
type fakeTxn struct{}

0 comments on commit 0b84de2

Please sign in to comment.