From 7089b59bd678260036bed0d54a838b4a8c445147 Mon Sep 17 00:00:00 2001
From: Gary Belvin <gdb@google.com>
Date: Wed, 6 Nov 2019 11:56:52 +0000
Subject: [PATCH] Implement randLog in keyserver

---
 core/keyserver/keyserver.go      | 21 +++++++++++++++++++--
 core/keyserver/revisions_test.go |  8 ++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/core/keyserver/keyserver.go b/core/keyserver/keyserver.go
index 441f4bd9b..d46a0c923 100644
--- a/core/keyserver/keyserver.go
+++ b/core/keyserver/keyserver.go
@@ -18,6 +18,7 @@ package keyserver
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"runtime"
 	"sync"
 	"time"
@@ -77,6 +78,8 @@ type MutationLogs interface {
 	// needed to do so.
 	ReadLog(ctx context.Context, directoryID string, logID int64, low, high time.Time,
 		limit int32) ([]*mutator.LogMessage, error)
+	// ListLogs returns a list of logs, optionally filtered by the writable bit.
+	ListLogs(ctx context.Context, directoryID string, writable bool) ([]int64, error)
 }
 
 // BatchReader reads batch definitions.
@@ -610,7 +613,6 @@ func (s *Server) QueueEntryUpdate(ctx context.Context, in *pb.UpdateEntryRequest
 	return s.BatchQueueUserUpdate(ctx, &pb.BatchQueueUserUpdateRequest{
 		DirectoryId: in.DirectoryId,
 		Updates:     []*pb.EntryUpdate{in.EntryUpdate},
-		LogId:       in.GetLogId(),
 	})
 }
 
@@ -657,7 +659,10 @@ func (s *Server) BatchQueueUserUpdate(ctx context.Context, in *pb.BatchQueueUser
 	tdone()
 
 	// Save mutation to the database.
-	wmLogID := in.GetLogId()
+	wmLogID, err := s.randLog(ctx, directory.DirectoryID)
+	if st := status.Convert(err); st.Code() != codes.OK {
+		return nil, status.Errorf(st.Code(), "Could not pick a log to write to: %v", err)
+	}
 	wmTime, err := s.logs.Send(ctx, directory.DirectoryID, wmLogID, in.Updates...)
 	if st := status.Convert(err); st.Code() != codes.OK {
 		glog.Errorf("mutations.Write failed: %v", err)
@@ -668,6 +673,18 @@ func (s *Server) BatchQueueUserUpdate(ctx context.Context, in *pb.BatchQueueUser
 	return &empty.Empty{}, nil
 }
 
+func (s *Server) randLog(ctx context.Context, directoryID string) (int64, error) {
+	// TODO(gbelvin): Cache these results.
+	writable := true
+	logIDs, err := s.logs.ListLogs(ctx, directoryID, writable)
+	if err != nil {
+		return 0, err
+	}
+
+	// Return a random log.
+	return logIDs[rand.Intn(len(logIDs))], nil
+}
+
 // GetDirectory returns all info tied to the specified directory.
 //
 // This API to get all necessary data needed to verify a particular
diff --git a/core/keyserver/revisions_test.go b/core/keyserver/revisions_test.go
index 19fd2ecaa..ea74b1d84 100644
--- a/core/keyserver/revisions_test.go
+++ b/core/keyserver/revisions_test.go
@@ -75,6 +75,14 @@ func (m *mutations) Send(ctx context.Context, dirID string, _ int64, mutation ..
 	return time.Time{}, errors.New("unimplemented")
 }
 
+func (m *mutations) ListLogs(ctx context.Context, dirID string, _ bool) ([]int64, error) {
+	logIDs := []int64{}
+	for id := range *m {
+		logIDs = append(logIDs, id)
+	}
+	return logIDs, nil
+}
+
 func (m *mutations) ReadLog(ctx context.Context, dirID string,
 	logID int64, low, high time.Time, batchSize int32) ([]*mutator.LogMessage, error) {
 	logShard := (*m)[logID]