refactored and added sth error-checking

2 files changed, 24 insertions, 15 deletions

diff --git a/handler.go b/handler.go
index 8f36e50..0bc5a53 100644
--- a/handler.go
+++ b/handler.go
@@ -9,7 +9,6 @@ import (
 
 	"github.com/golang/glog"
 	"github.com/google/trillian"
-	"github.com/google/trillian/types"
 )
 
 // appHandler implements the http.Handler interface, and contains a reference
@@ -211,16 +210,11 @@ func getSth(ctx context.Context, i *Instance, w http.ResponseWriter, _ *http.Req
 	if err != nil {
 		return http.StatusInternalServerError, fmt.Errorf("failed fetching signed tree head from Trillian backend: %v", err)
 	}
-	if trillianResponse.SignedLogRoot == nil {
-		return http.StatusInternalServerError, fmt.Errorf("Trillian returned no tree head")
-	}
 
-	var lr types.LogRootV1
-	if err := lr.UnmarshalBinary(trillianResponse.SignedLogRoot.GetLogRoot()); err != nil {
-		return http.StatusInternalServerError, fmt.Errorf("failed unmarshaling tree head: %v", err)
+	th, err := NewTreeHeadV1(i.LogParameters, trillianResponse.SignedLogRoot)
+	if err != nil {
+		return http.StatusInternalServerError, fmt.Errorf("failed creating tree head: %v", err)
 	}
-
-	th := NewTreeHeadV1(uint64(lr.TimestampNanos/1000/1000), uint64(lr.TreeSize), lr.RootHash)
 	sth, err := GenV1STH(i.LogParameters, th)
 	if err != nil {
 		return http.StatusInternalServerError, fmt.Errorf("failed creating signed tree head: %v", err)
diff --git a/type.go b/type.go
index 663ae3d..cee52bf 100644
--- a/type.go
+++ b/type.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/google/certificate-transparency-go/tls"
 	"github.com/google/trillian"
+	"github.com/google/trillian/types"
 )
 
 // StFormat defines a particular StItem type that is versioned
@@ -93,15 +94,29 @@ func NewSignedTreeHeadV1(th TreeHeadV1, logId, signature []byte) StItem {
 	}
 }
 
-func NewTreeHeadV1(timestamp, treeSize uint64, rootHash []byte) TreeHeadV1 {
+// NewTreeHead converts a Trillian-signed log root to a tree head without
+// verifying any signature.  In other words, Trillian <-> STFE is trusted.
+func NewTreeHeadV1(lp *LogParameters, slr *trillian.SignedLogRoot) (TreeHeadV1, error) {
+	if slr == nil {
+		return TreeHeadV1{}, fmt.Errorf("Trillian returned no tree head")
+	}
+
+	var lr types.LogRootV1
+	if err := lr.UnmarshalBinary(slr.GetLogRoot()); err != nil {
+		return TreeHeadV1{}, fmt.Errorf("failed unmarshaling Trillian slr: %v", err)
+	}
+	if lp.HashType.Size() != len(lr.RootHash) {
+		return TreeHeadV1{}, fmt.Errorf("invalid Trillian root hash: %v", lr.RootHash)
+	}
+
 	return TreeHeadV1{
-		Timestamp: timestamp,
-		TreeSize:  treeSize,
+		Timestamp: uint64(lr.TimestampNanos / 1000 / 1000),
+		TreeSize:  uint64(lr.TreeSize),
 		RootHash: NodeHash{
-			Data: rootHash,
+			Data: lr.RootHash,
 		},
-		Extension: nil,
-	}
+		Extension: nil, // no known extensions
+	}, nil
 }
 
 func NewSignedDebugInfoV1(logId, message, signature []byte) StItem {