Merge pull request #1218 from ChTimTsubasa/SSL_handshake_refactor

SSL handshake refactor

Author: mengranwo

src/include/common/internal_types.h

@@ -74,7 +74,7 @@ extern int TEST_TUPLES_PER_TILEGROUP;
 enum class CmpBool {
   FALSE = 0,
   TRUE = 1,
-  NULL_ = 2 // Note the underscore suffix
+  NULL_ = 2  // Note the underscore suffix
 };
 
 //===--------------------------------------------------------------------===//
@@ -115,7 +115,6 @@ std::string PostgresValueTypeToString(PostgresValueType type);
 PostgresValueType StringToPostgresValueType(const std::string &str);
 std::ostream &operator<<(std::ostream &os, const PostgresValueType &type);
 
-
 //===--------------------------------------------------------------------===//
 // Predicate Expression Operation Types
 //===--------------------------------------------------------------------===//
@@ -631,10 +630,12 @@ std::string DropTypeToString(DropType type);
 DropType StringToDropType(const std::string &str);
 std::ostream &operator<<(std::ostream &os, const DropType &type);
 
-template<class E> class EnumHash {
+template <class E>
+class EnumHash {
  public:
-  size_t operator()(const E&e) const {
-    return std::hash<typename std::underlying_type<E>::type>()(static_cast<typename std::underlying_type<E>::type>(e));
+  size_t operator()(const E &e) const {
+    return std::hash<typename std::underlying_type<E>::type>()(
+        static_cast<typename std::underlying_type<E>::type>(e));
   }
 };
 
@@ -665,9 +666,9 @@ enum class StatementType {
   ALTER = 12,                 // alter statement type
   TRANSACTION = 13,           // transaction statement type,
   COPY = 14,                  // copy type
-  ANALYZE = 15,                // analyze type
+  ANALYZE = 15,               // analyze type
   VARIABLE_SET = 16,          // variable set statement type
-  CREATE_FUNC = 17,            // create func statement type
+  CREATE_FUNC = 17,           // create func statement type
 };
 std::string StatementTypeToString(StatementType type);
 StatementType StringToStatementType(const std::string &str);
@@ -678,24 +679,24 @@ std::ostream &operator<<(std::ostream &os, const StatementType &type);
 //===--------------------------------------------------------------------===//
 
 enum class QueryType {
-  QUERY_BEGIN = 0,                // begin query
-  QUERY_COMMIT = 1,               // commit query
-  QUERY_ROLLBACK = 2,             // rollback query
-  QUERY_CREATE_TABLE = 3,               // create query
+  QUERY_BEGIN = 0,         // begin query
+  QUERY_COMMIT = 1,        // commit query
+  QUERY_ROLLBACK = 2,      // rollback query
+  QUERY_CREATE_TABLE = 3,  // create query
   QUERY_CREATE_DB = 4,
   QUERY_CREATE_INDEX = 5,
-  QUERY_DROP = 6,                // other queries
-  QUERY_INSERT = 7,               // insert query
-  QUERY_PREPARE = 8,	      // prepare query
-  QUERY_EXECUTE = 9, 	      // execute query
+  QUERY_DROP = 6,     // other queries
+  QUERY_INSERT = 7,   // insert query
+  QUERY_PREPARE = 8,  // prepare query
+  QUERY_EXECUTE = 9,  // execute query
   QUERY_UPDATE = 10,
   QUERY_DELETE = 11,
   QUERY_RENAME = 12,
   QUERY_ALTER = 13,
   QUERY_COPY = 14,
   QUERY_ANALYZE = 15,
-  QUERY_SET = 16,                  // set query
-  QUERY_SHOW = 17,                 // show query
+  QUERY_SET = 16,   // set query
+  QUERY_SHOW = 17,  // show query
   QUERY_SELECT = 18,
   QUERY_OTHER = 19,
   QUERY_INVALID = 20,
@@ -705,8 +706,11 @@ enum class QueryType {
 };
 std::string QueryTypeToString(QueryType query_type);
 QueryType StringToQueryType(std::string str);
-namespace parser{ class SQLStatement;}
-QueryType StatementTypeToQueryType(StatementType stmt_type, const parser::SQLStatement* sql_stmt);
+namespace parser {
+class SQLStatement;
+}
+QueryType StatementTypeToQueryType(StatementType stmt_type,
+                                   const parser::SQLStatement *sql_stmt);
 //===--------------------------------------------------------------------===//
 // Scan Direction Types
 //===--------------------------------------------------------------------===//
@@ -1014,7 +1018,8 @@ const int TRIGGER_TYPE_MAX = TRIGGER_TYPE_ROW | TRIGGER_TYPE_STATEMENT |
 
 // Statistics Collection Type
 // Disable or enable
-// TODO: This should probably be a collection level and not a boolean (enable/disable)
+// TODO: This should probably be a collection level and not a boolean
+// (enable/disable)
 enum class StatsType {
   // Disable statistics collection
   INVALID = INVALID_TYPE_ID,
@@ -1093,10 +1098,7 @@ static const int INVALID_FILE_DESCRIPTOR = -1;
 // Tuple serialization formats
 // ------------------------------------------------------------------
 
-enum class TupleSerializationFormat {
-  NATIVE = 0,
-  DR = 1
-};
+enum class TupleSerializationFormat { NATIVE = 0, DR = 1 };
 
 // ------------------------------------------------------------------
 // Entity types
@@ -1201,7 +1203,7 @@ std::ostream &operator<<(std::ostream &os, const RWType &type);
 
 // ItemPointer -> type
 typedef CuckooMap<ItemPointer, RWType, ItemPointerHasher, ItemPointerComparator>
-   ReadWriteSet;
+    ReadWriteSet;
 
 // this enum is to identify why the version should be GC'd.
 enum class GCVersionType {
@@ -1397,15 +1399,15 @@ enum class ProcessResult {
   COMPLETE,
   TERMINATE,
   PROCESSING,
-  MORE_DATA_REQUIRED
+  MORE_DATA_REQUIRED,
+  NEED_SSL_HANDSHAKE,
 };
 
 enum class NetworkProtocolType {
   POSTGRES_JDBC,
   POSTGRES_PSQL,
 };
 
-
 enum class SSLLevel {
   SSL_DISABLE = 0,
   SSL_PREFER = 1,

src/include/network/connection_handle.h

@@ -87,6 +87,10 @@ class ConnectionHandle {
   Transition ProcessWrite();
   Transition GetResult();
   Transition CloseSocket();
+  /**
+   * Flush out all the responses and do real SSL handshake
+   */
+  Transition ProcessWrite_SSLHandshake();
 
  private:
   /**
@@ -144,16 +148,7 @@ class ConnectionHandle {
   friend class ConnectionHandleFactory;
 
   ConnectionHandle(int sock_fd, ConnectionHandlerTask *handler,
-                   std::shared_ptr<Buffer> rbuf, std::shared_ptr<Buffer> wbuf,
-                   bool ssl_able);
-
-  ProcessResult ProcessInitial();
-
-  /**
-   * Extracts the header of a Postgres start up packet from the read socket
-   * buffer
-   */
-  static bool ReadStartupPacketHeader(Buffer &rbuf, InputPacket &rpkt);
+                   std::shared_ptr<Buffer> rbuf, std::shared_ptr<Buffer> wbuf);
 
   /**
    * Writes a packet's header (type, size) into the write buffer
@@ -171,6 +166,16 @@ class ConnectionHandle {
    */
   WriteState FlushWriteBuffer();
 
+  /**
+   * @brief: process SSL handshake to generate valid SSL
+   * connection context for further communications
+   * @return FINISH when the SSL handshake failed
+   *         PROCEED when the SSL handshake success
+   *         NEED_DATA when the SSL handshake is partially done due to network
+   *         latency
+   */
+  Transition SSLHandshake();
+
   /**
    * Set the socket to non-blocking mode
    */
@@ -190,6 +195,16 @@ class ConnectionHandle {
     setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof one);
   }
 
+  /**
+   * @brief: Determine if there is still responses in the buffer
+   * @return true if there is still responses to flush out in either wbuf or
+   * responses
+   */
+  inline bool HasResponse() {
+    return (protocol_handler_->responses_.size() != 0) ||
+           (wbuf_->buf_size != 0);
+  }
+
   int sock_fd_;                            // socket file descriptor
   struct event *network_event = nullptr;   // something to read from network
   struct event *workpool_event = nullptr;  // worker thread done the job
@@ -204,16 +219,8 @@ class ConnectionHandle {
   std::shared_ptr<Buffer> rbuf_;    // Socket's read buffer
   std::shared_ptr<Buffer> wbuf_;    // Socket's write buffer
   unsigned int next_response_ = 0;  // The next response in the response buffer
-  Client client_;
-  StateMachine state_machine_;
 
-  // TODO(Tianyi) Can we encapsulate these flags?
-  bool ssl_handshake_ = false;
-  bool finish_startup_packet_ = false;
-  bool ssl_able_;
-
-  // TODO(Tianyi) hide this in protocol handler
-  InputPacket initial_packet_;
+  StateMachine state_machine_;
 
   short curr_event_flag_;  // current libevent event flag
 };

src/include/network/connection_handle_factory.h

@@ -40,22 +40,21 @@ class ConnectionHandleFactory {
     // (probably also to-do: beat up the person who wrote this)
     PelotonServer::recent_connfd = conn_fd;
     auto it = reusable_handles_.find(conn_fd);
-    bool ssl_able = (PelotonServer::GetSSLLevel() != SSLLevel::SSL_DISABLE);
     if (it == reusable_handles_.end()) {
       // We are not using std::make_shared here because we want to keep
       // ConnectionHandle constructor
       // private to avoid unintentional use.
       auto handle = std::shared_ptr<ConnectionHandle>(
           new ConnectionHandle(conn_fd, handler, std::make_shared<Buffer>(),
-                               std::make_shared<Buffer>(), ssl_able));
+                               std::make_shared<Buffer>()));
       reusable_handles_[conn_fd] = handle;
       return handle;
     }
 
     it->second->rbuf_->Reset();
     it->second->wbuf_->Reset();
     std::shared_ptr<ConnectionHandle> new_handle(new ConnectionHandle(
-        conn_fd, handler, it->second->rbuf_, it->second->wbuf_, ssl_able));
+        conn_fd, handler, it->second->rbuf_, it->second->wbuf_));
     reusable_handles_[conn_fd] = new_handle;
     return new_handle;
   }

src/include/network/marshal.h

@@ -15,9 +15,9 @@
 #include <string>
 #include <vector>
 
+#include "common/internal_types.h"
 #include "common/logger.h"
 #include "common/macros.h"
-#include "common/internal_types.h"
 
 #define BUFFER_INIT_SIZE 100
 
@@ -138,6 +138,8 @@ struct OutputPacket {
   size_t ptr;                   // ByteBuf cursor, which is used for get and put
   NetworkMessageType msg_type;  // header
 
+  bool single_type_pkt;  // there would be only a pkt type being written to the
+                         // buffer when this flag is true
   bool skip_header_write;  // whether we should write header to socket wbuf
   size_t write_ptr;        // cursor used to write packet content to socket wbuf
 
@@ -146,28 +148,16 @@ struct OutputPacket {
     buf.resize(BUFFER_INIT_SIZE);
     buf.shrink_to_fit();
     buf.clear();
+    single_type_pkt = false;
     len = ptr = write_ptr = 0;
     msg_type = NetworkMessageType::NULL_COMMAND;
     skip_header_write = true;
   }
 };
 
-struct Client {
-  // Authentication details
-  std::string dbname;
-  std::string user;
-  std::unordered_map<std::string, std::string> cmdline_options;
-
-  inline void Reset() {
-    dbname.clear();
-    user.clear();
-    cmdline_options.clear();
-  }
-};
-
 /*
-* Marshallers
-*/
+ * Marshallers
+ */
 
 /* packet_put_byte - used to write a single byte into a packet */
 extern void PacketPutByte(OutputPacket *pkt, const uchar c);
@@ -186,22 +176,22 @@ extern void PacketPutCbytes(OutputPacket *pkt, const uchar *b, int len);
 extern void PacketPutString(OutputPacket *pkt, const std::string &data);
 
 /*
-* Unmarshallers
-*/
+ * Unmarshallers
+ */
 
 /* Copy len bytes from the position indicated by begin to an array */
 extern uchar *PacketCopyBytes(ByteBuf::const_iterator begin, int len);
 /*
-* packet_get_int -  Parse an int out of the head of the
-* 	packet. "base" bytes determine the number of bytes of integer
-* 	we are parsing out.
-*/
+ * packet_get_int -  Parse an int out of the head of the
+ * 	packet. "base" bytes determine the number of bytes of integer
+ * 	we are parsing out.
+ */
 extern int PacketGetInt(InputPacket *pkt, uchar base);
 
 /*
-* packet_get_string - parse out a string of size len.
-* 		if len=0? parse till the end of the string
-*/
+ * packet_get_string - parse out a string of size len.
+ * 		if len=0? parse till the end of the string
+ */
 extern void PacketGetString(InputPacket *pkt, size_t len, std::string &result);
 
 /* packet_get_bytes - Parse out "len" bytes of pkt as raw bytes */
@@ -211,9 +201,9 @@ extern void PacketGetBytes(InputPacket *pkt, size_t len, ByteBuf &result);
 extern void PacketGetByte(InputPacket *rpkt, uchar &result);
 
 /*
-* get_string_token - used to extract a string token
-* 		from an unsigned char vector
-*/
+ * get_string_token - used to extract a string token
+ * 		from an unsigned char vector
+ */
 extern void GetStringToken(InputPacket *pkt, std::string &result);
 
 }  // namespace network

src/include/network/network_state.h

@@ -21,10 +21,11 @@ namespace network {
 enum class ConnState {
   READ,        // State that reads data from the network
   WRITE,       // State the writes data to the network
-  WAIT,        // State for waiting for some event to happen
   PROCESS,     // State that runs the network protocol on received data
   CLOSING,     // State for closing the client connection
   GET_RESULT,  // State when triggered by worker thread that completes the task.
+  PROCESS_WRITE_SSL_HANDSHAKE,  // State to flush out responses and doing (Real)
+                                // SSL handshake
 };
 
 // TODO(tianyu): Convert use cases of this to just return Transition
@@ -46,7 +47,8 @@ enum class Transition {
   // TODO(tianyu) generalize this symbol, this is currently only used in process
   GET_RESULT,
   FINISH,
-  RETRY
+  RETRY,
+  NEED_SSL_HANDSHAKE,
 };
-}
-}
+}  // namespace network
+}  // namespace peloton