Add a flag to start specifying log index ranges for virtual indices. (#435)

This is part of the temporal sharding work. The flag is not
hooked up anywhere yet.

Signed-off-by: Dan Lorenc <dlorenc@google.com>

cmd/rekor-server/app/flags.go

+//
+// Copyright 2021 The Sigstore Authors.
+//
+// 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 app
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+type LogRange struct {
+	TreeID     uint64
+	TreeLength uint64
+}
+
+type LogRanges struct {
+	Ranges []LogRange
+}
+
+func (l *LogRanges) Set(s string) error {
+	ranges := strings.Split(s, ",")
+	l.Ranges = []LogRange{}
+
+	var err error
+	inputRanges := []LogRange{}
+
+	// Only go up to the second to last one, the last one is special cased beow
+	for _, r := range ranges[:len(ranges)-1] {
+		split := strings.SplitN(r, "=", 2)
+		if len(split) != 2 {
+			return fmt.Errorf("invalid range flag, expected two parts separated by an =, got %s", r)
+		}
+		lr := LogRange{}
+		lr.TreeID, err = strconv.ParseUint(split[0], 10, 64)
+		if err != nil {
+			return err
+		}
+		lr.TreeLength, err = strconv.ParseUint(split[1], 10, 64)
+		if err != nil {
+			return err
+		}
+
+		inputRanges = append(inputRanges, lr)
+	}
+
+	// The last entry is special and should not have a terminating range, because this is active.
+	lastRangeStr := ranges[len(ranges)-1]
+	lastTreeID, err := strconv.ParseUint(lastRangeStr, 10, 64)
+	if err != nil {
+		return err
+	}
+
+	inputRanges = append(inputRanges, LogRange{
+		TreeID: lastTreeID,
+	})
+
+	// Look for duplicate tree ids
+	TreeIDs := map[uint64]struct{}{}
+	for _, lr := range inputRanges {
+		if _, ok := TreeIDs[lr.TreeID]; ok {
+			return fmt.Errorf("duplicate tree id: %d", lr.TreeID)
+		}
+		TreeIDs[lr.TreeID] = struct{}{}
+	}
+
+	l.Ranges = inputRanges
+	return nil
+}
+
+func (l *LogRanges) String() string {
+	ranges := []string{}
+	for _, r := range l.Ranges {
+		ranges = append(ranges, fmt.Sprintf("%d=%d", r.TreeID, r.TreeLength))
+	}
+	return strings.Join(ranges, ",")
+}
+
+func (l *LogRanges) Type() string {
+	return "LogRanges"
+}
+
+func (l *LogRanges) ResolveVirtualIndex(index int) (uint64, uint64) {
+	indexLeft := index
+	for _, l := range l.Ranges {
+		if indexLeft < int(l.TreeLength) {
+			return l.TreeID, uint64(indexLeft)
+		}
+		indexLeft -= int(l.TreeLength)
+	}
+
+	// Return the last one!
+	return l.Ranges[len(l.Ranges)-1].TreeID, uint64(indexLeft)
+}

cmd/rekor-server/app/flags_test.go

+//
+// Copyright 2021 The Sigstore Authors.
+//
+// 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 app
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestLogRanges_Set(t *testing.T) {
+	tests := []struct {
+		name string
+		arg  string
+		want []LogRange
+	}{
+		{
+			name: "one, no length",
+			arg:  "1234",
+			want: []LogRange{
+				{
+					TreeID:     1234,
+					TreeLength: 0,
+				},
+			},
+		},
+		{
+			name: "two",
+			arg:  "1234=10,7234",
+			want: []LogRange{
+				{
+					TreeID:     1234,
+					TreeLength: 10,
+				},
+				{
+					TreeID:     7234,
+					TreeLength: 0,
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			l := &LogRanges{}
+			if err := l.Set(tt.arg); err != nil {
+				t.Errorf("LogRanges.Set() expected no error, got %v", err)
+			}
+
+			if diff := cmp.Diff(tt.want, l.Ranges); diff != "" {
+				t.Errorf(diff)
+			}
+		})
+	}
+}
+
+func TestLogRanges_SetErr(t *testing.T) {
+	tests := []struct {
+		name string
+		arg  string
+	}{
+		{
+			name: "one, length (error)",
+			arg:  "1234=10",
+		},
+		{
+			name: "two, length (error)",
+			arg:  "1234=10,7234=17",
+		},
+		{
+			name: "invalid",
+			arg:  "1234=10,7234-17",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			l := &LogRanges{}
+			if err := l.Set(tt.arg); err == nil {
+				t.Error("LogRanges.Set() expected error but got none")
+			}
+		})
+	}
+}
+
+func TestLogRanges_ResolveVirtualIndex(t *testing.T) {
+	lrs := LogRanges{
+		Ranges: []LogRange{
+			{TreeID: 1, TreeLength: 17},
+			{TreeID: 2, TreeLength: 1},
+			{TreeID: 3, TreeLength: 100},
+			{TreeID: 4},
+		},
+	}
+
+	for _, tt := range []struct {
+		Index      int
+		WantTreeID uint64
+		WantIndex  uint64
+	}{
+		{
+			Index:      3,
+			WantTreeID: 1, WantIndex: 3,
+		},
+		// This is the first (0th) entry in the next tree
+		{
+			Index:      17,
+			WantTreeID: 2, WantIndex: 0,
+		},
+		// Overflow
+		{
+			Index:      3000,
+			WantTreeID: 4, WantIndex: 2882,
+		},
+	} {
+		tree, index := lrs.ResolveVirtualIndex(tt.Index)
+		if tree != tt.WantTreeID {
+			t.Errorf("LogRanges.ResolveVirtualIndex() tree = %v, want %v", tree, tt.WantTreeID)
+		}
+		if index != tt.WantIndex {
+			t.Errorf("LogRanges.ResolveVirtualIndex() index = %v, want %v", index, tt.WantIndex)
+		}
+	}
+}

cmd/rekor-server/app/root.go

@@ -29,6 +29,7 @@ import (
 
 var cfgFile string
 var logType string
+var logRangeMap LogRanges
 
 // rootCmd represents the base command when called without any subcommands
 var rootCmd = &cobra.Command{
@@ -60,6 +61,8 @@ func init() {
 	rootCmd.PersistentFlags().String("trillian_log_server.address", "127.0.0.1", "Trillian log server address")
 	rootCmd.PersistentFlags().Uint16("trillian_log_server.port", 8090, "Trillian log server port")
 	rootCmd.PersistentFlags().Uint("trillian_log_server.tlog_id", 0, "Trillian tree id")
+	rootCmd.PersistentFlags().Var(&logRangeMap, "trillian_log_server.log_id_ranges", "ordered list of tree ids and ranges")
+
 	rootCmd.PersistentFlags().String("rekor_server.hostname", "rekor.sigstore.dev", "public hostname of instance")
 	rootCmd.PersistentFlags().String("rekor_server.address", "127.0.0.1", "Address to bind to")
 	rootCmd.PersistentFlags().String("rekor_server.signer", "memory", "Rekor signer to use. Current valid options include: [gcpkms, memory]")