github.com/google/trillian-examples@v0.0.0-20240520080811-0d40d35cef0e/binary_transparency/firmware/internal/ftmap/pipeline_test.go (about)

     1  // Copyright 2021 Google LLC. All Rights Reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package ftmap
    16  
    17  import (
    18  	"crypto/sha512"
    19  	"encoding/json"
    20  	"fmt"
    21  	"testing"
    22  
    23  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    24  	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
    25  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
    26  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
    27  	"github.com/google/trillian-examples/binary_transparency/firmware/api"
    28  	"github.com/google/trillian/experimental/batchmap"
    29  )
    30  
    31  func init() {
    32  	register.Function1x1(testLogToStringFn)
    33  	register.Function1x1(testRootToStringFn)
    34  }
    35  
    36  func TestCreate(t *testing.T) {
    37  	tests := []struct {
    38  		name   string
    39  		treeID int64
    40  		count  int64
    41  
    42  		wantRoot string
    43  		wantLogs []string
    44  	}{
    45  		{
    46  			name:   "One entry",
    47  			treeID: 12345,
    48  			count:  1,
    49  
    50  			wantRoot: "865bbd0034750dd6abe512470722fad51c85ba95d96316818e8ec8a2a55679c0",
    51  			wantLogs: []string{"dummy: [1]"},
    52  		},
    53  		{
    54  			name:   "All entries",
    55  			treeID: 12345,
    56  			count:  -1,
    57  
    58  			wantRoot: "daa0e0c66d69162abbe27ba9aa54a8bdb8850f1100e0626e45ea477cea4765e6",
    59  			wantLogs: []string{"dummy: [1 5 3]", "fish: [42]"},
    60  		},
    61  	}
    62  
    63  	inputLog := fakeLog{
    64  		leaves: []api.SignedStatement{
    65  			createFWSignedStatement("dummy", 1),
    66  			createFWSignedStatement("dummy", 5),
    67  			createFWSignedStatement("fish", 42),
    68  			createFWSignedStatement("dummy", 3),
    69  		},
    70  		head: []byte("this is just passed around"),
    71  	}
    72  	for _, test := range tests {
    73  		test := test
    74  		t.Run(test.name, func(t *testing.T) {
    75  			mb := NewMapBuilder(inputLog, test.treeID, 0)
    76  			p, s := beam.NewPipelineWithRoot()
    77  
    78  			result, err := mb.Create(s, test.count)
    79  			if err != nil {
    80  				t.Fatalf("failed to Create(): %v", err)
    81  			}
    82  
    83  			passert.Equals(s, beam.ParDo(s, testRootToStringFn, result.MapTiles), test.wantRoot)
    84  
    85  			passert.Equals(s, beam.ParDo(s, testLogToStringFn, result.DeviceLogs), beam.CreateList(s, test.wantLogs))
    86  
    87  			err = ptest.Run(p)
    88  			if err != nil {
    89  				t.Errorf("unexpected error: %v", err)
    90  			}
    91  		})
    92  	}
    93  }
    94  func testRootToStringFn(t *batchmap.Tile) string { return fmt.Sprintf("%x", t.RootHash) }
    95  func testLogToStringFn(l *api.DeviceReleaseLog) string {
    96  	return fmt.Sprintf("%s: %v", l.DeviceID, l.Revisions)
    97  }
    98  
    99  func createFW(device string, revision uint64) api.FirmwareMetadata {
   100  	image := fmt.Sprintf("this image is the firmware at revision %d for device %s.", revision, device)
   101  	imageHash := sha512.Sum512([]byte(image))
   102  	return api.FirmwareMetadata{
   103  		DeviceID:            device,
   104  		FirmwareRevision:    revision,
   105  		FirmwareImageSHA512: imageHash[:],
   106  	}
   107  }
   108  
   109  func createFWSignedStatement(device string, revision uint64) api.SignedStatement {
   110  	fwbs, _ := json.Marshal(createFW(device, revision))
   111  	return api.SignedStatement{
   112  		Type:      api.FirmwareMetadataType,
   113  		Statement: fwbs,
   114  	}
   115  }
   116  
   117  type fakeLog struct {
   118  	leaves []api.SignedStatement
   119  	head   []byte
   120  }
   121  
   122  func (l fakeLog) Head() ([]byte, int64, error) {
   123  	return l.head, int64(len(l.leaves)), nil
   124  }
   125  
   126  func (l fakeLog) Entries(s beam.Scope, start, end int64) beam.PCollection {
   127  	count := int(end - start)
   128  	entries := make([]InputLogLeaf, count)
   129  	for i := 0; i < count; i++ {
   130  		// This swallows the error, but the test will fail anyway. YOLO.
   131  		index := start + int64(i)
   132  		bs, _ := json.Marshal(l.leaves[int(index)])
   133  		entries[i] = InputLogLeaf{
   134  			Seq:  index,
   135  			Data: bs,
   136  		}
   137  	}
   138  	return beam.CreateList(s, entries)
   139  }