Add static domain security state generator tool.

net/tools/domain_security_preload_generator/preloaded_state_generator.cc

Index: net/tools/domain_security_preload_generator/preloaded_state_generator.cc
diff --git a/net/tools/domain_security_preload_generator/preloaded_state_generator.cc b/net/tools/domain_security_preload_generator/preloaded_state_generator.cc
new file mode 100644
index 0000000000000000000000000000000000000000..f85b982530970dc898bb028df25df2727bb590a0
--- /dev/null
+++ b/net/tools/domain_security_preload_generator/preloaded_state_generator.cc
@@ -0,0 +1,337 @@
+// Copyright (c) 2016 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "net/tools/domain_security_preload_generator/preloaded_state_generator.h"
+
+#include <iostream>
+
+#include <string>
+
+#include "base/strings/string_util.h"
+#include "base/strings/stringprintf.h"
+#include "net/tools/domain_security_preload_generator/cert_util.h"
+#include "net/tools/domain_security_preload_generator/huffman/huffman_frequency_tracker.h"
+#include "net/tools/domain_security_preload_generator/spki_hash.h"
+
+namespace net {
+
+namespace {
+
+static const char* kNewLine = "\n";

[agl, 2016/12/06 18:51:36]

ditto about [] vs *.

[martijnc, 2016/12/07 22:37:54]

Done.

+static const char* kIndent = "  ";
+
+std::string FormatSPKIName(const std::string& name) {
+  return "kSPKIHash_" + name;
+}
+
+std::string FormatAcceptedKeyName(const std::string& name) {
+  return "k" + name + "AcceptableCerts";
+}
+
+std::string FormatRejectedKeyName(const std::string& name) {
+  return "k" + name + "RejectedCerts";
+}
+
+std::string FormatReportURIName(const std::string& name) {
+  return "k" + name + "ReportURI";
+}
+
+std::string FormatVectorAsArray(const std::vector<uint8_t>& bytes) {
+  std::string output = "{";
+  output.append(kNewLine);
+  output.append(kIndent);
+  output.append(kIndent);
+
+  size_t bytes_on_current_line = 0;
+
+  for (size_t i = 0; i < bytes.size(); ++i) {
+    base::StringAppendF(&output, "0x%02x,", bytes[i]);
+
+    bytes_on_current_line++;
+    if (bytes_on_current_line >= 12 && (i + 1) < bytes.size()) {

[agl, 2016/12/06 18:51:36]

no need for parens around |i + 1|.

[martijnc, 2016/12/07 22:37:54]

Done.

+      output.append(kNewLine);
+      output.append(kIndent);
+      output.append(kIndent);
+
+      bytes_on_current_line = 0;
+    } else if ((i + 1) < bytes.size()) {

[agl, 2016/12/06 18:51:36]

no need for parens around |i + 1|.

[martijnc, 2016/12/07 22:37:54]

Done.

+      output.append(" ");
+    }
+  }
+
+  output.append(kNewLine);
+  output.append("}");
+
+  return output;
+}
+
+std::string WritePinsetList(const std::string& name,
+                            const std::vector<std::string>& pins) {
+  std::string output = "static const char* const " + name + "[] = {";
+  output.append(kNewLine);
+
+  for (const auto& pin_name : pins) {
+    output.append(kIndent);
+    output.append(kIndent);
+    output.append(FormatSPKIName(pin_name));
+    output.append(",");
+    output.append(kNewLine);
+  }
+
+  output.append(kIndent);
+  output.append(kIndent);
+  output.append("NULL,");
+  output.append(kNewLine);
+  output.append("};");
+
+  return output;
+}
+
+HuffmanRepresentationTable ApproximateHuffman(
+    const DomainSecurityEntries& entries) {
+  HuffmanFrequencyTracker tracker;
+  for (const auto& entry : entries) {
+    for (const auto& c : entry->hostname()) {
+      tracker.RecordUsage(c);
+    }
+
+    tracker.RecordUsage(TrieWriter::kTerminalValue);
+    tracker.RecordUsage(TrieWriter::kEndOfTableValue);
+  }
+
+  return tracker.ToHuffmanRepresentationTable();
+}
+
+}  // namespace
+
+PreloadedStateGenerator::PreloadedStateGenerator() {}
+
+PreloadedStateGenerator::~PreloadedStateGenerator() {}
+
+std::string PreloadedStateGenerator::Generate(
+    const std::string& preload_template,
+    const DomainSecurityEntries& entries,
+    const DomainIDList& domain_ids,
+    const Pinsets& pinsets,
+    bool verbose) {
+  template_.reset(new Template(preload_template));
+
+  NameIDMap domain_ids_map;
+  ProcessDomainIds(domain_ids, &domain_ids_map);
+
+  ProcessSPKIHashes(pinsets);
+
+  NameIDMap expect_ct_report_uri_map;
+  ProcessExpectCTURIs(entries, &expect_ct_report_uri_map);
+
+  NameIDMap expect_staple_report_uri_map;
+  ProcessExpectStapleURIs(entries, &expect_staple_report_uri_map);
+
+  NameIDMap pinsets_map;
+  ProcessPinsets(pinsets, &pinsets_map);
+
+  HuffmanRepresentationTable table = ApproximateHuffman(entries);
+  HuffmanFrequencyTracker tracker;
+  TrieWriter writer(table, domain_ids_map, expect_ct_report_uri_map,
+                    expect_staple_report_uri_map, pinsets_map, &tracker);
+  writer.WriteEntries(entries);
+  uint32_t initial_length = writer.position();
+
+  HuffmanRepresentationTable optimal_table =
+      tracker.ToHuffmanRepresentationTable();
+  TrieWriter new_writer(optimal_table, domain_ids_map, expect_ct_report_uri_map,
+                        expect_staple_report_uri_map, pinsets_map, nullptr);
+
+  int root_position = new_writer.WriteEntries(entries);
+  uint32_t new_length = new_writer.position();
+
+  std::vector<uint8_t> huffman_tree = tracker.AsVector();
+
+  new_writer.close();
+
+  template_->ReplaceTag("HUFFMAN_TREE", FormatVectorAsArray(huffman_tree));
+  template_->ReplaceTag("HSTS_TRIE", FormatVectorAsArray(new_writer.bytes()));
+
+  template_->ReplaceTag("HSTS_TRIE_BITS", std::to_string(new_length));
+  template_->ReplaceTag("HSTS_TRIE_ROOT", std::to_string(root_position));
+
+  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 template_->GetOutput();
+}
+
+void PreloadedStateGenerator::ProcessDomainIds(const DomainIDList& domain_ids,
+                                               NameIDMap* map) {
+  DCHECK(template_.get());
+
+  std::string output = "{";
+  output.append(kNewLine);
+
+  for (size_t i = 0; i < domain_ids.size(); ++i) {
+    const std::string& current = domain_ids.at(i);
+    output.append(kIndent);
+    output.append("DOMAIN_" + current + ",");
+    output.append(kNewLine);
+
+    map->insert(NameIDPair(current, i));
+  }
+
+  output.append(kIndent);
+  output.append("// Boundary value for UMA_HISTOGRAM_ENUMERATION.");
+  output.append(kNewLine);
+  output.append(kIndent);
+  output.append("DOMAIN_NUM_EVENTS,");
+  output.append(kNewLine);
+  output.append("}");
+
+  template_->ReplaceTag("DOMAIN_IDS", output);
+}
+
+void PreloadedStateGenerator::ProcessSPKIHashes(const Pinsets& pinset) {
+  DCHECK(template_.get());
+
+  std::string output;
+
+  const SPKIHashMap& hashes = pinset.spki_hashes();
+  for (const auto& current : hashes) {
+    const std::string& name = current.first;
+    const SPKIHash& hash = current.second;
+
+    output.append("static const char " + FormatSPKIName(name) + "[] =");
+    output.append(kNewLine);
+
+    for (size_t i = 0; i < hash.size() / 16; ++i) {
+      output.append(kIndent);
+      output.append(kIndent);
+      output.append("\"");
+
+      for (size_t j = i * 16; j < ((i + 1) * 16); ++j) {
+        base::StringAppendF(&output, "\\x%02x", hash.data()[j]);
+      }
+
+      output.append("\"");
+      if (i + 1 == hash.size() / 16) {
+        output.append(";");
+      }
+      output.append(kNewLine);
+    }
+
+    output.append(kNewLine);
+  }
+
+  base::TrimString(output, kNewLine, &output);
+  template_->ReplaceTag("SPKI_HASHES", output);
+}
+
+void PreloadedStateGenerator::ProcessExpectCTURIs(
+    const DomainSecurityEntries& entries,
+    NameIDMap* expect_ct_report_uri_map) {
+  std::string output = "{";
+  output.append(kNewLine);
+
+  for (const auto& entry : entries) {
+    const std::string& url = entry->expect_ct_report_uri();
+    if (entry->expect_ct() && url.size() &&
+        expect_ct_report_uri_map->find(url) ==
+            expect_ct_report_uri_map->cend()) {
+      output.append(kIndent);
+      output.append(kIndent);
+      output.append("\"" + entry->expect_ct_report_uri() + "\",");
+      output.append(kNewLine);
+
+      expect_ct_report_uri_map->insert(NameIDPair(
+          entry->expect_ct_report_uri(), expect_ct_report_uri_map->size()));
+    }
+  }
+
+  output.append("}");
+  template_->ReplaceTag("EXPECT_CT_REPORT_URIS", output);
+}
+
+void PreloadedStateGenerator::ProcessExpectStapleURIs(
+    const DomainSecurityEntries& entries,
+    NameIDMap* expect_staple_report_uri_map) {
+  std::string output = "{";
+  output.append(kNewLine);
+
+  for (const auto& entry : entries) {
+    const std::string& url = entry->expect_staple_report_uri();
+    if (entry->expect_staple() && url.size() &&
+        expect_staple_report_uri_map->find(url) ==
+            expect_staple_report_uri_map->cend()) {
+      output.append(kIndent);
+      output.append(kIndent);
+      output.append("\"" + entry->expect_staple_report_uri() + "\",");
+      output.append(kNewLine);
+
+      expect_staple_report_uri_map->insert(
+          NameIDPair(entry->expect_staple_report_uri(),
+                     expect_staple_report_uri_map->size()));
+    }
+  }
+
+  output.append("}");
+  template_->ReplaceTag("EXPECT_STAPLE_REPORT_URIS", output);
+}
+
+void PreloadedStateGenerator::ProcessPinsets(const Pinsets& pinset,
+                                             NameIDMap* pinset_map) {
+  std::string certs_output;
+  std::string pinsets_output = "{";
+  pinsets_output.append(kNewLine);
+
+  const PinsetMap& pinsets = pinset.pinsets();
+  for (const auto& current : pinsets) {
+    const std::unique_ptr<Pinset>& pinset = current.second;
+    std::string uppercased_name = pinset->name();
+    uppercased_name[0] = base::ToUpperASCII(uppercased_name[0]);
+
+    const std::string& accepted_pins_names =
+        FormatAcceptedKeyName(uppercased_name);
+    certs_output.append(
+        WritePinsetList(accepted_pins_names, pinset->static_spki_hashes()));
+    certs_output.append(kNewLine);
+
+    std::string rejected_pins_names = "kNoRejectedPublicKeys";
+    if (pinset->bad_static_spki_hashes().size()) {
+      rejected_pins_names = FormatRejectedKeyName(uppercased_name);
+      certs_output.append(WritePinsetList(rejected_pins_names,
+                                          pinset->bad_static_spki_hashes()));
+      certs_output.append(kNewLine);
+    }
+
+    std::string report_uri = "kNoReportURI";
+    if (pinset->report_uri().size()) {
+      report_uri = FormatReportURIName(uppercased_name);
+      certs_output.append("static const char " + report_uri + "[] = ");
+      certs_output.append("\"");
+      certs_output.append(pinset->report_uri());
+      certs_output.append("\";");
+      certs_output.append(kNewLine);
+    }
+    certs_output.append(kNewLine);
+
+    pinsets_output.append(kIndent);
+    pinsets_output.append(kIndent);
+    pinsets_output.append("{" + accepted_pins_names + ", " +
+                          rejected_pins_names + ", " + report_uri + "},");
+    pinsets_output.append(kNewLine);
+
+    pinset_map->insert(NameIDPair(pinset->name(), pinset_map->size()));
+  }
+
+  pinsets_output.append("}");
+
+  base::TrimString(certs_output, kNewLine, &certs_output);
+
+  template_->ReplaceTag("ACCEPTABLE_CERTS", certs_output);
+  template_->ReplaceTag("PINSETS", pinsets_output);
+}
+
+}  // namespace net