From bc31146c1f369902fcf72ccc30aadc913571bea4 Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Sat, 10 Sep 2016 22:02:46 +0900
Subject: [PATCH] nghttpx: Add tls_sni to mruby Nghttpx::Env class

tls_sni returns TLS SNI value which client sent in this TLS
connection.
---
 doc/nghttpx.h2r               |  4 ++++
 src/shrpx_client_handler.cc   |  4 ++++
 src/shrpx_client_handler.h    |  8 ++++++++
 src/shrpx_mruby_module_env.cc | 14 ++++++++++++++
 src/shrpx_ssl.cc              | 17 ++++++++++-------
 5 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/doc/nghttpx.h2r b/doc/nghttpx.h2r
index e7fd9b52..04ca05fd 100644
--- a/doc/nghttpx.h2r
+++ b/doc/nghttpx.h2r
@@ -291,6 +291,10 @@ respectively.
 
         Return true if TLS is used on the connection.
 
+    .. rb:attr_reader:: tls_sni
+
+        Return the TLS SNI value which client sent in this connection.
+
 .. rb:class:: Request
 
     Object to represent request from client.  The modification to
diff --git a/src/shrpx_client_handler.cc b/src/shrpx_client_handler.cc
index 82090076..4feb973c 100644
--- a/src/shrpx_client_handler.cc
+++ b/src/shrpx_client_handler.cc
@@ -1451,4 +1451,8 @@ const UpstreamAddr *ClientHandler::get_upstream_addr() const { return faddr_; }
 
 Connection *ClientHandler::get_connection() { return &conn_; };
 
+void ClientHandler::set_tls_sni(const StringRef &sni) { sni_ = sni.str(); }
+
+StringRef ClientHandler::get_tls_sni() const { return StringRef{sni_}; }
+
 } // namespace shrpx
diff --git a/src/shrpx_client_handler.h b/src/shrpx_client_handler.h
index 22b3c500..567972c4 100644
--- a/src/shrpx_client_handler.h
+++ b/src/shrpx_client_handler.h
@@ -156,6 +156,12 @@ public:
 
   Connection *get_connection();
 
+  // Stores |sni| which is TLS SNI extension value client sent in this
+  // connection.
+  void set_tls_sni(const StringRef &sni);
+  // Returns TLS SNI extension value client sent in this connection.
+  StringRef get_tls_sni() const;
+
 private:
   Connection conn_;
   ev_timer reneg_shutdown_timer_;
@@ -169,6 +175,8 @@ private:
   // The client address used in "for" parameter of Forwarded header
   // field.
   std::string forwarded_for_;
+  // lowercased TLS SNI which client sent.
+  std::string sni_;
   std::function<int(ClientHandler &)> read_, write_;
   std::function<int(ClientHandler &)> on_read_, on_write_;
   // Address of frontend listening socket
diff --git a/src/shrpx_mruby_module_env.cc b/src/shrpx_mruby_module_env.cc
index 95691046..7fabbedb 100644
--- a/src/shrpx_mruby_module_env.cc
+++ b/src/shrpx_mruby_module_env.cc
@@ -127,6 +127,18 @@ mrb_value env_get_tls_used(mrb_state *mrb, mrb_value self) {
 }
 } // namespace
 
+namespace {
+mrb_value env_get_tls_sni(mrb_state *mrb, mrb_value self) {
+  auto data = static_cast<MRubyAssocData *>(mrb->ud);
+  auto downstream = data->downstream;
+  auto upstream = downstream->get_upstream();
+  auto handler = upstream->get_client_handler();
+  auto sni = handler->get_tls_sni();
+
+  return mrb_str_new(mrb, sni.c_str(), sni.size());
+}
+} // namespace
+
 void init_env_class(mrb_state *mrb, RClass *module) {
   auto env_class =
       mrb_define_class_under(mrb, module, "Env", mrb->object_class);
@@ -144,6 +156,8 @@ void init_env_class(mrb_state *mrb, RClass *module) {
                     MRB_ARGS_NONE());
   mrb_define_method(mrb, env_class, "tls_used", env_get_tls_used,
                     MRB_ARGS_NONE());
+  mrb_define_method(mrb, env_class, "tls_sni", env_get_tls_sni,
+                    MRB_ARGS_NONE());
 }
 
 } // namespace mruby
diff --git a/src/shrpx_ssl.cc b/src/shrpx_ssl.cc
index 80e1cfbd..be7c57e8 100644
--- a/src/shrpx_ssl.cc
+++ b/src/shrpx_ssl.cc
@@ -151,12 +151,6 @@ int servername_callback(SSL *ssl, int *al, void *arg) {
   auto conn = static_cast<Connection *>(SSL_get_app_data(ssl));
   auto handler = static_cast<ClientHandler *>(conn->data);
   auto worker = handler->get_worker();
-  auto cert_tree = worker->get_cert_lookup_tree();
-  if (!cert_tree) {
-    return SSL_TLSEXT_ERR_OK;
-  }
-
-  std::array<uint8_t, NI_MAXHOST> buf;
 
   auto rawhost = SSL_get_servername(ssl, TLSEXT_NAMETYPE_host_name);
   if (rawhost == nullptr) {
@@ -165,16 +159,25 @@ int servername_callback(SSL *ssl, int *al, void *arg) {
 
   auto len = strlen(rawhost);
   // NI_MAXHOST includes terminal NULL.
-  if (len == 0 || len + 1 > buf.size()) {
+  if (len == 0 || len + 1 > NI_MAXHOST) {
     return SSL_TLSEXT_ERR_OK;
   }
 
+  std::array<uint8_t, NI_MAXHOST> buf;
+
   auto end_buf = std::copy_n(rawhost, len, std::begin(buf));
 
   util::inp_strlower(std::begin(buf), end_buf);
 
   auto hostname = StringRef{std::begin(buf), end_buf};
 
+  handler->set_tls_sni(hostname);
+
+  auto cert_tree = worker->get_cert_lookup_tree();
+  if (!cert_tree) {
+    return SSL_TLSEXT_ERR_OK;
+  }
+
   auto idx = cert_tree->lookup(hostname);
   if (idx == -1) {
     return SSL_TLSEXT_ERR_OK;