From 4f512d83bb46962cb8deaa604fbd68a7ca52e696 Mon Sep 17 00:00:00 2001 From: wangfeihuo Date: Thu, 7 Nov 2024 16:32:05 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=85=A8=E5=AF=86=E6=80=81?= =?UTF-8?q?=E6=97=B6=E7=9A=84=E9=83=A8=E5=88=86=E5=BC=82=E5=B8=B8=E6=89=93?= =?UTF-8?q?=E5=8D=B0=E4=BF=A1=E6=81=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/bin/psql/command.cpp | 9 +++++++- src/bin/psql/settings.h | 1 + src/bin/psql/startup.cpp | 20 +++++++++++++--- .../client_logic_common/statement_data.cpp | 3 +++ .../client_logic_processor/raw_value.cpp | 23 ++++++++++++++++++- .../libpq/client_logic_processor/raw_value.h | 1 + .../client_logic_processor/stmt_processor.cpp | 11 +++++++++ src/common/interfaces/libpq/fe-connect.cpp | 5 ++++ src/include/libpq/cl_state.h | 1 + 9 files changed, 69 insertions(+), 5 deletions(-) diff --git a/src/bin/psql/command.cpp b/src/bin/psql/command.cpp index c3bfd0e29..54f532702 100644 --- a/src/bin/psql/command.cpp +++ b/src/bin/psql/command.cpp @@ -1518,7 +1518,14 @@ static bool do_connect(char* dbname, char* user, char* host, char* port) values[7] = CONNECT_TIMEOUT; #ifdef HAVE_CE keywords[8] = "enable_ce"; - values[8] = (pset.enable_client_encryption) ? "1" : NULL; + if (!pset.enable_client_encryption_log) { + values[8] = (pset.enable_client_encryption) ? "1" : NULL; + } else { + values[8] = (pset.enable_client_encryption) ? "1_with_log" : NULL; + if (pset.enable_client_encryption) { + printf("do_connect with enable_client_encryption.\n"); + } + } #endif keywords[PARAMS_ARRAY_SIZE-1] = NULL; values[PARAMS_ARRAY_SIZE-1] = NULL; diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h index e4a34e665..2b49a44ff 100644 --- a/src/bin/psql/settings.h +++ b/src/bin/psql/settings.h @@ -98,6 +98,7 @@ typedef struct _psqlSettings { bool on_error_stop; bool quiet; bool enable_client_encryption; + bool enable_client_encryption_log; bool singleline; bool singlestep; bool maintance; diff --git a/src/bin/psql/startup.cpp b/src/bin/psql/startup.cpp index 36a808334..d7b390d66 100644 --- a/src/bin/psql/startup.cpp +++ b/src/bin/psql/startup.cpp @@ -578,7 +578,12 @@ int main(int argc, char* argv[]) values[8] = CONNECT_TIMEOUT; #ifdef HAVE_CE keywords[9] = "enable_ce"; - values[9] = (pset.enable_client_encryption) ? (char*)"1" : NULL; + if (!pset.enable_client_encryption_log) { + values[9] = (pset.enable_client_encryption) ? (char*)"1" : NULL; + } else { + values[9] = (pset.enable_client_encryption) ? (char*)"1_with_log" : NULL; + printf("startup with enable_client_encryption.\n"); + } #endif if (pset.maintance) { keywords[PARAMS_ARRAY_SIZE - 2] = "options"; @@ -1052,6 +1057,7 @@ static void parse_psql_options(int argc, char* const argv[], struct adhoc_opts* {"with-decryption", required_argument, NULL, 'D'}, {"with-module-params", required_argument, NULL, 'u'}, {"with-salt", required_argument, NULL, 1}, + {"enable_client_encryption_log", no_argument, NULL, '3'}, #if defined(USE_ASSERT_CHECKING) || defined(FASTCHECK) {"sql-parse", no_argument, NULL, 'g'}, #endif @@ -1077,10 +1083,10 @@ static void parse_psql_options(int argc, char* const argv[], struct adhoc_opts* rc = memset_s(options, sizeof(*options), 0, sizeof(*options)); check_memset_s(rc); - check_short_optOfVoid("aAc:d:eEf:F:gh:Hlk:L:mno:p:P:qCR:rsStT:U:v:W:VxXz?012", argc, argv); + check_short_optOfVoid("aAc:d:eEf:F:gh:Hlk:L:mno:p:P:qCR:rsStT:U:v:W:VxXz?0123", argc, argv); while ((c = getopt_long( - argc, argv, "aAc:d:D:eEf:F:gh:Hlk:u:L:mno:p:P:qCR:rsStT:U:v:W:VxXz?012", long_options, &optindex)) != -1) { + argc, argv, "aAc:d:D:eEf:F:gh:Hlk:u:L:mno:p:P:qCR:rsStT:U:v:W:VxXz?0123", long_options, &optindex)) != -1) { switch (c) { case 'a': if (!SetVariable(pset.vars, "ECHO", "all")) { @@ -1223,6 +1229,14 @@ static void parse_psql_options(int argc, char* const argv[], struct adhoc_opts* case 'C': pset.enable_client_encryption = true; break; + case '3': + pset.enable_client_encryption_log = true; + if (pset.enable_client_encryption) { + printf("running psql with client_encryption.\n"); + } else { + printf("running psql without client_encryption.\n"); + } + break; case 'r': #ifdef USE_READLINE useReadline = true; diff --git a/src/common/interfaces/libpq/client_logic_common/statement_data.cpp b/src/common/interfaces/libpq/client_logic_common/statement_data.cpp index 0351118ee..a35c77423 100644 --- a/src/common/interfaces/libpq/client_logic_common/statement_data.cpp +++ b/src/common/interfaces/libpq/client_logic_common/statement_data.cpp @@ -120,6 +120,7 @@ void StatementData::replace_raw_values() params.new_query = (char *)libpq_realloc(params.new_query, params.new_query_size, params.new_query_size + (new_size - original_size) + 1); if (params.new_query == NULL) { + fprintf(stderr, "cannot realloc memory for encrypt str.\n"); return; } } @@ -129,6 +130,8 @@ void StatementData::replace_raw_values() if (new_str != NULL) { check_memcpy_s(memcpy_s(params.new_query + raw_value->m_location, params.new_query_size - raw_value->m_location, new_str, new_size)); + } else { + fprintf(stderr, "cannot get the str after encrypt.\n"); } params.new_query[params.new_query_size] = '\0'; } diff --git a/src/common/interfaces/libpq/client_logic_processor/raw_value.cpp b/src/common/interfaces/libpq/client_logic_processor/raw_value.cpp index a1a466cbe..1f9e6e92e 100644 --- a/src/common/interfaces/libpq/client_logic_processor/raw_value.cpp +++ b/src/common/interfaces/libpq/client_logic_processor/raw_value.cpp @@ -182,7 +182,11 @@ bool RawValue::process(const ICachedColumn *cached_column, char *err_msg) 1; /* the \0 is counted in the orignal PQescapeByteaCe function, so we need -1 */ } - return true; + if (!m_conn->client_logic->enable_client_encryption_log) { + return true; + } else { + return check_processed_data(err_msg); + } } void RawValue::inc_ref_count() @@ -195,3 +199,20 @@ void RawValue::dec_ref_count() Assert(ref_count > 0); ref_count--; } + + +bool RawValue::check_processed_data(char *err_msg) +{ + if (m_processed_data_size != 0 && m_processed_data_size < 12 && + !(m_processed_data_size == 2 && m_processed_data[0] == '\\' && m_processed_data[1] == 'x')) { + check_sprintf_s(sprintf_s(err_msg, MAX_ERRMSG_LENGTH, "invalid processed_data[%s].", m_processed_data)); + return false; + } + if (strcmp((char*)m_data_value, (char*)m_processed_data) == 0) { + check_sprintf_s(sprintf_s(err_msg, MAX_ERRMSG_LENGTH, "invalid processed_data[%s].", m_processed_data)); + return false; + } + return true; +} + + diff --git a/src/common/interfaces/libpq/client_logic_processor/raw_value.h b/src/common/interfaces/libpq/client_logic_processor/raw_value.h index f70489a5e..984beac3b 100644 --- a/src/common/interfaces/libpq/client_logic_processor/raw_value.h +++ b/src/common/interfaces/libpq/client_logic_processor/raw_value.h @@ -57,6 +57,7 @@ public: bool safe_to_delete() { return ref_count == 0; } + bool check_processed_data(char *err_msg); public: bool m_is_param; diff --git a/src/common/interfaces/libpq/client_logic_processor/stmt_processor.cpp b/src/common/interfaces/libpq/client_logic_processor/stmt_processor.cpp index 92dd294ea..ebfd73bf7 100755 --- a/src/common/interfaces/libpq/client_logic_processor/stmt_processor.cpp +++ b/src/common/interfaces/libpq/client_logic_processor/stmt_processor.cpp @@ -2143,6 +2143,17 @@ bool Processor::run_pre_query(StatementData *statement_data, bool is_inner_query return false; } } + + if (conn->client_logic->enable_client_encryption_log == true) { + foreach (stmt_iter, stmts) { + Node *stmt = (Node *)lfirst(stmt_iter); + size_t size = statement_data->conn->client_logic->rawValuesForReplace->size(); + if (IsA(stmt, InsertStmt) || IsA(stmt, CopyStmt)) { + printf("rawValuesForReplace size is %zu.\n", size); + } + } + } + statement_data->replace_raw_values(); if (!is_inner_query) { free_memory(); diff --git a/src/common/interfaces/libpq/fe-connect.cpp b/src/common/interfaces/libpq/fe-connect.cpp index 3b4d907e7..f0afcf1dc 100644 --- a/src/common/interfaces/libpq/fe-connect.cpp +++ b/src/common/interfaces/libpq/fe-connect.cpp @@ -983,6 +983,11 @@ static void fillPGconn(PGconn* conn, PQconninfoOption* connOptions) tmp = conninfo_getval(connOptions, "enable_ce"); if (tmp != NULL && strcmp(tmp, "1") == 0) { conn->client_logic->enable_client_encryption = true; + conn->client_logic->enable_client_encryption_log = false; + } else if (tmp != NULL && strcmp(tmp, "1_with_log") == 0) { + conn->client_logic->enable_client_encryption = true; + conn->client_logic->enable_client_encryption_log = true; + printf("enable_ce has been enabled.\n"); } else { conn->client_logic->enable_client_encryption = false; } diff --git a/src/include/libpq/cl_state.h b/src/include/libpq/cl_state.h index abad99fb6..b885ee90b 100644 --- a/src/include/libpq/cl_state.h +++ b/src/include/libpq/cl_state.h @@ -79,6 +79,7 @@ public: size_t get_rec_origial_ids_length(const Oid typid, const char* pname) const; PGconn* m_conn; bool enable_client_encryption; + bool enable_client_encryption_log; bool disable_once; PreparedStatementsList *preparedStatements; PreparedStatementsList *pendingStatements;