From fbb8a926d8cc510a3ad870b96e1bd38b510034fd Mon Sep 17 00:00:00 2001
From: Neel Goyal <ngoyal@fb.com>
Date: Tue, 15 Dec 2015 06:35:14 -0800
Subject: [PATCH] Update SSLContext to use folly::Random and
 discrete_distribution

Summary:
Change next pick logic to use `std::discrete_distribution`.  Use folly::Random::DefaultGenerator as the generator.
Update the TARGETS to have sslcontext depend on random (which was the build break in D2741855)

Reviewed By: siyengar

Differential Revision: D2749166

fb-gh-sync-id: 7a6b400173b1e8c69a961f9e75e97d4777d5467a
---
 folly/io/async/SSLContext.cpp | 27 +++++++++++----------------
 folly/io/async/SSLContext.h   | 19 +++++++++++++------
 2 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/folly/io/async/SSLContext.cpp b/folly/io/async/SSLContext.cpp
index 7ab01c30..4e8ea69f 100644
--- a/folly/io/async/SSLContext.cpp
+++ b/folly/io/async/SSLContext.cpp
@@ -84,6 +84,10 @@ SSLContext::SSLContext(SSLVersion version) {
   SSL_CTX_set_tlsext_servername_callback(ctx_, baseServerNameOpenSSLCallback);
   SSL_CTX_set_tlsext_servername_arg(ctx_, this);
 #endif
+
+#ifdef OPENSSL_NPN_NEGOTIATED
+  Random::seed(nextProtocolPicker_);
+#endif
 }
 
 SSLContext::~SSLContext() {
@@ -374,16 +378,16 @@ bool SSLContext::setRandomizedAdvertisedNextProtocols(
       dst += protoLength;
     }
     total_weight += item.weight;
-    advertised_item.probability = item.weight;
     advertisedNextProtocols_.push_back(advertised_item);
+    advertisedNextProtocolWeights_.push_back(item.weight);
   }
   if (total_weight == 0) {
     deleteNextProtocolsStrings();
     return false;
   }
-  for (auto& advertised_item : advertisedNextProtocols_) {
-    advertised_item.probability /= total_weight;
-  }
+  nextProtocolDistribution_ =
+      std::discrete_distribution<>(advertisedNextProtocolWeights_.begin(),
+                                   advertisedNextProtocolWeights_.end());
   if ((uint8_t)protocolType & (uint8_t)NextProtocolType::NPN) {
     SSL_CTX_set_next_protos_advertised_cb(
         ctx_, advertisedNextProtocolCallback, this);
@@ -406,6 +410,7 @@ void SSLContext::deleteNextProtocolsStrings() {
     delete[] protocols.protocols;
   }
   advertisedNextProtocols_.clear();
+  advertisedNextProtocolWeights_.clear();
 }
 
 void SSLContext::unsetNextProtocols() {
@@ -419,18 +424,8 @@ void SSLContext::unsetNextProtocols() {
 }
 
 size_t SSLContext::pickNextProtocols() {
-  unsigned char random_byte;
-  RAND_bytes(&random_byte, 1);
-  double random_value = random_byte / 255.0;
-  double sum = 0;
-  for (size_t i = 0; i < advertisedNextProtocols_.size(); ++i) {
-    sum += advertisedNextProtocols_[i].probability;
-    if (sum < random_value && i + 1 < advertisedNextProtocols_.size()) {
-      continue;
-    }
-    return i;
-  }
-  CHECK(false) << "Failed to pickNextProtocols";
+  CHECK(!advertisedNextProtocols_.empty()) << "Failed to pickNextProtocols";
+  return nextProtocolDistribution_(nextProtocolPicker_);
 }
 
 int SSLContext::advertisedNextProtocolCallback(SSL* ssl,
diff --git a/folly/io/async/SSLContext.h b/folly/io/async/SSLContext.h
index 316fc1e3..e20b093b 100644
--- a/folly/io/async/SSLContext.h
+++ b/folly/io/async/SSLContext.h
@@ -22,6 +22,7 @@
 #include <vector>
 #include <memory>
 #include <string>
+#include <random>
 
 #include <openssl/ssl.h>
 #include <openssl/tls1.h>
@@ -35,6 +36,8 @@
 #include <folly/folly-config.h>
 #endif
 
+#include <folly/Random.h>
+
 namespace folly {
 
 /**
@@ -87,12 +90,6 @@ class SSLContext {
     std::list<std::string> protocols;
   };
 
-  struct AdvertisedNextProtocolsItem {
-    unsigned char* protocols;
-    unsigned length;
-    double probability;
-  };
-
   // Function that selects a client protocol given the server's list
   using ClientProtocolFilterCallback = bool (*)(unsigned char**, unsigned int*,
                                         const unsigned char*, unsigned int);
@@ -458,10 +455,20 @@ class SSLContext {
   static bool initialized_;
 
 #ifdef OPENSSL_NPN_NEGOTIATED
+
+  struct AdvertisedNextProtocolsItem {
+    unsigned char* protocols;
+    unsigned length;
+  };
+
   /**
    * Wire-format list of advertised protocols for use in NPN.
    */
   std::vector<AdvertisedNextProtocolsItem> advertisedNextProtocols_;
+  std::vector<int> advertisedNextProtocolWeights_;
+  std::discrete_distribution<int> nextProtocolDistribution_;
+  Random::DefaultGenerator nextProtocolPicker_;
+
   static int sNextProtocolsExDataIndex_;
 
   static int advertisedNextProtocolCallback(SSL* ssl,
-- 
2.34.1