Specify public key for inactive shards in shard config (#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>

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

@@ -74,6 +74,13 @@ func NewGetPublicKeyParamsWithHTTPClient(client *http.Client) *GetPublicKeyParam
    Typically these are written to a http.Request.
 */
 type GetPublicKeyParams struct {
+
+	/* TreeID.
+
+	   The tree ID of the tree you wish to get a public key for
+	*/
+	TreeID *string
+
 	timeout    time.Duration
 	Context    context.Context
 	HTTPClient *http.Client
@@ -127,6 +134,17 @@ func (o *GetPublicKeyParams) SetHTTPClient(client *http.Client) {
 	o.HTTPClient = client
 }
 
+// WithTreeID adds the treeID to the get public key params
+func (o *GetPublicKeyParams) WithTreeID(treeID *string) *GetPublicKeyParams {
+	o.SetTreeID(treeID)
+	return o
+}
+
+// SetTreeID adds the treeId to the get public key params
+func (o *GetPublicKeyParams) SetTreeID(treeID *string) {
+	o.TreeID = treeID
+}
+
 // WriteToRequest writes these params to a swagger request
 func (o *GetPublicKeyParams) WriteToRequest(r runtime.ClientRequest, reg strfmt.Registry) error {
 
@@ -135,6 +153,23 @@ func (o *GetPublicKeyParams) WriteToRequest(r runtime.ClientRequest, reg strfmt.
 	}
 	var res []error
 
+	if o.TreeID != nil {
+
+		// query param treeID
+		var qrTreeID string
+
+		if o.TreeID != nil {
+			qrTreeID = *o.TreeID
+		}
+		qTreeID := qrTreeID
+		if qTreeID != "" {
+
+			if err := r.SetQueryParam("treeID", qTreeID); err != nil {
+				return err
+			}
+		}
+	}
+
 	if len(res) > 0 {
 		return errors.CompositeValidationError(res...)
 	}

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

@@ -25,7 +25,10 @@ import (
 	"net/http"
 
 	"github.com/go-openapi/errors"
+	"github.com/go-openapi/runtime"
 	"github.com/go-openapi/runtime/middleware"
+	"github.com/go-openapi/strfmt"
+	"github.com/go-openapi/validate"
 )
 
 // NewGetPublicKeyParams creates a new GetPublicKeyParams object
@@ -44,6 +47,12 @@ type GetPublicKeyParams struct {
 
 	// HTTP Request Object
 	HTTPRequest *http.Request `json:"-"`
+
+	/*The tree ID of the tree you wish to get a public key for
+	  Pattern: ^[0-9]+$
+	  In: query
+	*/
+	TreeID *string
 }
 
 // BindRequest both binds and validates a request, it assumes that complex things implement a Validatable(strfmt.Registry) error interface
@@ -55,8 +64,46 @@ func (o *GetPublicKeyParams) BindRequest(r *http.Request, route *middleware.Matc
 
 	o.HTTPRequest = r
 
+	qs := runtime.Values(r.URL.Query())
+
+	qTreeID, qhkTreeID, _ := qs.GetOK("treeID")
+	if err := o.bindTreeID(qTreeID, qhkTreeID, route.Formats); err != nil {
+		res = append(res, err)
+	}
 	if len(res) > 0 {
 		return errors.CompositeValidationError(res...)
 	}
 	return nil
 }
+
+// bindTreeID binds and validates parameter TreeID from query.
+func (o *GetPublicKeyParams) bindTreeID(rawData []string, hasKey bool, formats strfmt.Registry) error {
+	var raw string
+	if len(rawData) > 0 {
+		raw = rawData[len(rawData)-1]
+	}
+
+	// Required: false
+	// AllowEmptyValue: false
+
+	if raw == "" { // empty values pass all other validations
+		return nil
+	}
+	o.TreeID = &raw
+
+	if err := o.validateTreeID(formats); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// validateTreeID carries on validations for parameter TreeID
+func (o *GetPublicKeyParams) validateTreeID(formats strfmt.Registry) error {
+
+	if err := validate.Pattern("treeID", "query", *o.TreeID, `^[0-9]+$`); err != nil {
+		return err
+	}
+
+	return nil
+}

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

@@ -29,7 +29,11 @@ import (
 
 // GetPublicKeyURL generates an URL for the get public key operation
 type GetPublicKeyURL struct {
+	TreeID *string
+
 	_basePath string
+	// avoid unkeyed usage
+	_ struct{}
 }
 
 // WithBasePath sets the base path for this url builder, only required when it's different from the
@@ -56,6 +60,18 @@ func (o *GetPublicKeyURL) Build() (*url.URL, error) {
 	_basePath := o._basePath
 	_result.Path = golangswaggerpaths.Join(_basePath, _path)
 
+	qs := make(url.Values)
+
+	var treeIDQ string
+	if o.TreeID != nil {
+		treeIDQ = *o.TreeID
+	}
+	if treeIDQ != "" {
+		qs.Set("treeID", treeIDQ)
+	}
+
+	_result.RawQuery = qs.Encode()
+
 	return &_result, nil
 }
 

pkg/sharding/ranges.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/sharding/ranges_test.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)
+			}
+		})
+	}
+}

