Fix Huffman tree construction.

net/tools/transport_security_state_generator/huffman/huffman_builder.cc

 // 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/transport_security_state_generator/huffman/huffman_builder.h"
 
 #include <algorithm>
 
 #include "base/logging.h"
+#include "base/memory/ptr_util.h"
 
 namespace net {
 
 namespace transport_security_state {
 
 namespace {
 
 class HuffmanNode {
  public:
   HuffmanNode(uint8_t value,
@@ skipping 27 matching lines @@
   return lhs->count() < rhs->count();
 }
 
 }  // namespace
 
 HuffmanBuilder::HuffmanBuilder() {}
 
 HuffmanBuilder::~HuffmanBuilder() {}
 
 void HuffmanBuilder::RecordUsage(uint8_t character) {
-  counts_[character] += 1;
+  DCHECK(character < 128);
+  counts_[character & 127] += 1;
 }
 
 HuffmanRepresentationTable HuffmanBuilder::ToTable() {
   HuffmanRepresentationTable table;
   std::unique_ptr<HuffmanNode> node(BuildTree());
 
   TreeToTable(node.get(), 0, 0, &table);
   return table;
 }
 
@@ skipping 24 matching lines @@
 uint32_t HuffmanBuilder::WriteToVector(HuffmanNode* node,
                                        std::vector<uint8_t>* vector) {
   uint8_t left_value;
   uint8_t right_value;
   uint32_t child_position;
 
   if (node->left()->IsLeaf()) {
     left_value = 128 | node->left()->value();
   } else {
     child_position = WriteToVector(node->left().get(), vector);
-    CHECK(child_position < 512) << "huffman tree too large";
+    DCHECK(child_position < 512) << "huffman tree too large";
     left_value = child_position / 2;
   }
 
   if (node->right()->IsLeaf()) {
     right_value = 128 | node->right()->value();
   } else {
     child_position = WriteToVector(node->right().get(), vector);
-    CHECK(child_position < 512) << "huffman tree to large";
+    DCHECK(child_position < 512) << "huffman tree to large";
     right_value = child_position / 2;
   }
 
   uint32_t position = static_cast<uint32_t>(vector->size());
   vector->push_back(left_value);
   vector->push_back(right_value);
   return position;
 }
 
 std::unique_ptr<HuffmanNode> HuffmanBuilder::BuildTree() {
   std::vector<std::unique_ptr<HuffmanNode>> nodes;
   nodes.reserve(counts_.size());
 
   for (const auto& item : counts_) {
-    if (item.second > 0) {

[martijnc, 2017/02/15 22:09:09]

This check is always true.

Only RecordUsage() writes to |counts_| and only increments the value. If a
character is never used it won't be in |counts_| so the code wouldn't iterate
over it. If a character is used, its count would be 1 or higher.

-      std::unique_ptr<HuffmanNode> node(
-          new HuffmanNode(item.first, item.second, nullptr, nullptr));
-      nodes.push_back(std::move(node));
-    }
+    nodes.push_back(base::MakeUnique<HuffmanNode>(item.first, item.second,
+                                                  nullptr, nullptr));
   }
 
-  if (nodes.size() < 2) {
-    return std::move(nodes[0]);
+  // At least 2 entries are required for everything to work properly. Add
+  // arbitrary values to fill the tree.
+  for (uint8_t i = 0; nodes.size() < 2 && i < 2; ++i) {
+    for (const auto& node : nodes) {
+      if (node->value() == i) {
+        break;
+      }
+    }
+
+    nodes.push_back(base::MakeUnique<HuffmanNode>(i, 0, nullptr, nullptr));
   }
 
   std::stable_sort(nodes.begin(), nodes.end(), CompareNodes);
 
   while (nodes.size() > 1) {
     std::unique_ptr<HuffmanNode> a = std::move(nodes[0]);
     std::unique_ptr<HuffmanNode> b = std::move(nodes[1]);
 
     uint32_t count_a = a->count();
     uint32_t count_b = b->count();
 
     std::unique_ptr<HuffmanNode> parent(
         new HuffmanNode(0, count_a + count_b, std::move(a), std::move(b)));
 
     nodes.erase(nodes.begin());
     nodes[0] = std::move(parent);
 
     std::stable_sort(nodes.begin(), nodes.end(), CompareNodes);
   }
 
   return std::move(nodes[0]);
 }
 
 }  // namespace transport_security_state
 
 }  // namespace net