ClickHouse/base/poco/MongoDB/src/Database.cpp
2023-10-09 14:46:48 +00:00

483 lines
13 KiB
C++

//
// Database.cpp
//
// Library: MongoDB
// Package: MongoDB
// Module: Database
//
// Copyright (c) 2012, Applied Informatics Software Engineering GmbH.
// and Contributors.
//
// SPDX-License-Identifier: BSL-1.0
//
#include "Poco/MongoDB/Database.h"
#include "Poco/MongoDB/Binary.h"
#include "Poco/MD5Engine.h"
#include "Poco/SHA1Engine.h"
#include "Poco/PBKDF2Engine.h"
#include "Poco/HMACEngine.h"
#include "Poco/Base64Decoder.h"
#include "Poco/MemoryStream.h"
#include "Poco/StreamCopier.h"
#include "Poco/Exception.h"
#include "Poco/RandomStream.h"
#include "Poco/Random.h"
#include "Poco/Format.h"
#include "Poco/NumberParser.h"
#include <sstream>
#include <map>
namespace Poco {
namespace MongoDB {
const std::string Database::AUTH_MONGODB_CR("MONGODB-CR");
const std::string Database::AUTH_SCRAM_SHA1("SCRAM-SHA-1");
namespace
{
std::map<std::string, std::string> parseKeyValueList(const std::string& str)
{
std::map<std::string, std::string> kvm;
std::string::const_iterator it = str.begin();
std::string::const_iterator end = str.end();
while (it != end)
{
std::string k;
std::string v;
while (it != end && *it != '=') k += *it++;
if (it != end) ++it;
while (it != end && *it != ',') v += *it++;
if (it != end) ++it;
kvm[k] = v;
}
return kvm;
}
std::string decodeBase64(const std::string& base64)
{
Poco::MemoryInputStream istr(base64.data(), base64.size());
Poco::Base64Decoder decoder(istr);
std::string result;
Poco::StreamCopier::copyToString(decoder, result);
return result;
}
std::string encodeBase64(const std::string& data)
{
std::ostringstream ostr;
Poco::Base64Encoder encoder(ostr);
encoder.rdbuf()->setLineLength(0);
encoder << data;
encoder.close();
return ostr.str();
}
std::string digestToBinaryString(Poco::DigestEngine& engine)
{
Poco::DigestEngine::Digest d = engine.digest();
return std::string(reinterpret_cast<const char*>(&d[0]), d.size());
}
std::string digestToHexString(Poco::DigestEngine& engine)
{
Poco::DigestEngine::Digest d = engine.digest();
return Poco::DigestEngine::digestToHex(d);
}
std::string digestToBase64(Poco::DigestEngine& engine)
{
return encodeBase64(digestToBinaryString(engine));
}
std::string hashCredentials(const std::string& username, const std::string& password)
{
Poco::MD5Engine md5;
md5.update(username);
md5.update(std::string(":mongo:"));
md5.update(password);
return digestToHexString(md5);
}
std::string createNonce()
{
Poco::MD5Engine md5;
Poco::RandomInputStream randomStream;
Poco::Random random;
for (int i = 0; i < 4; i++)
{
md5.update(randomStream.get());
md5.update(random.nextChar());
}
return digestToHexString(md5);
}
}
Database::Database(const std::string& db):
_dbname(db)
{
}
Database::~Database()
{
}
bool Database::authenticate(Connection& connection, const std::string& username, const std::string& password, const std::string& method)
{
if (username.empty()) throw Poco::InvalidArgumentException("empty username");
if (password.empty()) throw Poco::InvalidArgumentException("empty password");
if (method == AUTH_MONGODB_CR)
return authCR(connection, username, password);
else if (method == AUTH_SCRAM_SHA1)
return authSCRAM(connection, username, password);
else
throw Poco::InvalidArgumentException("authentication method", method);
}
bool Database::authCR(Connection& connection, const std::string& username, const std::string& password)
{
std::string nonce;
Poco::SharedPtr<QueryRequest> pCommand = createCommand();
pCommand->selector().add<Poco::Int32>("getnonce", 1);
ResponseMessage response;
connection.sendRequest(*pCommand, response);
if (response.documents().size() > 0)
{
Document::Ptr pDoc = response.documents()[0];
if (pDoc->getInteger("ok") != 1) return false;
nonce = pDoc->get<std::string>("nonce", "");
if (nonce.empty()) throw Poco::ProtocolException("no nonce received");
}
else throw Poco::ProtocolException("empty response for getnonce");
std::string credsDigest = hashCredentials(username, password);
Poco::MD5Engine md5;
md5.update(nonce);
md5.update(username);
md5.update(credsDigest);
std::string key = digestToHexString(md5);
pCommand = createCommand();
pCommand->selector()
.add<Poco::Int32>("authenticate", 1)
.add<std::string>("user", username)
.add<std::string>("nonce", nonce)
.add<std::string>("key", key);
connection.sendRequest(*pCommand, response);
if (response.documents().size() > 0)
{
Document::Ptr pDoc = response.documents()[0];
return pDoc->getInteger("ok") == 1;
}
else throw Poco::ProtocolException("empty response for authenticate");
}
bool Database::authSCRAM(Connection& connection, const std::string& username, const std::string& password)
{
std::string clientNonce(createNonce());
std::string clientFirstMsg = Poco::format("n=%s,r=%s", username, clientNonce);
Poco::SharedPtr<QueryRequest> pCommand = createCommand();
pCommand->selector()
.add<Poco::Int32>("saslStart", 1)
.add<std::string>("mechanism", AUTH_SCRAM_SHA1)
.add<Binary::Ptr>("payload", new Binary(Poco::format("n,,%s", clientFirstMsg)));
ResponseMessage response;
connection.sendRequest(*pCommand, response);
Int32 conversationId = 0;
std::string serverFirstMsg;
if (response.documents().size() > 0)
{
Document::Ptr pDoc = response.documents()[0];
if (pDoc->getInteger("ok") == 1)
{
Binary::Ptr pPayload = pDoc->get<Binary::Ptr>("payload");
serverFirstMsg = pPayload->toRawString();
conversationId = pDoc->get<Int32>("conversationId");
}
else
{
if (pDoc->exists("errmsg"))
{
const Poco::MongoDB::Element::Ptr value = pDoc->get("errmsg");
auto message = static_cast<const Poco::MongoDB::ConcreteElement<std::string> &>(*value).value();
throw Poco::RuntimeException(message);
}
else
return false;
}
}
else throw Poco::ProtocolException("empty response for saslStart");
std::map<std::string, std::string> kvm = parseKeyValueList(serverFirstMsg);
const std::string serverNonce = kvm["r"];
const std::string salt = decodeBase64(kvm["s"]);
const unsigned iterations = Poco::NumberParser::parseUnsigned(kvm["i"]);
const Poco::UInt32 dkLen = 20;
std::string hashedPassword = hashCredentials(username, password);
Poco::PBKDF2Engine<Poco::HMACEngine<Poco::SHA1Engine> > pbkdf2(salt, iterations, dkLen);
pbkdf2.update(hashedPassword);
std::string saltedPassword = digestToBinaryString(pbkdf2);
std::string clientFinalNoProof = Poco::format("c=biws,r=%s", serverNonce);
std::string authMessage = Poco::format("%s,%s,%s", clientFirstMsg, serverFirstMsg, clientFinalNoProof);
Poco::HMACEngine<Poco::SHA1Engine> hmacKey(saltedPassword);
hmacKey.update(std::string("Client Key"));
std::string clientKey = digestToBinaryString(hmacKey);
Poco::SHA1Engine sha1;
sha1.update(clientKey);
std::string storedKey = digestToBinaryString(sha1);
Poco::HMACEngine<Poco::SHA1Engine> hmacSig(storedKey);
hmacSig.update(authMessage);
std::string clientSignature = digestToBinaryString(hmacSig);
std::string clientProof(clientKey);
for (std::size_t i = 0; i < clientProof.size(); i++)
{
clientProof[i] ^= clientSignature[i];
}
std::string clientFinal = Poco::format("%s,p=%s", clientFinalNoProof, encodeBase64(clientProof));
pCommand = createCommand();
pCommand->selector()
.add<Poco::Int32>("saslContinue", 1)
.add<Poco::Int32>("conversationId", conversationId)
.add<Binary::Ptr>("payload", new Binary(clientFinal));
std::string serverSecondMsg;
connection.sendRequest(*pCommand, response);
if (response.documents().size() > 0)
{
Document::Ptr pDoc = response.documents()[0];
if (pDoc->getInteger("ok") == 1)
{
Binary::Ptr pPayload = pDoc->get<Binary::Ptr>("payload");
serverSecondMsg = pPayload->toRawString();
}
else
{
if (pDoc->exists("errmsg"))
{
const Poco::MongoDB::Element::Ptr value = pDoc->get("errmsg");
auto message = static_cast<const Poco::MongoDB::ConcreteElement<std::string> &>(*value).value();
throw Poco::RuntimeException(message);
}
else
return false;
}
}
else throw Poco::ProtocolException("empty response for saslContinue");
Poco::HMACEngine<Poco::SHA1Engine> hmacSKey(saltedPassword);
hmacSKey.update(std::string("Server Key"));
std::string serverKey = digestToBinaryString(hmacSKey);
Poco::HMACEngine<Poco::SHA1Engine> hmacSSig(serverKey);
hmacSSig.update(authMessage);
std::string serverSignature = digestToBase64(hmacSSig);
kvm = parseKeyValueList(serverSecondMsg);
std::string serverSignatureReceived = kvm["v"];
if (serverSignature != serverSignatureReceived)
throw Poco::ProtocolException("server signature verification failed");
pCommand = createCommand();
pCommand->selector()
.add<Poco::Int32>("saslContinue", 1)
.add<Poco::Int32>("conversationId", conversationId)
.add<Binary::Ptr>("payload", new Binary);
connection.sendRequest(*pCommand, response);
if (response.documents().size() > 0)
{
Document::Ptr pDoc = response.documents()[0];
if (pDoc->getInteger("ok") == 1)
{
return true;
}
else
{
if (pDoc->exists("errmsg"))
{
const Poco::MongoDB::Element::Ptr value = pDoc->get("errmsg");
auto message = static_cast<const Poco::MongoDB::ConcreteElement<std::string> &>(*value).value();
throw Poco::RuntimeException(message);
}
else
return false;
}
}
else throw Poco::ProtocolException("empty response for saslContinue");
}
Document::Ptr Database::queryBuildInfo(Connection& connection) const
{
// build info can be issued on "config" system database
Poco::SharedPtr<Poco::MongoDB::QueryRequest> request = createCommand();
request->selector().add("buildInfo", 1);
Poco::MongoDB::ResponseMessage response;
connection.sendRequest(*request, response);
Document::Ptr buildInfo;
if ( response.documents().size() > 0 )
{
buildInfo = response.documents()[0];
}
else
{
throw Poco::ProtocolException("Didn't get a response from the buildinfo command");
}
return buildInfo;
}
Document::Ptr Database::queryServerHello(Connection& connection, bool old) const
{
// hello can be issued on "config" system database
Poco::SharedPtr<Poco::MongoDB::QueryRequest> request = createCommand();
// 'hello' command was previously called 'isMaster'
std::string command_name;
if (old)
command_name = "isMaster";
else
command_name = "hello";
request->selector().add(command_name, 1);
Poco::MongoDB::ResponseMessage response;
connection.sendRequest(*request, response);
Document::Ptr hello;
if ( response.documents().size() > 0 )
{
hello = response.documents()[0];
}
else
{
throw Poco::ProtocolException("Didn't get a response from the hello command");
}
return hello;
}
Int64 Database::count(Connection& connection, const std::string& collectionName) const
{
Poco::SharedPtr<Poco::MongoDB::QueryRequest> countRequest = createCountRequest(collectionName);
Poco::MongoDB::ResponseMessage response;
connection.sendRequest(*countRequest, response);
if (response.documents().size() > 0)
{
Poco::MongoDB::Document::Ptr doc = response.documents()[0];
return doc->getInteger("n");
}
return -1;
}
Poco::MongoDB::Document::Ptr Database::ensureIndex(Connection& connection, const std::string& collection, const std::string& indexName, Poco::MongoDB::Document::Ptr keys, bool unique, bool background, int version, int ttl)
{
Poco::MongoDB::Document::Ptr index = new Poco::MongoDB::Document();
index->add("ns", _dbname + "." + collection);
index->add("name", indexName);
index->add("key", keys);
if (version > 0)
{
index->add("version", version);
}
if (unique)
{
index->add("unique", true);
}
if (background)
{
index->add("background", true);
}
if (ttl > 0)
{
index->add("expireAfterSeconds", ttl);
}
Poco::SharedPtr<Poco::MongoDB::InsertRequest> insertRequest = createInsertRequest("system.indexes");
insertRequest->documents().push_back(index);
connection.sendRequest(*insertRequest);
return getLastErrorDoc(connection);
}
Document::Ptr Database::getLastErrorDoc(Connection& connection) const
{
Document::Ptr errorDoc;
Poco::SharedPtr<Poco::MongoDB::QueryRequest> request = createCommand();
request->setNumberToReturn(1);
request->selector().add("getLastError", 1);
Poco::MongoDB::ResponseMessage response;
connection.sendRequest(*request, response);
if (response.documents().size() > 0)
{
errorDoc = response.documents()[0];
}
return errorDoc;
}
std::string Database::getLastError(Connection& connection) const
{
Document::Ptr errorDoc = getLastErrorDoc(connection);
if (!errorDoc.isNull() && errorDoc->isType<std::string>("err"))
{
return errorDoc->get<std::string>("err");
}
return "";
}
Poco::SharedPtr<Poco::MongoDB::QueryRequest> Database::createCountRequest(const std::string& collectionName) const
{
Poco::SharedPtr<Poco::MongoDB::QueryRequest> request = createCommand();
request->setNumberToReturn(1);
request->selector().add("count", collectionName);
return request;
}
} } // namespace Poco::MongoDB