rocksplicator/common/ssl_context_manager.cpp at master · pinterest/rocksplicator · GitHub - github.com

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
/// Copyright 2019 Pinterest Inc.
///
/// Licensed under the Apache License, Version 2.0 (the "License");
/// you may not use this file except in compliance with the License.
/// You may obtain a copy of the License at
///
/// http://www.apache.org/licenses/LICENSE-2.0

/// Unless required by applicable law or agreed to in writing, software
/// distributed under the License is distributed on an "AS IS" BASIS,
/// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
/// See the License for the specific language governing permissions and
/// limitations under the License.

//
// @author bol (bol@pinterest.com)
//

#include "common/ssl_context_manager.h"

#include <memory>
#include <string>

#include "common/future_util.h"
#include "common/stats/stats.h"
#include "folly/io/async/SSLContext.h"

DEFINE_string(tls_certfile, "", "Certificate file path location for TLS");

DEFINE_string(tls_trusted_certfile, "",
              "Trusted certificate file path location for TLS");

DEFINE_string(tls_keyfile, "", "Key file path location for TLS");

namespace {

// create a new SSLContext from files, return nullptr on error
std::shared_ptr<folly::SSLContext> loadSSLContext() {
  auto ctx = std::make_shared<folly::SSLContext>();

  try {
    ctx->loadCertificate(FLAGS_tls_certfile.c_str());
    ctx->loadPrivateKey(FLAGS_tls_keyfile.c_str());
    ctx->loadTrustedCertificates(FLAGS_tls_trusted_certfile.c_str());
  } catch (const std::exception& ex) {
    static const std::string kSSLContextRefreshFailures =
      "ssl_context_refresh_failures";
    common::Stats::get()->Incr(kSSLContextRefreshFailures);

    LOG(ERROR) << "Failed to refresh SSLContext: " << ex.what();
    return nullptr;
  }

  return ctx;
}

// schedule periodial refresh for cur_ctx
void scheduleRefresh(std::shared_ptr<folly::SSLContext>* cur_ctx) {
  auto delayed_future =
    common::GenerateDelayedFuture(std::chrono::seconds(60 * 60));
#if __GNUC__ >= 8
  std::move(delayed_future).then([cur_ctx] (auto&&) {
#else
  delayed_future.then([cur_ctx] () {
#endif
      auto new_ctx = loadSSLContext();
      if (new_ctx) {
        LOG(INFO) << "Got a new SSLContext, swapping";
        static const std::string kSSLContextRefreshTimes =
          "ssl_context_refresh_times";
        common::Stats::get()->Incr(kSSLContextRefreshTimes);
        std::atomic_exchange_explicit(cur_ctx, new_ctx,
                                      std::memory_order_release);
      }

      scheduleRefresh(cur_ctx);
    });
}

std::once_flag schedule_flag;

}  // namespace


namespace common {

const std::shared_ptr<folly::SSLContext>* getSSLContext() {
  static const bool no_tls = FLAGS_tls_certfile.empty() ||
    FLAGS_tls_trusted_certfile.empty() ||
    FLAGS_tls_keyfile.empty();

  // always return nullptr if any of the gflags not set
  if (no_tls) {
    return nullptr;
  }

  static auto ctx = loadSSLContext();


  // always return nullptr if failed to load the context for the first time
  // static object ctx can be accessed from multiple thread, so we have to use
  // atomic functions to read/write it
  if (std::atomic_load_explicit(&ctx, std::memory_order_acquire) == nullptr) {
    return nullptr;
  }

  // succeeded to load the context, schedule for future refresh once
  std::call_once(schedule_flag, scheduleRefresh, &ctx);

  return &ctx;
}

}  // namespace common