github.com/google/trillian-examples@v0.0.0-20240520080811-0d40d35cef0e/binary_transparency/firmware/internal/ftmap/aggregate.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  	"reflect"
    22  
    23  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    24  
    25  	"github.com/google/trillian-examples/binary_transparency/firmware/api"
    26  	"github.com/google/trillian/experimental/batchmap"
    27  	"github.com/google/trillian/merkle/coniks"
    28  	"github.com/google/trillian/merkle/smt/node"
    29  )
    30  
    31  func init() {
    32  	beam.RegisterFunction(aggregationFn)
    33  	beam.RegisterFunction(annotationLogIndexFn)
    34  	beam.RegisterFunction(logEntryIndexFn)
    35  	beam.RegisterType(reflect.TypeOf((*api.AggregatedFirmware)(nil)).Elem())
    36  	beam.RegisterType(reflect.TypeOf((*aggregatedFirmwareHashFn)(nil)).Elem())
    37  }
    38  
    39  // Aggregate will output an entry for each firmware entry in the input log, which includes the
    40  // firmware metadata, along with an aggregated representation of the annotations for it. The
    41  // rules are:
    42  //   - AnnotationMalware: `Good` is true providing there are no malware annotations that claim the
    43  //     firmware is bad.
    44  func Aggregate(s beam.Scope, treeID int64, fws, annotationMalwares beam.PCollection) (beam.PCollection, beam.PCollection) {
    45  	keyedFws := beam.ParDo(s, logEntryIndexFn, fws)
    46  	keyedAnns := beam.ParDo(s, annotationLogIndexFn, annotationMalwares)
    47  	annotations := beam.ParDo(s, aggregationFn, beam.CoGroupByKey(s, keyedFws, keyedAnns))
    48  	return beam.ParDo(s, &aggregatedFirmwareHashFn{treeID}, annotations), annotations
    49  }
    50  
    51  func logEntryIndexFn(l *firmwareLogEntry) (uint64, *firmwareLogEntry) { return uint64(l.Index), l }
    52  
    53  func annotationLogIndexFn(a *annotationMalwareLogEntry) (uint64, *annotationMalwareLogEntry) {
    54  	return a.Annotation.FirmwareID.LogIndex, a
    55  }
    56  
    57  func aggregationFn(fwIndex uint64, fwit func(**firmwareLogEntry) bool, amit func(**annotationMalwareLogEntry) bool) (*api.AggregatedFirmware, error) {
    58  	// There will be exactly one firmware entry for the log index.
    59  	var fwle *firmwareLogEntry
    60  	if !fwit(&fwle) {
    61  		return nil, fmt.Errorf("aggregationFn for %d found no firmware", fwIndex)
    62  	}
    63  
    64  	// And 0-many annotations. The FW is good as long as no annotations say that it is not.
    65  	good := true
    66  	var amle *annotationMalwareLogEntry
    67  	for amit(&amle) {
    68  		good = good && amle.Annotation.Good
    69  	}
    70  
    71  	return &api.AggregatedFirmware{
    72  		Index: fwIndex,
    73  		Good:  good,
    74  	}, nil
    75  }
    76  
    77  type aggregatedFirmwareHashFn struct {
    78  	TreeID int64
    79  }
    80  
    81  func (fn *aggregatedFirmwareHashFn) ProcessElement(afw *api.AggregatedFirmware) (*batchmap.Entry, error) {
    82  	// We use the Index of the FW Log Entry as the key because that's unique and what we matched on.
    83  	// The client already has this information from the inclusion proof so it's known information.
    84  	// Using `afw.Firmware.FirmwareImageSHA512` is another logical choice, but there's nothing to
    85  	// prevent multiple entries in the log with the same FW hash, and thus we'd have key conflicts.
    86  	// If we did want to use image hash as the key then we'd need to restructure FW annotations to
    87  	// remove the `LogIndex` from `FirmwareId`.
    88  	key := []byte(fmt.Sprintf("summary:%d", afw.Index))
    89  	value, err := json.Marshal(afw)
    90  	if err != nil {
    91  		return nil, fmt.Errorf("marshaling AggregatedFirmware failed: %v", err)
    92  	}
    93  
    94  	kbs := sha512.Sum512_256(key)
    95  	leafID := node.NewID(string(kbs[:]), 256)
    96  	return &batchmap.Entry{
    97  		HashKey:   kbs[:],
    98  		HashValue: coniks.Default.HashLeaf(fn.TreeID, leafID, value),
    99  	}, nil
   100  }