Specify public key for inactive shards in shard config (sigstore#746)

* Specify public key for each inactive shard in config

Signed-off-by: Priya Wadhwa <priya@chainguard.dev>

* Updated the integration test

Signed-off-by: Priya Wadhwa <priya@chainguard.dev>

* Add debugging to the sharding test

Signed-off-by: Priya Wadhwa <priya@chainguard.dev>

* Add debugging

Signed-off-by: Priya Wadhwa <priya@chainguard.dev>

Author: priyawadhwa

openapi.yaml

@@ -93,6 +93,12 @@ paths:
       operationId: getPublicKey
       tags:
         - pubkey
+      parameters:
+        - in: query
+          name: treeID
+          type: string
+          pattern: '^[0-9]+$'
+          description: The tree ID of the tree you wish to get a public key for
       produces:
         - application/x-pem-file
       responses:

pkg/api/public_key.go

@@ -17,10 +17,20 @@ limitations under the License.
 package api
 
 import (
+	"net/http"
+
 	"github.com/go-openapi/runtime/middleware"
+	"github.com/go-openapi/swag"
 	"github.com/sigstore/rekor/pkg/generated/restapi/operations/pubkey"
 )
 
 func GetPublicKeyHandler(params pubkey.GetPublicKeyParams) middleware.Responder {
-	return pubkey.NewGetPublicKeyOK().WithPayload(api.pubkey)
+	ctx := params.HTTPRequest.Context()
+	treeID := swag.StringValue(params.TreeID)
+	tc := NewTrillianClient(ctx)
+	pk, err := tc.ranges.PublicKey(api.pubkey, treeID)
+	if err != nil {
+		return handleRekorAPIError(params, http.StatusBadRequest, err, "")
+	}
+	return pubkey.NewGetPublicKeyOK().WithPayload(pk)
 }

pkg/generated/client/pubkey/get_public_key_parameters.go

@@ -16,9 +16,11 @@
 package sharding
 
 import (
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"io/ioutil"
+	"strconv"
 	"strings"
 
 	"github.com/ghodss/yaml"
@@ -33,8 +35,10 @@ type LogRanges struct {
 type Ranges []LogRange
 
 type LogRange struct {
-	TreeID     int64 `yaml:"treeID"`
-	TreeLength int64 `yaml:"treeLength"`
+	TreeID           int64  `yaml:"treeID"`
+	TreeLength       int64  `yaml:"treeLength"`
+	EncodedPublicKey string `yaml:"encodedPublicKey"`
+	decodedPublicKey string
 }
 
 func NewLogRanges(path string, treeID uint) (LogRanges, error) {
@@ -54,6 +58,14 @@ func NewLogRanges(path string, treeID uint) (LogRanges, error) {
 	if err := yaml.Unmarshal(contents, &ranges); err != nil {
 		return LogRanges{}, err
 	}
+	for i, r := range ranges {
+		decoded, err := base64.StdEncoding.DecodeString(r.EncodedPublicKey)
+		if err != nil {
+			return LogRanges{}, err
+		}
+		r.decodedPublicKey = string(decoded)
+		ranges[i] = r
+	}
 	return LogRanges{
 		inactive: ranges,
 		active:   int64(treeID),
@@ -119,3 +131,30 @@ func (l *LogRanges) String() string {
 	ranges = append(ranges, fmt.Sprintf("active=%d", l.active))
 	return strings.Join(ranges, ",")
 }
+
+// PublicKey returns the associated public key for the given Tree ID
+// and returns the active public key by default
+func (l *LogRanges) PublicKey(activePublicKey, treeID string) (string, error) {
+	// if no tree ID is specified, assume the active tree
+	if treeID == "" {
+		return activePublicKey, nil
+	}
+	tid, err := strconv.Atoi(treeID)
+	if err != nil {
+		return "", err
+	}
+
+	for _, i := range l.inactive {
+		if int(i.TreeID) == tid {
+			if i.decodedPublicKey != "" {
+				return i.decodedPublicKey, nil
+			}
+			// assume the active public key if one wasn't provided
+			return activePublicKey, nil
+		}
+	}
+	if tid == int(l.active) {
+		return activePublicKey, nil
+	}
+	return "", fmt.Errorf("%d is not a valid tree ID and doesn't have an associated public key", tid)
+}

pkg/generated/restapi/operations/pubkey/get_public_key_parameters.go

@@ -26,6 +26,7 @@ func TestNewLogRanges(t *testing.T) {
 	contents := `
 - treeID: 0001
   treeLength: 3
+  encodedPublicKey: c2hhcmRpbmcK
 - treeID: 0002
   treeLength: 4`
 	file := filepath.Join(t.TempDir(), "sharding-config")
@@ -36,8 +37,10 @@ func TestNewLogRanges(t *testing.T) {
 	expected := LogRanges{
 		inactive: []LogRange{
 			{
-				TreeID:     1,
-				TreeLength: 3,
+				TreeID:           1,
+				TreeLength:       3,
+				EncodedPublicKey: "c2hhcmRpbmcK",
+				decodedPublicKey: "sharding\n",
 			}, {
 				TreeID:     2,
 				TreeLength: 4,
@@ -94,3 +97,62 @@ func TestLogRanges_ResolveVirtualIndex(t *testing.T) {
 		}
 	}
 }
+
+func TestPublicKey(t *testing.T) {
+	ranges := LogRanges{
+		active: 45,
+		inactive: []LogRange{
+			{
+				TreeID:           10,
+				TreeLength:       10,
+				decodedPublicKey: "sharding",
+			}, {
+				TreeID:     20,
+				TreeLength: 20,
+			},
+		},
+	}
+	activePubKey := "activekey"
+	tests := []struct {
+		description    string
+		treeID         string
+		expectedPubKey string
+		shouldErr      bool
+	}{
+		{
+			description:    "empty tree ID",
+			expectedPubKey: "activekey",
+		}, {
+			description:    "tree id with decoded public key",
+			treeID:         "10",
+			expectedPubKey: "sharding",
+		}, {
+			description:    "tree id without decoded public key",
+			treeID:         "20",
+			expectedPubKey: "activekey",
+		}, {
+			description: "invalid tree id",
+			treeID:      "34",
+			shouldErr:   true,
+		}, {
+			description:    "pass in active tree id",
+			treeID:         "45",
+			expectedPubKey: "activekey",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.description, func(t *testing.T) {
+			got, err := ranges.PublicKey(activePubKey, test.treeID)
+			if err != nil && !test.shouldErr {
+				t.Fatal(err)
+			}
+			if test.shouldErr {
+				return
+			}
+			if got != test.expectedPubKey {
+				t.Fatalf("got %s doesn't match expected %s", got, test.expectedPubKey)
+			}
+		})
+	}
+}