Allow specifying an exit code for the 'quit' command

Summary:
This patch adds the possibility to specify an exit code when calling quit.
We accept any int, even though it depends on the user what happens if the int is
out of the range of what the operating system supports as exit codes.

Fixes rdar://problem/38452312

Reviewers: davide, jingham, clayborg

Reviewed By: jingham

Subscribers: clayborg, jingham, lldb-commits

Differential Revision: https://reviews.llvm.org/D48659

llvm-svn: 336824

lldb/include/lldb/API/SBCommandInterpreter.h

@@ -207,6 +207,25 @@ public:
 
   void SetPromptOnQuit(bool b);
 
+  //----------------------------------------------------------------------
+  /// Sets whether the command interpreter should allow custom exit codes
+  /// for the 'quit' command.
+  //----------------------------------------------------------------------
+  void AllowExitCodeOnQuit(bool allow);
+
+  //----------------------------------------------------------------------
+  /// Returns true if the user has called the 'quit' command with a custom exit
+  /// code.
+  //----------------------------------------------------------------------
+  bool HasCustomQuitExitCode();
+
+  //----------------------------------------------------------------------
+  /// Returns the exit code that the user has specified when running the
+  /// 'quit' command. Returns 0 if the user hasn't called 'quit' at all or
+  /// without a custom exit code.
+  //----------------------------------------------------------------------
+  int GetQuitStatus();
+
   //----------------------------------------------------------------------
   /// Resolve the command just as HandleCommand would, expanding abbreviations
   /// and aliases.  If successful, result->GetOutput has the full expansion.

lldb/include/lldb/Interpreter/CommandInterpreter.h

@@ -455,6 +455,30 @@ public:
 
   void SetPromptOnQuit(bool b);
 
+  //------------------------------------------------------------------
+  /// Specify if the command interpreter should allow that the user can
+  /// specify a custom exit code when calling 'quit'.
+  //------------------------------------------------------------------
+  void AllowExitCodeOnQuit(bool allow);
+
+  //------------------------------------------------------------------
+  /// Sets the exit code for the quit command.
+  /// @param[in] exit_code
+  ///     The exit code that the driver should return on exit.
+  /// @return True if the exit code was successfully set; false if the
+  ///         interpreter doesn't allow custom exit codes.
+  /// @see AllowExitCodeOnQuit
+  //------------------------------------------------------------------
+  LLVM_NODISCARD bool SetQuitExitCode(int exit_code);
+
+  //------------------------------------------------------------------
+  /// Returns the exit code that the user has specified when running the
+  /// 'quit' command.
+  /// @param[out] exited
+  ///     Set to true if the user has called quit with a custom exit code.
+  //------------------------------------------------------------------
+  int GetQuitExitCode(bool &exited) const;
+
   void ResolveCommand(const char *command_line, CommandReturnObject &result);
 
   bool GetStopCmdSourceOnError() const;
@@ -558,6 +582,12 @@ private:
   uint32_t m_num_errors;
   bool m_quit_requested;
   bool m_stopped_for_crash;
+
+  // The exit code the user has requested when calling the 'quit' command.
+  // No value means the user hasn't set a custom exit code so far.
+  llvm::Optional<int> m_quit_exit_code;
+  // If the driver is accepts custom exit codes for the 'quit' command.
+  bool m_allow_exit_code = false;
 };
 
 } // namespace lldb_private

lldb/lit/Quit/TestQuitExitCode-30.test

@@ -0,0 +1,3 @@
+# UNSUPPORTED: windows
+# RUN: python %S/expect_exit_code.py 226 %lldb -b -s %s
+q -30

lldb/lit/Quit/TestQuitExitCode0.test

@@ -0,0 +1,3 @@
+# UNSUPPORTED: windows
+# RUN: %lldb -b -s %s
+q 0

lldb/lit/Quit/TestQuitExitCode30.test

@@ -0,0 +1,3 @@
+# UNSUPPORTED: windows
+# RUN: python %S/expect_exit_code.py 30  %lldb -b -s %s
+q 30

lldb/lit/Quit/TestQuitExitCodeHex0.test

@@ -0,0 +1,3 @@
+# UNSUPPORTED: windows
+# RUN: %lldb -b -s %s
+q 0x0

lldb/lit/Quit/TestQuitExitCodeHexA.test

@@ -0,0 +1,3 @@
+# UNSUPPORTED: windows
+# RUN: python %S/expect_exit_code.py 10 %lldb -b -s %s
+q 0xA

lldb/lit/Quit/TestQuitExitCodeImplicit0.test

@@ -0,0 +1,3 @@
+# UNSUPPORTED: windows
+# RUN: %lldb -b -s %s
+q

lldb/lit/Quit/TestQuitExitCodeNonInt.test

