From 9a85b8eb06f0b845c32abf9a0f1b2fe40477ba2a Mon Sep 17 00:00:00 2001
From: Johan Wikman <johan.wikman@mariadb.com>
Date: Thu, 12 Jan 2017 13:45:44 +0200
Subject: [PATCH] Masking: Handle multi results

---
 .../modules/filter/masking/maskingfilter.cc   |  2 +-
 .../filter/masking/maskingfiltersession.cc    | 18 ++++++--
 .../filter/masking/maskingfiltersession.hh    | 11 +++++
 server/modules/filter/masking/mysql.hh        | 44 +++++++++++++++++++
 4 files changed, 70 insertions(+), 5 deletions(-)

diff --git a/server/modules/filter/masking/maskingfilter.cc b/server/modules/filter/masking/maskingfilter.cc
index 1aa92e344..710cc44b6 100644
--- a/server/modules/filter/masking/maskingfilter.cc
+++ b/server/modules/filter/masking/maskingfilter.cc
@@ -153,7 +153,7 @@ void MaskingFilter::diagnostics(DCB* pDcb)
 // static
 uint64_t MaskingFilter::getCapabilities()
 {
-    return RCAP_TYPE_STMT_INPUT | RCAP_TYPE_STMT_OUTPUT;
+    return RCAP_TYPE_STMT_INPUT | RCAP_TYPE_CONTIGUOUS_OUTPUT;
 }
 
 std::tr1::shared_ptr<MaskingRules> MaskingFilter::rules() const
diff --git a/server/modules/filter/masking/maskingfiltersession.cc b/server/modules/filter/masking/maskingfiltersession.cc
index acda3a3a2..91494fcd7 100644
--- a/server/modules/filter/masking/maskingfiltersession.cc
+++ b/server/modules/filter/masking/maskingfiltersession.cc
@@ -103,6 +103,7 @@ void MaskingFilterSession::handle_response(GWBUF* pPacket)
     switch (response.type())
     {
     case 0x00: // OK
+        // We'll end up here also in the case of a multi-result.
     case 0xff: // ERR
     case 0xfb: // GET_MORE_CLIENT_DATA/SEND_MORE_CLIENT_DATA
         m_state = EXPECTING_NOTHING;
@@ -144,8 +145,6 @@ void MaskingFilterSession::handle_field(GWBUF* pPacket)
         // All fields have been read.
         m_state = EXPECTING_FIELD_EOF;
     }
-
-    MXS_NOTICE("Stats: %s", column_def.to_string().c_str());
 }
 
 void MaskingFilterSession::handle_eof(GWBUF* pPacket)
@@ -199,8 +198,19 @@ void MaskingFilterSession::handle_row(GWBUF* pPacket)
     {
     case ComPacket::EOF_PACKET:
         // EOF after last row.
-        MXS_NOTICE("EOF after last row received.");
-        m_state = EXPECTING_NOTHING;
+        {
+            ComEOF eof(response);
+
+            if (eof.status() & SERVER_MORE_RESULTS_EXIST)
+            {
+                m_res.reset_multi();
+                m_state = EXPECTING_RESPONSE;
+            }
+            else
+            {
+                m_state = EXPECTING_NOTHING;
+            }
+        }
         break;
 
     default:
diff --git a/server/modules/filter/masking/maskingfiltersession.hh b/server/modules/filter/masking/maskingfiltersession.hh
index 7e7f59c22..710dcc048 100644
--- a/server/modules/filter/masking/maskingfiltersession.hh
+++ b/server/modules/filter/masking/maskingfiltersession.hh
@@ -67,16 +67,24 @@ private:
             : m_command(0)
             , m_nTotal_fields(0)
             , m_index(0)
+            , m_multi_result(false)
         {}
 
         void reset(uint8_t command, const SMaskingRules& sRules)
         {
             m_command = command;
             m_sRules = sRules;
+
+            reset_multi();
+        }
+
+        void reset_multi()
+        {
             m_nTotal_fields = 0;
             m_types.clear();
             m_rules.clear();
             m_index = 0;
+            m_multi_result = true;
         }
 
         uint8_t command() const
@@ -89,6 +97,8 @@ private:
             return m_sRules;
         }
 
+        bool is_multi_result() const { return m_multi_result; }
+
         uint32_t total_fields() const { return m_nTotal_fields; }
 
         void set_total_fields(uint32_t n) { m_nTotal_fields = n; }
@@ -124,6 +134,7 @@ private:
         std::vector<enum_field_types>          m_types;         /*<! The column types. */
         std::vector<const MaskingRules::Rule*> m_rules;         /*<! The rules applied for columns. */
         size_t                                 m_index;         /*<! Index to the current rule.*/
+        bool                                   m_multi_result;  /*<! Are we processing multi-results. */
     };
 
     const MaskingFilter& m_filter;
diff --git a/server/modules/filter/masking/mysql.hh b/server/modules/filter/masking/mysql.hh
index 26b26fcf7..e6131bab1 100644
--- a/server/modules/filter/masking/mysql.hh
+++ b/server/modules/filter/masking/mysql.hh
@@ -565,6 +565,50 @@ protected:
     uint8_t m_type;
 };
 
+class ComEOF : public ComResponse
+{
+public:
+    ComEOF(GWBUF* pPacket)
+        : ComResponse(pPacket)
+    {
+        ss_dassert(m_type == EOF_PACKET);
+
+        extract_payload();
+    }
+
+    ComEOF(const ComResponse& response)
+        : ComResponse(response)
+    {
+        ss_dassert(m_type == EOF_PACKET);
+
+        extract_payload();
+    }
+
+    uint16_t warnings() const
+    {
+        return m_warnings;
+    }
+
+    uint16_t status() const
+    {
+        return m_status;
+    }
+
+private:
+    void extract_payload()
+    {
+        m_warnings = *m_pData++;
+        m_warnings += (*m_pData++ << 8);
+
+        m_status = *m_pData++;
+        m_status += (*m_pData++ << 8);
+    }
+
+private:
+    uint16_t m_warnings;
+    uint16_t m_status;
+};
+
 /**
  * @class ComRequest
  *