From 18cf995ef9426c4f99dfee8330aa63d7d89fcaef Mon Sep 17 00:00:00 2001
From: Dan Lorenc <dlorenc@google.com>
Date: Fri, 9 Jul 2021 21:31:57 -0500
Subject: [PATCH] Add index keys for in-toto provenance objects.

Signed-off-by: Dan Lorenc <dlorenc@google.com>
---
 pkg/types/intoto/v0.0.1/entry.go      | 36 ++++++++++++++++-
 pkg/types/intoto/v0.0.1/entry_test.go | 58 +++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/pkg/types/intoto/v0.0.1/entry.go b/pkg/types/intoto/v0.0.1/entry.go
index f9e44fb8e..de3d10192 100644
--- a/pkg/types/intoto/v0.0.1/entry.go
+++ b/pkg/types/intoto/v0.0.1/entry.go
@@ -20,10 +20,12 @@ import (
 	"context"
 	"crypto"
 	"crypto/sha256"
+	"encoding/base64"
 	"encoding/hex"
 	"encoding/json"
 	"errors"
 
+	"github.com/in-toto/in-toto-golang/in_toto"
 	"github.com/in-toto/in-toto-golang/pkg/ssl"
 	"github.com/spf13/viper"
 
@@ -69,11 +71,39 @@ func (v V001Entry) IndexKeys() []string {
 	var result []string
 
 	h := sha256.Sum256([]byte(v.env.Payload))
-	payloadKey := "sha256:" + string(h[:])
+	payloadKey := "sha256:" + hex.EncodeToString(h[:])
 	result = append(result, payloadKey)
+
+	switch v.env.PayloadType {
+	case in_toto.PayloadType:
+		statement, err := parseStatement(v.env.Payload)
+		if err != nil {
+			log.Logger.Info("invalid id in_toto Statement")
+			return result
+		}
+		for _, s := range statement.Subject {
+			for alg, ds := range s.Digest {
+				result = append(result, alg+":"+ds)
+			}
+		}
+	default:
+		log.Logger.Infof("Unknown in_toto Statement Type: %s", v.env.PayloadType)
+	}
 	return result
 }
 
+func parseStatement(p string) (*in_toto.Statement, error) {
+	ps := in_toto.Statement{}
+	payload, err := base64.StdEncoding.DecodeString(p)
+	if err != nil {
+		return nil, err
+	}
+	if err := json.Unmarshal(payload, &ps); err != nil {
+		return nil, err
+	}
+	return &ps, nil
+}
+
 func (v *V001Entry) Unmarshal(pe models.ProposedEntry) error {
 	it, ok := pe.(*models.Intoto)
 	if !ok {
@@ -155,6 +185,10 @@ func (v *V001Entry) Validate() error {
 		return err
 	}
 
+	if v.IntotoObj.Content.Envelope == "" {
+		return nil
+	}
+
 	if err := json.Unmarshal([]byte(v.IntotoObj.Content.Envelope), &v.env); err != nil {
 		return err
 	}
diff --git a/pkg/types/intoto/v0.0.1/entry_test.go b/pkg/types/intoto/v0.0.1/entry_test.go
index 98c3c9d68..408868b5d 100644
--- a/pkg/types/intoto/v0.0.1/entry_test.go
+++ b/pkg/types/intoto/v0.0.1/entry_test.go
@@ -22,6 +22,8 @@ import (
 	"crypto/rand"
 	"crypto/sha256"
 	"crypto/x509"
+	"encoding/base64"
+	"encoding/hex"
 	"encoding/json"
 	"encoding/pem"
 	"fmt"
@@ -184,3 +186,59 @@ func TestV001Entry_Unmarshal(t *testing.T) {
 		})
 	}
 }
+
+func TestV001Entry_IndexKeys(t *testing.T) {
+
+	tests := []struct {
+		name      string
+		statement in_toto.Statement
+		want      []string
+	}{
+		{
+			name: "standard",
+			want: []string{},
+			statement: in_toto.Statement{
+				Predicate: "hello",
+			},
+		},
+		{
+			name: "subject",
+			want: []string{"sha256:foo"},
+			statement: in_toto.Statement{
+				StatementHeader: in_toto.StatementHeader{
+					Subject: []in_toto.Subject{
+						{
+							Name: "foo",
+							Digest: map[string]string{
+								"sha256": "foo",
+							},
+						},
+					},
+				},
+				Predicate: "hello",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b, err := json.Marshal(tt.statement)
+			if err != nil {
+				t.Fatal(err)
+			}
+			payload := base64.StdEncoding.EncodeToString(b)
+			v := V001Entry{
+				env: ssl.Envelope{
+					Payload:     payload,
+					PayloadType: in_toto.PayloadType,
+				},
+			}
+			sha := sha256.Sum256([]byte(payload))
+			// Always start with the hash
+			want := []string{"sha256:" + hex.EncodeToString(sha[:])}
+			want = append(want, tt.want...)
+			if got := v.IndexKeys(); !reflect.DeepEqual(got, want) {
+				t.Errorf("V001Entry.IndexKeys() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}