pw_software_update: Parse bundles

Bug: 456
Change-Id: I21e6f8a4e69a037e4c028fc59f9049be1c3b8cdb
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/61341
Reviewed-by: David Rogers <davidrogers@google.com>
Reviewed-by: Ewout van Bekkum <ewout@google.com>
Commit-Queue: Yecheng Zhao <zyecheng@google.com>
diff --git a/pw_protobuf/public/pw_protobuf/message.h b/pw_protobuf/public/pw_protobuf/message.h
index 7ad73e9..86e5fab 100644
--- a/pw_protobuf/public/pw_protobuf/message.h
+++ b/pw_protobuf/public/pw_protobuf/message.h
@@ -104,6 +104,8 @@
   Status status() { return reader_.status(); }
 
   // Check whether the bytes value equals the given `bytes`.
+  // TODO(pwbug/456): Should this return `bool`? In the case of error, is it ok
+  // to just return false?
   Result<bool> Equal(ConstByteSpan bytes);
 
  private:
@@ -229,9 +231,9 @@
     // template specialization.
     template <typename FieldType>
     FieldType As() {
-      protobuf::StreamDecoder decoder(field_reader_.Reset());
+      StreamDecoder decoder(field_reader_.Reset());
       PW_TRY(decoder.Next());
-      Result<protobuf::StreamDecoder::Bounds> payload_bounds =
+      Result<StreamDecoder::Bounds> payload_bounds =
           decoder.GetLengthDelimitedPayloadBounds();
       PW_TRY(payload_bounds.status());
       // The bounds is relative to the given stream::IntervalReader. Convert
diff --git a/pw_software_update/BUILD.bazel b/pw_software_update/BUILD.bazel
index 861302e..5ae1d87 100644
--- a/pw_software_update/BUILD.bazel
+++ b/pw_software_update/BUILD.bazel
@@ -34,6 +34,7 @@
 
 pw_cc_library(
     name = "update_bundle",
+    srcs = ["update_bundle.cc"],
     hdrs = [
         "public/pw_software_update/config.h",
         "public/pw_software_update/update_bundle.h",
diff --git a/pw_software_update/BUILD.gn b/pw_software_update/BUILD.gn
index 6f431d7..7a8a420 100644
--- a/pw_software_update/BUILD.gn
+++ b/pw_software_update/BUILD.gn
@@ -52,10 +52,6 @@
   }
 }
 
-pw_test_group("tests") {
-  tests = [ ":update_bundle_test" ]
-}
-
 pw_doc_group("docs") {
   sources = [ "docs.rst" ]
 }
@@ -83,6 +79,7 @@
   pw_source_set("update_bundle") {
     public_configs = [ ":public_include_path" ]
     public_deps = [
+      ":protos.pwpb",
       ":update_backend",
       dir_pw_blob_store,
       dir_pw_kvs,
@@ -92,6 +89,7 @@
       "public/pw_software_update/config.h",
       "public/pw_software_update/update_bundle.h",
     ]
+    sources = [ "update_bundle.cc" ]
   }
 } else {
   group("rpc_service") {
@@ -100,14 +98,36 @@
   }
 }
 
+pw_python_action("generate_test_bundle") {
+  header_output = "$target_gen_dir/$target_name/test_bundles.h"
+  script = "py/pw_software_update/generate_test_bundle.py"
+  python_deps = [
+    ":protos.python",
+    "py",
+  ]
+  outputs = [ header_output ]
+  args = [ rebase_path(header_output) ]
+}
+
+config("generated_test_bundle_include") {
+  _generated_outputs = get_target_outputs(":generate_test_bundle")
+  include_dirs = [ get_path_info(_generated_outputs[0], "dir") ]
+}
+
 pw_test("update_bundle_test") {
   enable_if =
       dir_pw_third_party_nanopb != "" && dir_pw_third_party_protobuf != ""
   sources = [ "update_bundle_test.cc" ]
   public_deps = [
+    ":generate_test_bundle",
     ":rpc_service",
     ":update_bundle",
     "$dir_pw_kvs:fake_flash",
     "$dir_pw_kvs:fake_flash_test_key_value_store",
   ]
+  configs = [ ":generated_test_bundle_include" ]
+}
+
+pw_test_group("tests") {
+  tests = [ ":update_bundle_test" ]
 }
