Better tests for types package (#1070)

* Add nil check to types.UnmarshalBinary * types nits
google · Mar 26, 2018 · facffc9 · facffc9
1 parent 734f47a
commit facffc9
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 2 deletions.
diff --git a/types/logroot.go b/types/logroot.go
@@ -60,7 +60,10 @@ type LogRoot struct {
 // the LOG_ROOT_FORMAT_V1 tag, and populates the caller with the deserialized
 // *LogRootV1.
 func (l *LogRootV1) UnmarshalBinary(logRootBytes []byte) error {
-	if logRootBytes == nil {
+	if len(logRootBytes) < 3 {
+		return fmt.Errorf("logRootBytes too short")
+	}
+	if l == nil {
 		return fmt.Errorf("nil log root")
 	}
 	version := binary.BigEndian.Uint16(logRootBytes)

diff --git a/types/logroot_test.go b/types/logroot_test.go
@@ -65,6 +65,11 @@ func TestUnmarshalLogRoot(t *testing.T) {
 			logRoot: []byte{0, 1, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5},
 			wantErr: true,
 		},
+		{
+			// Incorrect type.
+			logRoot: []byte{0},
+			wantErr: true,
+		},
 		{logRoot: []byte("foo"), wantErr: true},
 		{logRoot: nil, wantErr: true},
 	} {
@@ -75,6 +80,12 @@ func TestUnmarshalLogRoot(t *testing.T) {
 			t.Errorf("UnmarshalBinary(): %v, wantErr %v", err, want)
 		}
 	}
+
+	// Unmarshaling to a nil should throw an error.
+	var nilPtr *LogRootV1
+	if err := nilPtr.UnmarshalBinary(MustMarshalLogRoot(&LogRootV1{})); err == nil {
+		t.Errorf("nil.UnmarshalBinary(): %v, want err", err)
+	}
 }
 
 func MustMarshalLogRoot(root *LogRootV1) []byte {

diff --git a/types/maproot.go b/types/maproot.go
@@ -54,7 +54,10 @@ type MapRoot struct {
 // UnmarshalBinary verifies that mapRootBytes is a TLS serialized MapRoot,
 // has the MAP_ROOT_FORMAT_V1 tag, and returns the deserialized *MapRootV1.
 func (m *MapRootV1) UnmarshalBinary(mapRootBytes []byte) error {
-	if mapRootBytes == nil {
+	if len(mapRootBytes) < 3 {
+		return fmt.Errorf("mapRootBytes too short")
+	}
+	if m == nil {
 		return fmt.Errorf("nil map root")
 	}
 	version := binary.BigEndian.Uint16(mapRootBytes)

diff --git a/types/maproot_test.go b/types/maproot_test.go
@@ -74,6 +74,11 @@ func TestUnmarshalMapRoot(t *testing.T) {
 			mapRoot: []byte{0, 1, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5},
 			wantErr: true,
 		},
+		{
+			// Incorrect type
+			mapRoot: []byte{0},
+			wantErr: true,
+		},
 		{mapRoot: []byte("foo"), wantErr: true},
 		{mapRoot: nil, wantErr: true},
 	} {
@@ -86,4 +91,9 @@ func TestUnmarshalMapRoot(t *testing.T) {
 			t.Errorf("UnmarshalBinary(): \n%#v, want: \n%#v", got, want)
 		}
 	}
+	// Unmarshaling to a nil should throw an error.
+	var nilPtr *MapRootV1
+	if err := nilPtr.UnmarshalBinary(MustMarshalMapRoot(&MapRootV1{})); err == nil {
+		t.Errorf("nil.UnmarshalBinary(): %v, want err", err)
+	}
 }