Support JWT in clickhouse-client

This commit is contained in:
Konstantin Bogdanov 2024-04-22 03:36:34 +02:00
parent ef34ec874a
commit f2583db8a1
Signed by: thevar1able
GPG Key ID: DB399448D9FE52F1
12 changed files with 110 additions and 175 deletions

View File

@ -944,6 +944,7 @@ void Client::addOptions(OptionsDescription & options_description)
("ssh-key-file", po::value<std::string>(), "File containing the SSH private key for authenticate with the server.")
("ssh-key-passphrase", po::value<std::string>(), "Passphrase for the SSH private key specified by --ssh-key-file.")
("quota_key", po::value<std::string>(), "A string to differentiate quotas when the user have keyed quotas configured on server")
("jwt", po::value<std::string>(), "Use JWT for authentication")
("max_client_network_bandwidth", po::value<int>(), "the maximum speed of data exchange over the network for the client in bytes per second.")
("compression", po::value<bool>(), "enable or disable compression (enabled by default for remote communication and disabled for localhost communication).")
@ -1102,6 +1103,11 @@ void Client::processOptions(const OptionsDescription & options_description,
config().setBool("no-warnings", true);
if (options.count("fake-drop"))
config().setString("ignore_drop_queries_probability", "1");
if (options.count("jwt"))
{
config().setString("jwt", options["jwt"].as<std::string>());
config().setString("user", "");
}
if (options.count("accept-invalid-certificate"))
{
config().setString("openSSL.client.invalidCertificateHandler.name", "AcceptCertificateHandler");

View File

@ -109,6 +109,7 @@ namespace ErrorCodes
extern const int USER_SESSION_LIMIT_EXCEEDED;
extern const int NOT_IMPLEMENTED;
extern const int CANNOT_READ_FROM_FILE_DESCRIPTOR;
extern const int USER_EXPIRED;
}
}
@ -916,7 +917,7 @@ bool ClientBase::isSyncInsertWithData(const ASTInsertQuery & insert_query, const
return !settings.async_insert;
}
void ClientBase::processTextAsSingleQuery(const String & full_query)
bool ClientBase::processTextAsSingleQuery(const String & full_query)
{
/// Some parts of a query (result output and formatting) are executed
/// client-side. Thus we need to parse the query.
@ -928,7 +929,7 @@ void ClientBase::processTextAsSingleQuery(const String & full_query)
ignore_error);
if (!parsed_query)
return;
return is_interactive;
String query_to_execute;
@ -952,9 +953,10 @@ void ClientBase::processTextAsSingleQuery(const String & full_query)
else
query_to_execute = full_query;
bool continue_repl = is_interactive;
try
{
processParsedSingleQuery(full_query, query_to_execute, parsed_query, echo_queries);
continue_repl = processParsedSingleQuery(full_query, query_to_execute, parsed_query, echo_queries);
}
catch (Exception & e)
{
@ -965,6 +967,8 @@ void ClientBase::processTextAsSingleQuery(const String & full_query)
if (have_error)
processError(full_query);
return continue_repl;
}
void ClientBase::processOrdinaryQuery(const String & query_to_execute, ASTPtr parsed_query)
@ -1867,7 +1871,7 @@ void ClientBase::cancelQuery()
cancelled = true;
}
void ClientBase::processParsedSingleQuery(const String & full_query, const String & query_to_execute,
bool ClientBase::processParsedSingleQuery(const String & full_query, const String & query_to_execute,
ASTPtr parsed_query, std::optional<bool> echo_query_, bool report_error)
{
resetOutput();
@ -2030,6 +2034,15 @@ void ClientBase::processParsedSingleQuery(const String & full_query, const Strin
connection->setDefaultDatabase(new_database);
}
}
else
{
if (server_exception && server_exception->code() == ErrorCodes::USER_EXPIRED)
{
if (report_error)
processError(full_query);
return false;
}
}
/// Always print last block (if it was not printed already)
if (profile_events.last_block)
@ -2064,6 +2077,8 @@ void ClientBase::processParsedSingleQuery(const String & full_query, const Strin
if (have_error && report_error)
processError(full_query);
return is_interactive;
}
@ -2263,14 +2278,15 @@ bool ClientBase::executeMultiQuery(const String & all_queries_text)
// Echo all queries if asked; makes for a more readable reference file.
echo_query = test_hint.echoQueries().value_or(echo_query);
bool continue_repl = is_interactive;
try
{
processParsedSingleQuery(full_query, query_to_execute, parsed_query, echo_query, false);
continue_repl = processParsedSingleQuery(full_query, query_to_execute, parsed_query, echo_query, false);
}
catch (...)
{
// Surprisingly, this is a client error. A server error would
// have been reported without throwing (see onReceiveSeverException()).
// have been reported without throwing (see onReceiveExceptionFromServer()).
client_exception = std::make_unique<Exception>(getCurrentExceptionMessageAndPattern(print_stack_trace), getCurrentExceptionCode());
have_error = true;
}
@ -2370,7 +2386,7 @@ bool ClientBase::executeMultiQuery(const String & all_queries_text)
// Stop processing queries if needed.
if (have_error && !ignore_error)
return is_interactive;
return continue_repl;
this_query_begin = this_query_end;
break;
@ -2400,9 +2416,7 @@ bool ClientBase::processQueryText(const String & text)
if (!is_multiquery)
{
assert(!query_fuzzer_runs);
processTextAsSingleQuery(text);
return true;
return processTextAsSingleQuery(text);
}
if (query_fuzzer_runs)

View File

@ -94,8 +94,8 @@ protected:
void processOrdinaryQuery(const String & query_to_execute, ASTPtr parsed_query);
void processInsertQuery(const String & query_to_execute, ASTPtr parsed_query);
void processTextAsSingleQuery(const String & full_query);
void processParsedSingleQuery(const String & full_query, const String & query_to_execute,
bool processTextAsSingleQuery(const String & full_query);
bool processParsedSingleQuery(const String & full_query, const String & query_to_execute,
ASTPtr parsed_query, std::optional<bool> echo_query_ = {}, bool report_error = false);
static void adjustQueryEnd(const char *& this_query_end, const char * all_queries_end, uint32_t max_parser_depth, uint32_t max_parser_backtracks);

View File

@ -74,6 +74,7 @@ Connection::Connection(const String & host_, UInt16 port_,
const String & default_database_,
const String & user_, const String & password_,
[[maybe_unused]] const SSHKey & ssh_private_key_,
const String & jwt_,
const String & quota_key_,
const String & cluster_,
const String & cluster_secret_,
@ -86,6 +87,7 @@ Connection::Connection(const String & host_, UInt16 port_,
, ssh_private_key(ssh_private_key_)
#endif
, quota_key(quota_key_)
, jwt(jwt_)
, cluster(cluster_)
, cluster_secret(cluster_secret_)
, client_name(client_name_)
@ -341,6 +343,11 @@ void Connection::sendHello()
performHandshakeForSSHAuth();
}
#endif
else if (!jwt.empty())
{
writeStringBinary(EncodedUserInfo::JWT_AUTHENTICAION_MARKER, *out);
writeStringBinary(jwt, *out);
}
else
{
writeStringBinary(user, *out);
@ -1310,6 +1317,7 @@ ServerConnectionPtr Connection::createConnection(const ConnectionParameters & pa
parameters.user,
parameters.password,
parameters.ssh_private_key,
parameters.jwt,
parameters.quota_key,
"", /* cluster */
"", /* cluster_secret */

View File

@ -53,6 +53,7 @@ public:
const String & default_database_,
const String & user_, const String & password_,
const SSHKey & ssh_private_key_,
const String & jwt_,
const String & quota_key_,
const String & cluster_,
const String & cluster_secret_,
@ -173,6 +174,7 @@ private:
SSHKey ssh_private_key;
#endif
String quota_key;
String jwt;
/// For inter-server authorization
String cluster;

View File

@ -54,26 +54,33 @@ ConnectionParameters::ConnectionParameters(const Poco::Util::AbstractConfigurati
if (!config.has("ssh-key-file"))
{
bool password_prompt = false;
if (config.getBool("ask-password", false))
if (config.has("jwt"))
{
if (config.has("password"))
throw Exception(ErrorCodes::BAD_ARGUMENTS, "Specified both --password and --ask-password. Remove one of them");
password_prompt = true;
jwt = config.getString("jwt");
}
else
{
password = config.getString("password", "");
/// if the value of --password is omitted, the password will be set implicitly to "\n"
if (password == ASK_PASSWORD)
bool password_prompt = false;
if (config.getBool("ask-password", false))
{
if (config.has("password"))
throw Exception(ErrorCodes::BAD_ARGUMENTS, "Specified both --password and --ask-password. Remove one of them");
password_prompt = true;
}
if (password_prompt)
{
std::string prompt{"Password for user (" + user + "): "};
char buf[1000] = {};
if (auto * result = readpassphrase(prompt.c_str(), buf, sizeof(buf), 0))
password = result;
}
else
{
password = config.getString("password", "");
/// if the value of --password is omitted, the password will be set implicitly to "\n"
if (password == ASK_PASSWORD)
password_prompt = true;
}
if (password_prompt)
{
std::string prompt{"Password for user (" + user + "): "};
char buf[1000] = {};
if (auto * result = readpassphrase(prompt.c_str(), buf, sizeof(buf), 0))
password = result;
}
}
}
else

View File

@ -22,6 +22,7 @@ struct ConnectionParameters
std::string password;
std::string quota_key;
SSHKey ssh_private_key;
std::string jwt;
Protocol::Secure security = Protocol::Secure::Disable;
Protocol::Compression compression = Protocol::Compression::Enable;
ConnectionTimeouts timeouts;
@ -30,7 +31,7 @@ struct ConnectionParameters
ConnectionParameters(const Poco::Util::AbstractConfiguration & config, std::string host);
ConnectionParameters(const Poco::Util::AbstractConfiguration & config, std::string host, std::optional<UInt16> port);
static UInt16 getPortFromConfig(const Poco::Util::AbstractConfiguration & config, std::string connection_host);
static UInt16 getPortFromConfig(const Poco::Util::AbstractConfiguration & config, String connection_host);
/// Ask to enter the user's password if password option contains this value.
/// "\n" is used because there is hardly a chance that a user would use '\n' as password.

View File

@ -123,7 +123,7 @@ protected:
{
return std::make_shared<Connection>(
host, port,
default_database, user, password, SSHKey(), quota_key,
default_database, user, password, SSHKey(), "", quota_key,
cluster, cluster_secret,
client_name, compression, secure);
}

View File

@ -62,6 +62,7 @@ const char USER_INTERSERVER_MARKER[] = " INTERSERVER SECRET ";
/// Marker for SSH-keys-based authentication (passed as the user name)
const char SSH_KEY_AUTHENTICAION_MARKER[] = " SSH KEY AUTHENTICATION ";
const char JWT_AUTHENTICAION_MARKER[] = " JWT AUTHENTICATION ";
};

View File

@ -90,6 +90,7 @@ message QueryInfo {
string user_name = 9;
string password = 10;
string quota = 11;
string jwt = 25;
// Works exactly like sessions in the HTTP protocol.
string session_id = 12;

View File

@ -5656,7 +5656,7 @@ std::optional<QueryPipeline> StorageReplicatedMergeTree::distributedWriteFromClu
{
auto connection = std::make_shared<Connection>(
node.host_name, node.port, query_context->getGlobalContext()->getCurrentDatabase(),
node.user, node.password, SSHKey(), node.quota_key, node.cluster, node.cluster_secret,
node.user, node.password, SSHKey(), "", node.quota_key, node.cluster, node.cluster_secret,
"ParallelInsertSelectInititiator",
node.compression,
node.secure

File diff suppressed because one or more lines are too long