diff --git a/pw_software_update/public/pw_software_update/config.h b/pw_software_update/public/pw_software_update/config.h
index 25a27ab..8de358d 100644
--- a/pw_software_update/public/pw_software_update/config.h
+++ b/pw_software_update/public/pw_software_update/config.h
@@ -14,6 +14,11 @@
 #pragma once
 
 // The log level to use for this module. Logs below this level are omitted.
+#define PW_LOG_MODULE_NAME "PWSU"
 #ifndef PW_SOFTWARE_UPDATE_CONFIG_LOG_LEVEL
 #define PW_SOFTWARE_UPDATE_CONFIG_LOG_LEVEL PW_LOG_LEVEL_DEBUG
+
+// The size of the buffer to create on stack for streaming manifest data from
+// the bundle reader.
+#define WRITE_MANIFEST_STREAM_PIPE_BUFFER_SIZE 8
 #endif  // PW_SOFTWARE_UPDATE_CONFIG_LOG_LEVEL
diff --git a/pw_software_update/public/pw_software_update/update_bundle.h b/pw_software_update/public/pw_software_update/update_bundle.h
index 5f56a21..17fff95 100644
--- a/pw_software_update/public/pw_software_update/update_bundle.h
+++ b/pw_software_update/public/pw_software_update/update_bundle.h
@@ -17,14 +17,16 @@
 #include <cstddef>
 
 #include "pw_blob_store/blob_store.h"
+#include "pw_protobuf/map_utils.h"
+#include "pw_protobuf/message.h"
 #include "pw_software_update/update_backend.h"
+#include "pw_stream/memory_stream.h"
 
 namespace pw::software_update {
 
-// TODO(pwbug/456): Place-holder declaration for now. To be imlemented
+// TODO(pwbug/456): Place-holder declaration for now. To be implemented
 // and moved elsewhere.
-class ElementPayloadReader {};
-class Manifest;
+class Manifest {};
 
 // UpdateBundle is responsible for parsing, verifying and providing
 // target payload access of a software update bundle. It takes the following as
@@ -74,10 +76,10 @@
  public:
   // UpdateBundle
   // update_bundle - The software update bundle data on storage.
-  // helper - project-specific BundledUpdateHelper
+  // backend - project-specific BundledUpdateBackend
   UpdateBundle(blob_store::BlobStore& update_bundle,
                BundledUpdateBackend& backend)
-      : bundle_(update_bundle), backend_(backend) {}
+      : bundle_(update_bundle), backend_(backend), bundle_reader_(bundle_) {}
 
   // Opens and verifies the software update bundle, using the TUF process.
   //
@@ -116,11 +118,13 @@
   // Returns:
   // A reader instance for the target file.
   // TODO(pwbug/456): Figure out a way to propagate error.
-  ElementPayloadReader GetTargetPayload(std::string_view target_file);
+  stream::IntervalReader GetTargetPayload(std::string_view target_file);
 
  private:
   blob_store::BlobStore& bundle_;
   BundledUpdateBackend& backend_;
+  blob_store::BlobStore::BlobReader bundle_reader_;
+  protobuf::Message decoder_;
 };
 
 }  // namespace pw::software_update
\ No newline at end of file
diff --git a/pw_software_update/py/BUILD.gn b/pw_software_update/py/BUILD.gn
index bc76091..242c4fb 100644
--- a/pw_software_update/py/BUILD.gn
+++ b/pw_software_update/py/BUILD.gn
@@ -22,11 +22,10 @@
       name = "pw_software_update"
       version = "0.0.1"
     }
-    options = {
-    }
   }
   sources = [
     "pw_software_update/__init__.py",
+    "pw_software_update/generate_test_bundle.py",
     "pw_software_update/metadata.py",
     "pw_software_update/update_bundle.py",
   ]