@@ -0,0 +1,4 @@
+# UNSUPPORTED: windows
+# RUN: %lldb -b -s %s 2>&1 | FileCheck %s
+q str
+// CHECK: Couldn't parse 'str'

lldb/lit/Quit/TestQuitExitCodeTooManyArgs.test

@@ -0,0 +1,4 @@
+# UNSUPPORTED: windows
+# RUN: %lldb -b -s %s 2>&1 | FileCheck %s
+q 1 2
+// CHECK: Too many arguments for 'quit'

lldb/lit/Quit/expect_exit_code.py

@@ -0,0 +1,16 @@
+#!/usr/bin/env python2
+
+import subprocess
+import sys
+
+args = sys.argv
+
+expected_exit_code = args[1]
+
+args = args[2:]
+print("Running " + (" ".join(args)))
+real_exit_code = subprocess.call(args)
+
+if str(real_exit_code) != expected_exit_code:
+  print("Got exit code %d but expected %s" % (real_exit_code, expected_exit_code))
+  exit(1)

lldb/lit/Quit/lit.local.cfg

@@ -0,0 +1 @@
+config.suffixes = ['.test']

lldb/packages/Python/lldbsuite/test/quit/TestQuit.py

@@ -0,0 +1,32 @@
+"""
+Test lldb's quit command.
+"""
+
+from __future__ import print_function
+
+import lldb
+from lldbsuite.test.decorators import *
+from lldbsuite.test.lldbtest import *
+from lldbsuite.test import lldbutil
+
+
+class QuitCommandTestCase(TestBase):
+
+    mydir = TestBase.compute_mydir(__file__)
+
+    @no_debug_info_test
+    def test_quit_exit_code_disallow(self):
+        self.ci.AllowExitCodeOnQuit(False)
+        self.expect(
+            "quit 20",
+            substrs=[
+                "error: The current driver doesn't allow custom exit codes for the quit command"],
+            error=True)
+        self.assertFalse(self.ci.HasCustomQuitExitCode())
+
+    @no_debug_info_test
+    def test_quit_exit_code_allow(self):
+        self.ci.AllowExitCodeOnQuit(True)
+        self.runCmd("quit 10", check=False)
+        self.assertTrue(self.ci.HasCustomQuitExitCode())
+        self.assertEqual(self.ci.GetQuitStatus(), 10)

lldb/scripts/interface/SBCommandInterpreter.i

@@ -155,6 +155,15 @@ public:
     void
     SetPromptOnQuit(bool b);
 
+    void
+    AllowExitCodeOnQuit(bool b);
+
+    bool
+    HasCustomQuitExitCode();
+
+    int
+    GetQuitStatus();
+
     void
     ResolveCommand(const char *command_line, SBCommandReturnObject &result);
 

lldb/source/API/SBCommandInterpreter.cpp

@@ -379,6 +379,23 @@ void SBCommandInterpreter::SetPromptOnQuit(bool b) {
     m_opaque_ptr->SetPromptOnQuit(b);
 }
 
+void SBCommandInterpreter::AllowExitCodeOnQuit(bool allow) {
+  if (m_opaque_ptr)
+    m_opaque_ptr->AllowExitCodeOnQuit(allow);
+}
+
+bool SBCommandInterpreter::HasCustomQuitExitCode() {
+  bool exited = false;
+  if (m_opaque_ptr)
+    m_opaque_ptr->GetQuitExitCode(exited);
+  return exited;
+}
+
+int SBCommandInterpreter::GetQuitStatus() {
+  bool exited = false;
+  return (m_opaque_ptr ? m_opaque_ptr->GetQuitExitCode(exited) : 0);
+}
+
 void SBCommandInterpreter::ResolveCommand(const char *command_line,
                                           SBCommandReturnObject &result) {
   result.Clear();

lldb/source/Commands/CommandObjectQuit.cpp

@@ -16,6 +16,7 @@
 #include "lldb/Interpreter/CommandInterpreter.h"
 #include "lldb/Interpreter/CommandReturnObject.h"
 #include "lldb/Target/Process.h"
+#include "lldb/Utility/StreamString.h"
 
 using namespace lldb;
 using namespace lldb_private;
@@ -26,7 +27,7 @@ using namespace lldb_private;
 
 CommandObjectQuit::CommandObjectQuit(CommandInterpreter &interpreter)
     : CommandObjectParsed(interpreter, "quit", "Quit the LLDB debugger.",
-                          "quit") {}
+                          "quit [exit-code]") {}
 
 CommandObjectQuit::~CommandObjectQuit() {}
 
@@ -77,6 +78,41 @@ bool CommandObjectQuit::DoExecute(Args &command, CommandReturnObject &result) {
       return false;
     }
   }