tests/sharding-e2e-test.sh

@@ -48,22 +48,41 @@ function check_log_index () {
   fi
 }
 
-count=0
-
-echo -n "waiting up to 60 sec for system to start"
-until [ $(docker-compose ps | grep -c "(healthy)") == 3 ];
-do
-    if [ $count -eq 6 ]; then
-       echo "! timeout reached"
-       exit 1
-    else
-       echo -n "."
-       sleep 10
-       let 'count+=1'
-    fi
-done
-
-echo
+function stringsMatch () {
+  one=$1
+  two=$2
+
+  if [[ "$one" == "$two" ]]; then
+    echo "Strings match"
+  else
+    echo "$one and $two don't match but should"
+    exit 1
+  fi
+}
+
+function waitForRekorServer () {
+  count=0
+
+  echo -n "waiting up to 60 sec for system to start"
+  until [ $(docker-compose ps | grep -c "(healthy)") == 3 ];
+  do
+      if [ $count -eq 6 ]; then
+        echo "! timeout reached"
+        REKOR_CONTAINER_ID=$(docker ps --filter name=rekor-server --format {{.ID}})
+        docker logs $REKOR_CONTAINER_ID
+        exit 1
+      else
+        echo -n "."
+        sleep 10
+        let 'count+=1'
+      fi
+  done
+
+  echo
+}
+
+echo "Waiting for rekor server to come up..."
+waitForRekorServer
 
 # Add some things to the tlog :)
 pushd tests
@@ -94,7 +113,10 @@ SHARD_TREE_ID=$(createtree --admin_server localhost:8090)
 echo "the new shard ID is $SHARD_TREE_ID"
 
 # Once more
-$REKOR_CLI loginfo --rekor_server http://localhost:3000 
+$REKOR_CLI loginfo --rekor_server http://localhost:3000
+
+# Get the public key for the active tree for later
+ENCODED_PUBLIC_KEY=$(curl http://localhost:3000/api/v1/log/publicKey | base64 -w 0)
 
 # Spin down the rekor server
 echo "stopping the rekor server..."
@@ -107,8 +129,10 @@ SHARDING_CONFIG=sharding-config.yaml
 cat << EOF > $SHARDING_CONFIG
 - treeID: $INITIAL_TREE_ID
   treeLength: 3
+  encodedPublicKey: $ENCODED_PUBLIC_KEY
 EOF
 
+cat $SHARDING_CONFIG
 
 COMPOSE_FILE=docker-compose-sharding.yaml
 cat << EOF > $COMPOSE_FILE
@@ -152,18 +176,13 @@ EOF
 # Spin up the new Rekor
 
 docker-compose -f $COMPOSE_FILE up  -d
-sleep 15
+waitForRekorServer
 $REKOR_CLI loginfo --rekor_server http://localhost:3000 
 
 # Make sure we are pointing to the new tree now
-TREE_ID=$($REKOR_CLI loginfo --rekor_server http://localhost:3000  --format json)
+TREE_ID=$($REKOR_CLI loginfo --rekor_server http://localhost:3000  --format json | jq -r .TreeID)
 # Check that the SHARD_TREE_ID is a substring of the `$REKOR_CLI loginfo` output
-if [[ "$TREE_ID" == *"$SHARD_TREE_ID"* ]]; then
-  echo "Rekor server is now pointing to the new shard"
-else
-  echo "Rekor server is not pointing to the new shard"
-  exit 1
-fi
+stringsMatch $TREE_ID $SHARD_TREE_ID
 
 # Now, if we run $REKOR_CLI get --log_index 2 again, it should grab the log index
 # from Shard 0
@@ -181,7 +200,25 @@ $REKOR_CLI logproof --last-size 2 --tree-id $INITIAL_TREE_ID --rekor_server http
 # And the logproof for the now active shard
 $REKOR_CLI logproof --last-size 1 --rekor_server http://localhost:3000
 
+echo "Getting public key for inactive shard..."
+GOT_PUB_KEY=$(curl "http://localhost:3000/api/v1/log/publicKey?treeID=$INITIAL_TREE_ID" | base64 -w 0)
+echo "Got encoded public key $GOT_PUB_KEY, making sure this matches the public key we got earlier..."
+stringsMatch $ENCODED_PUBLIC_KEY $GOT_PUB_KEY
+
+echo "Getting the public key for the active tree..."
+NEW_PUB_KEY=$(curl "http://localhost:3000/api/v1/log/publicKey" | base64 -w 0)
+echo "Making sure the public key for the active shard is different from the inactive shard..."
+if [[ "$ENCODED_PUBLIC_KEY" == "$NEW_PUB_KEY" ]]; then
+    echo
+    echo "Active tree public key should be different from inactive shard public key but isn't..."
+    echo "Inactive Shard Public Key: $ENCODED_PUBLIC_KEY"
+    echo "Active Shard Public Key: $NEW_PUB_KEY"
+    exit 1
+fi
+
+
 # TODO: Try to get the entry via Entry ID (Tree ID in hex + UUID)
 UUID=$($REKOR_CLI get --log-index 2 --rekor_server http://localhost:3000 --format json | jq -r .UUID)
 
+
 echo "Test passed successfully :)"