diff --git a/pw_software_update/py/pw_software_update/generate_test_bundle.py b/pw_software_update/py/pw_software_update/generate_test_bundle.py
new file mode 100644
index 0000000..2f1a263
--- /dev/null
+++ b/pw_software_update/py/pw_software_update/generate_test_bundle.py
@@ -0,0 +1,144 @@
+# Copyright 2021 The Pigweed Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not
+# use this file except in compliance with the License. You may obtain a copy of
+# the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.
+"""Script for generating test bundles"""
+
+import argparse
+import subprocess
+import sys
+from typing import Dict
+
+from pw_software_update import tuf_pb2, metadata, update_bundle_pb2
+
+HEADER = """// Copyright 2021 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy
+// of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+
+#pragma once
+
+#include "pw_bytes/span.h"
+
+"""
+
+
+def byte_array_declaration(data: bytes, name: str) -> str:
+    """Generates a byte C array declaration for a byte array"""
+    type_name = '[[maybe_unused]] const std::byte'
+    byte_str = ''.join([f'std::byte{{0x{b:02x}}},' for b in data])
+    array_body = f'{{{byte_str}}}'
+    return f'{type_name} {name}[] = {array_body};'
+
+
+def proto_array_declaration(proto, name: str) -> str:
+    """Generates a byte array declaration for a proto"""
+    return byte_array_declaration(proto.SerializeToString(), name)
+
+
+def sign_metadata(role: str, metadata_proto):
+    """Signs a serialized metadata for a given role"""
+    signed = getattr(tuf_pb2, f'Signed{role.capitalize()}Metadata')()
+    serialized = metadata_proto.SerializeToString()
+    setattr(signed, f'serialized_{role}_metadata', serialized)
+    # TODO(pwbug:456): Generate signatures.
+
+    return signed
+
+
+class Bundle:
+    """A helper for test UpdateBundle generation"""
+    def __init__(self):
+        self._payloads: Dict[str, bytes] = {}
+
+    def add_payload(self, name: str, payload: bytes) -> None:
+        """Adds a payload to the bundle"""
+        self._payloads[name] = payload
+
+    def generate_targets_metadata(self) -> tuf_pb2.TargetsMetadata:
+        """Generates the targets metadata"""
+        targets = metadata.gen_targets_metadata(self._payloads)
+        return targets
+
+    def generate_signed_metadta(self, role: str):
+        """Generate a signed metadata of a given role"""
+        unsigned = getattr(self, f'generate_{role}_metadata')()
+        return sign_metadata(role, unsigned)
+
+    def generate_bundle(self) -> update_bundle_pb2.UpdateBundle:
+        """Generate the update bundle"""
+        bundle = update_bundle_pb2.UpdateBundle()
+        bundle.targets_metadata['targets'].CopyFrom(
+            self.generate_signed_metadta('targets'))
+        for name, payload in self._payloads.items():
+            bundle.target_payloads[name] = payload
+
+        return bundle
+
+    def generate_manifest(self) -> update_bundle_pb2.Manifest:
+        """Generates the manifest"""
+        manifest = update_bundle_pb2.Manifest()
+        manifest.targets_metadata['targets'].CopyFrom(
+            self.generate_targets_metadata())
+        return manifest
+
+
+def parse_args():
+    """Setup argparse."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("output_header",
+                        help="output path of the generated C header")
+    return parser.parse_args()
+
+
+def main() -> int:
+    """Main"""
+    args = parse_args()
+
+    test_bundle = Bundle()
+
+    # Add some update files.
+    test_bundle.add_payload('file1', 'file 1 content'.encode())  # type: ignore
+    test_bundle.add_payload('file2', 'file 2 content'.encode())  # type: ignore
+
+    update_bundle_proto = test_bundle.generate_bundle()
+    manifest_proto = test_bundle.generate_manifest()
+
+    with open(args.output_header, 'w') as header:
+        header.write(HEADER)
+        header.write('namespace {')
+        header.write(
+            proto_array_declaration(update_bundle_proto, 'kTestBundle'))
+        header.write(
+            proto_array_declaration(manifest_proto, 'kTestBundleManifest'))
+        header.write('}')
+
+    subprocess.run([
+        'clang-format',
+        '-i',
+        args.output_header,
+    ], check=True)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/pw_software_update/update_bundle.cc b/pw_software_update/update_bundle.cc
new file mode 100644
index 0000000..fd9e894
--- /dev/null
+++ b/pw_software_update/update_bundle.cc
@@ -0,0 +1,137 @@
+// Copyright 2021 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+
+#include "pw_software_update/config.h"
+
+#define PW_LOG_LEVEL PW_SOFTWARE_UPDATE_CONFIG_LOG_LEVEL
+
+#include <cstddef>
+#include <string_view>
+
+#include "pw_log/log.h"
+#include "pw_protobuf/message.h"
+#include "pw_result/result.h"
+#include "pw_software_update/update_bundle.h"
+#include "pw_software_update/update_bundle.pwpb.h"
+#include "pw_stream/interval_reader.h"
+#include "pw_stream/memory_stream.h"
+
+namespace pw::software_update {
+namespace {
+
+constexpr std::string_view kTopLevelTargetsName = "targets";
+
+}
+
+Status UpdateBundle::OpenAndVerify(const Manifest&) {
+  PW_TRY(bundle_.Init());
+  PW_TRY(bundle_reader_.Open());
+  decoder_ =
+      protobuf::Message(bundle_reader_, bundle_reader_.ConservativeReadLimit());
+  (void)backend_;
+  // TODO(pw_bug/456): Implement verification logic.
+  return OkStatus();
+}
+
+// Get the target element corresponding to `target_file`
+stream::IntervalReader UpdateBundle::GetTargetPayload(
+    std::string_view target_file_name) {
+  protobuf::StringToBytesMap target_payloads =
+      decoder_.AsStringToBytesMap(static_cast<uint32_t>(
+          pw_software_update::UpdateBundle::Fields::TARGET_PAYLOADS));
+  PW_TRY(target_payloads.status());
+  protobuf::Bytes payload = target_payloads[target_file_name];
+  PW_TRY(payload.status());
+  return payload.GetBytesReader();
+}
+
+Result<bool> UpdateBundle::IsTargetPayloadIncluded(
+    std::string_view target_file_name) {
+  // TODO(pwbug/456): Perform personalization check first. If the target
+  // is personalized out. Don't need to proceed.
+
+  protobuf::StringToMessageMap signed_targets_metadata_map =
+      decoder_.AsStringToMessageMap(static_cast<uint32_t>(
+          pw_software_update::UpdateBundle::Fields::TARGETS_METADATA));
+  PW_TRY(signed_targets_metadata_map.status());
+
+  // There should only be one element in the map, which is the top-level
+  // targets metadata.
+  protobuf::Message signed_targets_metadata =
+      signed_targets_metadata_map[kTopLevelTargetsName];
+  PW_TRY(signed_targets_metadata.status());
+
+  protobuf::Message metadata = signed_targets_metadata.AsMessage(
+      static_cast<uint32_t>(pw_software_update::SignedTargetsMetadata::Fields::
+                                SERIALIZED_TARGETS_METADATA));
+  PW_TRY(metadata.status());
+
+  protobuf::RepeatedMessages target_files =
+      metadata.AsRepeatedMessages(static_cast<uint32_t>(
+          pw_software_update::TargetsMetadata::Fields::TARGET_FILES));
+  PW_TRY(target_files.status());
+
+  for (protobuf::Message target_file : target_files) {
+    protobuf::String name = target_file.AsString(static_cast<uint32_t>(
+        pw_software_update::TargetFile::Fields::FILE_NAME));
+    PW_TRY(name.status());
+    Result<bool> file_name_matches = name.Equal(target_file_name);
+    PW_TRY(file_name_matches.status());
+    if (file_name_matches.value()) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+Status UpdateBundle::WriteManifest(stream::Writer& staged_manifest_writer) {
+  protobuf::StringToMessageMap signed_targets_metadata_map =
+      decoder_.AsStringToMessageMap(static_cast<uint32_t>(
+          pw_software_update::UpdateBundle::Fields::TARGETS_METADATA));
+  PW_TRY(signed_targets_metadata_map.status());
+
+  // There should only be one element in the map, which is the top-level
+  // targets metadata.
+  protobuf::Message signed_targets_metadata =
+      signed_targets_metadata_map[kTopLevelTargetsName];
+  PW_TRY(signed_targets_metadata.status());
+
+  protobuf::Bytes metadata = signed_targets_metadata.AsBytes(
+      static_cast<uint32_t>(pw_software_update::SignedTargetsMetadata::Fields::
+                                SERIALIZED_TARGETS_METADATA));
+  PW_TRY(metadata.status());
+
+  stream::MemoryReader name_reader(
+      std::as_bytes(std::span(kTopLevelTargetsName)));
+  stream::IntervalReader metadata_reader = metadata.GetBytesReader();
+
+  std::byte stream_pipe_buffer[WRITE_MANIFEST_STREAM_PIPE_BUFFER_SIZE];
+  return protobuf::WriteProtoStringToBytesMapEntry(
+      static_cast<uint32_t>(
+          pw_software_update::Manifest::Fields::TARGETS_METADATA),
+      name_reader,
+      kTopLevelTargetsName.size(),
+      metadata_reader,
+      metadata_reader.interval_size(),
+      stream_pipe_buffer,
+      staged_manifest_writer);
+}
+
+Status UpdateBundle::Close() {
+  // TODO(pwbug/456): To be implemented.
+  return bundle_reader_.Close();
+}
+
+}  // namespace pw::software_update
diff --git a/pw_software_update/update_bundle_test.cc b/pw_software_update/update_bundle_test.cc
index 4522cac..e454828 100644
--- a/pw_software_update/update_bundle_test.cc
+++ b/pw_software_update/update_bundle_test.cc
@@ -18,6 +18,10 @@
 #include "pw_kvs/fake_flash_memory.h"
 #include "pw_kvs/test_key_value_store.h"
 #include "pw_software_update/update_backend.h"
+#include "pw_software_update/update_bundle.pwpb.h"
+#include "test_bundles.h"
+
+#define ASSERT_OK(status) ASSERT_EQ(OkStatus(), status)
 
 namespace pw::software_update {
 namespace {
@@ -44,6 +48,14 @@
 
   BundledUpdateBackend& backend() { return backend_; }
 
+  void StageTestBundle(ConstByteSpan bundle_data) {
+    ASSERT_OK(bundle_blob_.Init());
+    blob_store::BlobStore::BlobWriter blob_writer(bundle_blob());
+    ASSERT_OK(blob_writer.Open());
+    ASSERT_OK(blob_writer.Write(bundle_data));
+    ASSERT_OK(blob_writer.Close());
+  }
+
  private:
   kvs::FakeFlashMemoryBuffer<kSectorSize, kSectorCount> blob_flash_;
   kvs::FlashPartition blob_partition_;
@@ -53,6 +65,73 @@
 
 }  // namespace
 
-TEST_F(UpdateBundleTest, Create) { UpdateBundle(bundle_blob(), backend()); }
+TEST_F(UpdateBundleTest, GetTargetPayload) {
+  StageTestBundle(kTestBundle);
+  UpdateBundle update_bundle(bundle_blob(), backend());
+
+  Manifest current_manifest;
+  ASSERT_OK(update_bundle.OpenAndVerify(current_manifest));
+
+  {
+    stream::IntervalReader res = update_bundle.GetTargetPayload("file1");
+    ASSERT_OK(res.status());
+
+    const char kExpectedContent[] = "file 1 content";
+    char read_buffer[sizeof(kExpectedContent) + 1] = {0};
+    ASSERT_TRUE(res.Read(read_buffer, sizeof(kExpectedContent)).ok());
+    ASSERT_STREQ(read_buffer, kExpectedContent);
+  }
+
+  {
+    stream::IntervalReader res = update_bundle.GetTargetPayload("file2");
+    ASSERT_OK(res.status());
+
+    const char kExpectedContent[] = "file 2 content";
+    char read_buffer[sizeof(kExpectedContent) + 1] = {0};
+    ASSERT_TRUE(res.Read(read_buffer, sizeof(kExpectedContent)).ok());
+    ASSERT_STREQ(read_buffer, kExpectedContent);
+  }
+
+  {
+    stream::IntervalReader res = update_bundle.GetTargetPayload("non-exist");
+    ASSERT_EQ(res.status(), Status::NotFound());
+  }
+}
+
+TEST_F(UpdateBundleTest, IsTargetPayloadIncluded) {
+  StageTestBundle(kTestBundle);
+  UpdateBundle update_bundle(bundle_blob(), backend());
+
+  Manifest current_manifest;
+  ASSERT_OK(update_bundle.OpenAndVerify(current_manifest));
+
+  Result<bool> res = update_bundle.IsTargetPayloadIncluded("file1");
+  ASSERT_OK(res.status());
+  ASSERT_TRUE(res.value());
+
+  res = update_bundle.IsTargetPayloadIncluded("file2");
+  ASSERT_OK(res.status());
+  ASSERT_TRUE(res.value());
+
+  res = update_bundle.IsTargetPayloadIncluded("non-exist");
+  ASSERT_OK(res.status());
+  ASSERT_FALSE(res.value());
+}
+
+TEST_F(UpdateBundleTest, WriteManifest) {
+  StageTestBundle(kTestBundle);
+  UpdateBundle update_bundle(bundle_blob(), backend());
+
+  Manifest current_manifest;
+  ASSERT_OK(update_bundle.OpenAndVerify(current_manifest));
+
+  std::byte manifest_buffer[sizeof(kTestBundleManifest)];
+  stream::MemoryWriter manifest_writer(manifest_buffer);
+  ASSERT_OK(update_bundle.WriteManifest(manifest_writer));
+
+  ASSERT_EQ(
+      memcmp(manifest_buffer, kTestBundleManifest, sizeof(kTestBundleManifest)),
+      0);
+}
 
 }  // namespace pw::software_update