+
+  if (command.GetArgumentCount() > 1) {
+    result.AppendError("Too many arguments for 'quit'. Only an optional exit "
+                       "code is allowed");
+    result.SetStatus(eReturnStatusFailed);
+    return false;
+  }
+
+  if (command.GetArgumentCount() > 1) {
+    result.AppendError("Too many arguments for 'quit'. Only an optional exit "
+                       "code is allowed");
+    result.SetStatus(eReturnStatusFailed);
+    return false;
+  }
+
+  // We parse the exit code argument if there is one.
+  if (command.GetArgumentCount() == 1) {
+    llvm::StringRef arg = command.GetArgumentAtIndex(0);
+    int exit_code;
+    if (arg.getAsInteger(/*autodetect radix*/ 0, exit_code)) {
+      lldb_private::StreamString s;
+      std::string arg_str = arg.str();
+      s.Printf("Couldn't parse '%s' as integer for exit code.", arg_str.data());
+      result.AppendError(s.GetString());
+      result.SetStatus(eReturnStatusFailed);
+      return false;
+    }
+    if (!m_interpreter.SetQuitExitCode(exit_code)) {
+      result.AppendError("The current driver doesn't allow custom exit codes"
+                         " for the quit command.");
+      result.SetStatus(eReturnStatusFailed);
+      return false;
+    }
+  }
+
   const uint32_t event_type =
       CommandInterpreter::eBroadcastBitQuitCommandReceived;
   m_interpreter.BroadcastEvent(event_type);

lldb/source/Interpreter/CommandInterpreter.cpp

@@ -144,6 +144,26 @@ void CommandInterpreter::SetPromptOnQuit(bool b) {
   m_collection_sp->SetPropertyAtIndexAsBoolean(nullptr, idx, b);
 }
 
+void CommandInterpreter::AllowExitCodeOnQuit(bool allow) {
+  m_allow_exit_code = allow;
+  if (!allow)
+    m_quit_exit_code.reset();
+}
+
+bool CommandInterpreter::SetQuitExitCode(int exit_code) {
+  if (!m_allow_exit_code)
+    return false;
+  m_quit_exit_code = exit_code;
+  return true;
+}
+
+int CommandInterpreter::GetQuitExitCode(bool &exited) const {
+  exited = m_quit_exit_code.hasValue();
+  if (exited)
+    return *m_quit_exit_code;
+  return 0;
+}
+
 void CommandInterpreter::ResolveCommand(const char *command_line,
                                         CommandReturnObject &result) {
   std::string command = command_line;

lldb/tools/driver/Driver.cpp

@@ -962,7 +962,7 @@ std::string EscapeString(std::string arg) {
   return '"' + arg + '"';
 }
 
-void Driver::MainLoop() {
+int Driver::MainLoop() {
   if (::tcgetattr(STDIN_FILENO, &g_old_stdin_termios) == 0) {
     g_old_stdin_termios_is_valid = true;
     atexit(reset_stdin_termios);
@@ -1001,6 +1001,10 @@ void Driver::MainLoop() {
     result.PutOutput(m_debugger.GetOutputFileHandle());
   }
 
+  // We allow the user to specify an exit code when calling quit which we will
+  // return when exiting.
+  m_debugger.GetCommandInterpreter().AllowExitCodeOnQuit(true);
+
   // Now we handle options we got from the command line
   SBStream commands_stream;
 
@@ -1159,7 +1163,9 @@ void Driver::MainLoop() {
   reset_stdin_termios();
   fclose(stdin);
 
+  int exit_code = sb_interpreter.GetQuitStatus();
   SBDebugger::Destroy(m_debugger);
+  return exit_code;
 }
 
 void Driver::ResizeWindow(unsigned short col) {
@@ -1237,6 +1243,7 @@ main(int argc, char const *argv[])
   signal(SIGCONT, sigcont_handler);
 #endif
 
+  int exit_code = 0;
   // Create a scope for driver so that the driver object will destroy itself
   // before SBDebugger::Terminate() is called.
   {
@@ -1245,14 +1252,15 @@ main(int argc, char const *argv[])
     bool exiting = false;
     SBError error(driver.ParseArgs(argc, argv, stdout, exiting));
     if (error.Fail()) {
+      exit_code = 1;
       const char *error_cstr = error.GetCString();
       if (error_cstr)
         ::fprintf(stderr, "error: %s\n", error_cstr);
     } else if (!exiting) {
-      driver.MainLoop();
+      exit_code = driver.MainLoop();
     }
   }
 
   SBDebugger::Terminate();
-  return 0;
+  return exit_code;
 }

lldb/tools/driver/Driver.h

@@ -37,7 +37,10 @@ public:
 
   virtual ~Driver();
 
-  void MainLoop();
+  /// Runs the main loop.
+  ///
+  /// @return The exit code that the process should return.
+  int MainLoop();
 
   lldb::SBError ParseArgs(int argc, const char *argv[], FILE *out_fh,
                           bool &do_exit);