Improve error handling of the transport security state generator.

net/tools/transport_security_state_generator/preloaded_state_generator.cc

Index: net/tools/transport_security_state_generator/preloaded_state_generator.cc
diff --git a/net/tools/transport_security_state_generator/preloaded_state_generator.cc b/net/tools/transport_security_state_generator/preloaded_state_generator.cc
index 64e89890ac76506ff85f910a8bdf7febca645b74..2d1102dd069b952b0cee0948b2457e9409105efb 100644
--- a/net/tools/transport_security_state_generator/preloaded_state_generator.cc
+++ b/net/tools/transport_security_state_generator/preloaded_state_generator.cc
@@ -4,8 +4,6 @@
 
 #include "net/tools/transport_security_state_generator/preloaded_state_generator.h"
 
-#include <iostream>
-
 #include <string>
 
 #include "base/strings/string_util.h"
@@ -129,27 +127,28 @@ PreloadedStateGenerator::PreloadedStateGenerator() {}
 
 PreloadedStateGenerator::~PreloadedStateGenerator() {}
 
-std::string PreloadedStateGenerator::Generate(
+bool PreloadedStateGenerator::Generate(
     const std::string& preload_template,
     const TransportSecurityStateEntries& entries,
     const DomainIDList& domain_ids,
     const Pinsets& pinsets,
-    bool verbose) {
-  std::string output = preload_template;
+    std::string* output) {
+  DCHECK(output);
+  std::string generated = preload_template;
 
   NameIDMap domain_ids_map;
-  ProcessDomainIds(domain_ids, &domain_ids_map, &output);
+  ProcessDomainIds(domain_ids, &domain_ids_map, &generated);
 
-  ProcessSPKIHashes(pinsets, &output);
+  ProcessSPKIHashes(pinsets, &generated);
 
   NameIDMap expect_ct_report_uri_map;
-  ProcessExpectCTURIs(entries, &expect_ct_report_uri_map, &output);
+  ProcessExpectCTURIs(entries, &expect_ct_report_uri_map, &generated);
 
   NameIDMap expect_staple_report_uri_map;
-  ProcessExpectStapleURIs(entries, &expect_staple_report_uri_map, &output);
+  ProcessExpectStapleURIs(entries, &expect_staple_report_uri_map, &generated);
 
   NameIDMap pinsets_map;
-  ProcessPinsets(pinsets, &pinsets_map, &output);
+  ProcessPinsets(pinsets, &pinsets_map, &generated);
 
   // The trie generation process is ran twice, the first time using an
   // approximate Huffman table. During this first run, the correct character
@@ -161,34 +160,31 @@ std::string PreloadedStateGenerator::Generate(
   TrieWriter writer(table, domain_ids_map, expect_ct_report_uri_map,
                     expect_staple_report_uri_map, pinsets_map,
                     &huffman_builder);
-  writer.WriteEntries(entries);
-  uint32_t initial_length = writer.position();
+  uint32_t root_position;
+  if (!writer.WriteEntries(entries, &root_position)) {
+    return false;
+  }
 
   HuffmanRepresentationTable optimal_table = huffman_builder.ToTable();
   TrieWriter new_writer(optimal_table, domain_ids_map, expect_ct_report_uri_map,
                         expect_staple_report_uri_map, pinsets_map, nullptr);
 
-  uint32_t root_position = new_writer.WriteEntries(entries);
-  uint32_t new_length = new_writer.position();
+  if (!new_writer.WriteEntries(entries, &root_position)) {
+    return false;
+  }
 
+  uint32_t new_length = new_writer.position();
   std::vector<uint8_t> huffman_tree = huffman_builder.ToVector();
-
   new_writer.Flush();
 
-  ReplaceTag("HUFFMAN_TREE", FormatVectorAsArray(huffman_tree), &output);
-  ReplaceTag("HSTS_TRIE", FormatVectorAsArray(new_writer.bytes()), &output);
+  ReplaceTag("HUFFMAN_TREE", FormatVectorAsArray(huffman_tree), &generated);
+  ReplaceTag("HSTS_TRIE", FormatVectorAsArray(new_writer.bytes()), &generated);
 
-  ReplaceTag("HSTS_TRIE_BITS", std::to_string(new_length), &output);
-  ReplaceTag("HSTS_TRIE_ROOT", std::to_string(root_position), &output);
+  ReplaceTag("HSTS_TRIE_BITS", std::to_string(new_length), &generated);
+  ReplaceTag("HSTS_TRIE_ROOT", std::to_string(root_position), &generated);
 
-  if (verbose) {
-    std::cout << "Saved " << std::to_string(initial_length - new_length)
-              << " bits by using accurate Huffman counts." << std::endl;
-    std::cout << "Bit length " << std::to_string(new_length) << std::endl;
-    std::cout << "Root position " << std::to_string(root_position) << std::endl;
-  }
-
-  return output;
+  *output = generated;
+  return true;
 }
 
 void PreloadedStateGenerator::ProcessDomainIds(const DomainIDList& domain_ids,
@@ -275,6 +271,11 @@ void PreloadedStateGenerator::ProcessExpectCTURIs(
     }
   }
 
+  output.append(kIndent);
+  output.append(kIndent);
+  output.append("NULL,");
+  output.append(kNewLine);
+
   output.append("}");
   ReplaceTag("EXPECT_CT_REPORT_URIS", output, tpl);
 }
@@ -302,6 +303,11 @@ void PreloadedStateGenerator::ProcessExpectStapleURIs(
     }
   }
 
+  output.append(kIndent);
+  output.append(kIndent);
+  output.append("NULL,");
+  output.append(kNewLine);
+
   output.append("}");
   ReplaceTag("EXPECT_STAPLE_REPORT_URIS", output, tpl);
 }