From cd8cce6dcb3d63bbf4bf882379540ed8b719dfc4 Mon Sep 17 00:00:00 2001
From: Sara Golemon <sara.golemon@mongodb.com>
Date: Fri, 28 Aug 2020 17:02:46 +0000
Subject: SERVER-50605 Add logMessage test-only command

(cherry picked from commit cbdf4deaa4ef4352750893ab0b4b276b86e3026f)
---
 src/mongo/db/commands/generic.cpp               | 47 ++++++++++++++++++
 src/mongo/db/initialize_server_global_state.cpp |  1 +
 src/mongo/logger/logstream_builder.cpp          |  2 +
 src/mongo/logger/logstream_builder.h            | 66 ++++++++++++++++++++-----
 4 files changed, 104 insertions(+), 12 deletions(-)

(limited to 'src')

diff --git a/src/mongo/db/commands/generic.cpp b/src/mongo/db/commands/generic.cpp
index c1312768ee1..af647671660 100644
--- a/src/mongo/db/commands/generic.cpp
+++ b/src/mongo/db/commands/generic.cpp
@@ -500,6 +500,53 @@ public:
 MONGO_FP_DECLARE(crashOnShutdown);
 int* volatile illegalAddress;  // NOLINT - used for fail point only
 
+class CmdLogMessage : public BasicCommand {
+public:
+    CmdLogMessage() : BasicCommand("logMessage") {}
+
+    void help(stringstream& help) const final {
+        help << "Send a message to the server log";
+    }
+
+    bool supportsWriteConcern(const BSONObj& cmd) const final {
+        return false;
+    }
+
+    bool slaveOk() const final {
+        return true;
+    }
+
+    bool adminOnly() const final {
+        return true;
+    }
+
+    void addRequiredPrivileges(const std::string& dbname,
+                               const BSONObj& cmdObj,
+                               std::vector<Privilege>* out) final {
+        out->push_back(
+            Privilege(ResourcePattern::forClusterResource(), ActionType::applicationMessage));
+    }
+
+    bool run(OperationContext* opCtx,
+             const std::string& ns,
+             const BSONObj& cmdObj,
+             BSONObjBuilder& result) final {
+        auto msgElem = cmdObj["logMessage"];
+        uassert(ErrorCodes::BadValue, "logMessage must be a string", msgElem.type() == String);
+
+        log() << "logMessage: " << msgElem.valueStringData();
+        return true;
+    }
+};
+
+MONGO_INITIALIZER(RegisterCmdLogMessage)(InitializerContext* context) {
+    if (Command::testCommandsEnabled) {
+        // Leaked intentionally: a Command registers itself when constructed.
+        new CmdLogMessage();
+    }
+    return Status::OK();
+}
+
 }  // namespace
 
 void CmdShutdown::addRequiredPrivileges(const std::string& dbname,
diff --git a/src/mongo/db/initialize_server_global_state.cpp b/src/mongo/db/initialize_server_global_state.cpp
index ac6e6580bed..56d9830f54d 100644
--- a/src/mongo/db/initialize_server_global_state.cpp
+++ b/src/mongo/db/initialize_server_global_state.cpp
@@ -316,6 +316,7 @@ MONGO_INITIALIZER_GENERAL(
 
     logger::globalLogDomain()->attachAppender(
         logger::MessageLogDomain::AppenderAutoPtr(new RamLogAppender(RamLog::get("global"))));
+    logger::LogstreamBuilder::setNewlineEscape();
 
     return Status::OK();
 }
diff --git a/src/mongo/logger/logstream_builder.cpp b/src/mongo/logger/logstream_builder.cpp
index 201130579e9..77ed984a9b3 100644
--- a/src/mongo/logger/logstream_builder.cpp
+++ b/src/mongo/logger/logstream_builder.cpp
@@ -71,6 +71,8 @@ struct ThreadOstreamCacheFinalizer {
 
 namespace logger {
 
+bool LogstreamBuilder::newlineEscape = false;
+
 LogstreamBuilder::LogstreamBuilder(MessageLogDomain* domain,
                                    StringData contextName,
                                    LogSeverity severity)
diff --git a/src/mongo/logger/logstream_builder.h b/src/mongo/logger/logstream_builder.h
index 2185ed1661d..62269a50783 100644
--- a/src/mongo/logger/logstream_builder.h
+++ b/src/mongo/logger/logstream_builder.h
@@ -124,24 +124,19 @@ public:
     }
 
     LogstreamBuilder& operator<<(const char* x) {
-        stream() << x;
-        return *this;
+        return appendEscapedString(x);
     }
     LogstreamBuilder& operator<<(const std::string& x) {
-        stream() << x;
-        return *this;
+        return appendEscapedString(x);
     }
     LogstreamBuilder& operator<<(StringData x) {
-        stream() << x;
-        return *this;
+        return appendEscapedString(x);
     }
     LogstreamBuilder& operator<<(char* x) {
-        stream() << x;
-        return *this;
+        return appendEscapedString(x);
     }
     LogstreamBuilder& operator<<(char x) {
-        stream() << x;
-        return *this;
+        return appendEscapedChar(x);
     }
     LogstreamBuilder& operator<<(int x) {
         stream() << x;
@@ -210,8 +205,7 @@ public:
 
     template <typename T>
     LogstreamBuilder& operator<<(const T& x) {
-        stream() << x.toString();
-        return *this;
+        return appendEscapedString(x.toString());
     }
 
     LogstreamBuilder& operator<<(std::ostream& (*manip)(std::ostream&)) {
@@ -229,9 +223,57 @@ public:
      */
     void operator<<(Tee* tee);
 
+    /**
+     * Enable newilne escaping behavior.
+     */
+    static void setNewlineEscape() {
+        newlineEscape = true;
+    }
+
 private:
+    /**
+     * Global option to escape newlines.
+     */
+    static bool newlineEscape;
+
     void makeStream();
 
+    // Append a string explicitly escaping \r and \n only.
+    LogstreamBuilder& appendEscapedString(StringData x) {
+        if (newlineEscape) {
+            auto idx = std::min(x.find('\r'), x.find('\n'));
+            while (idx != std::string::npos) {
+                if (idx > 0) {
+                    stream() << x.substr(0, idx);
+                }
+                appendEscapedChar(x[idx]);
+                x = x.substr(idx + 1);
+                idx = std::min(x.find('\r'), x.find('\n'));
+            }
+        }
+
+        stream() << x;
+        return *this;
+    }
+
+    LogstreamBuilder& appendEscapedChar(char ch) {
+        if (newlineEscape) {
+            switch (ch) {
+                case '\n':
+                    stream() << "\\n";
+                    break;
+                case '\r':
+                    stream() << "\\r";
+                    break;
+                default:
+                    stream() << ch;
+            }
+        } else {
+            stream() << ch;
+        }
+        return *this;
+    }
+
     MessageLogDomain* _domain;
     std::string _contextName;
     LogSeverity _severity;
-- 
cgit v1.2.1