From ed8e38f15442dd6e93de652f405e8b36fbae07f2 Mon Sep 17 00:00:00 2001
From: levy <zlw5307@163.com>
Date: Tue, 8 Sep 2020 15:38:49 +0800
Subject: [PATCH] feat(security): server_negotiation handle challenge response

---
 src/runtime/security/sasl_server_wrapper.cpp | 13 +++++-
 src/runtime/security/server_negotiation.cpp  | 15 ++++++-
 src/runtime/security/server_negotiation.h    |  2 +
 src/runtime/test/server_negotiation_test.cpp | 43 ++++++++++++++++++++
 4 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/src/runtime/security/sasl_server_wrapper.cpp b/src/runtime/security/sasl_server_wrapper.cpp
index ef3e892fc5..7ba6517ef5 100644
--- a/src/runtime/security/sasl_server_wrapper.cpp
+++ b/src/runtime/security/sasl_server_wrapper.cpp
@@ -58,8 +58,17 @@ error_s sasl_server_wrapper::start(const std::string &mechanism,
 
 error_s sasl_server_wrapper::step(const std::string &input, std::string &output)
 {
-    // TBD(zlw)
-    return error_s::make(ERR_OK);
+    FAIL_POINT_INJECT_F("sasl_server_wrapper_step", [](dsn::string_view str) {
+        error_code err = error_code::try_get(str.data(), ERR_UNKNOWN);
+        return error_s::make(err);
+    });
+
+    const char *msg = nullptr;
+    unsigned msg_len = 0;
+    int sasl_err = sasl_server_step(_conn, input.c_str(), input.length(), &msg, &msg_len);
+
+    output.assign(msg, msg_len);
+    return wrap_error(sasl_err);
 }
 } // namespace security
 } // namespace dsn
diff --git a/src/runtime/security/server_negotiation.cpp b/src/runtime/security/server_negotiation.cpp
index 807689ce53..4ee839fb4d 100644
--- a/src/runtime/security/server_negotiation.cpp
+++ b/src/runtime/security/server_negotiation.cpp
@@ -53,7 +53,7 @@ void server_negotiation::handle_request(negotiation_rpc rpc)
         on_initiate(rpc);
         break;
     case negotiation_status::type::SASL_CHALLENGE:
-        // TBD(zlw)
+        on_challenge_resp(rpc);
         break;
     default:
         fail_negotiation();
@@ -115,6 +115,19 @@ void server_negotiation::on_initiate(negotiation_rpc rpc)
     return do_challenge(rpc, err_s, start_output);
 }
 
+void server_negotiation::on_challenge_resp(negotiation_rpc rpc)
+{
+    const negotiation_request &request = rpc.request();
+    if (!check_status(request.status, negotiation_status::type::SASL_CHALLENGE_RESP)) {
+        fail_negotiation();
+        return;
+    }
+
+    std::string resp_msg;
+    error_s err_s = _sasl->step(request.msg, resp_msg);
+    return do_challenge(rpc, err_s, resp_msg);
+}
+
 void server_negotiation::do_challenge(negotiation_rpc rpc,
                                       error_s err_s,
                                       const std::string &resp_msg)
diff --git a/src/runtime/security/server_negotiation.h b/src/runtime/security/server_negotiation.h
index a8d97afb20..0e39888b01 100644
--- a/src/runtime/security/server_negotiation.h
+++ b/src/runtime/security/server_negotiation.h
@@ -37,6 +37,8 @@ class server_negotiation : public negotiation
     void on_list_mechanisms(negotiation_rpc rpc);
     void on_select_mechanism(negotiation_rpc rpc);
     void on_initiate(negotiation_rpc rpc);
+    void on_challenge_resp(negotiation_rpc rpc);
+
     void do_challenge(negotiation_rpc rpc, error_s err_s, const std::string &resp_msg);
 
     friend class server_negotiation_test;
diff --git a/src/runtime/test/server_negotiation_test.cpp b/src/runtime/test/server_negotiation_test.cpp
index abc69e64fb..3524be8fc8 100644
--- a/src/runtime/test/server_negotiation_test.cpp
+++ b/src/runtime/test/server_negotiation_test.cpp
@@ -50,6 +50,8 @@ class server_negotiation_test : public testing::Test
 
     void on_initiate(negotiation_rpc rpc) { _srv_negotiation->on_initiate(rpc); }
 
+    void on_challenge_resp(negotiation_rpc rpc) { _srv_negotiation->on_challenge_resp(rpc); }
+
     negotiation_status::type get_negotiation_status() { return _srv_negotiation->_status; }
 
     // _sim_session is used for holding the sim_rpc_session which is created in ctor,
@@ -178,5 +180,46 @@ TEST_F(server_negotiation_test, on_initiate)
         }
     }
 }
+
+TEST_F(server_negotiation_test, on_challenge_resp)
+{
+    struct
+    {
+        std::string sasl_step_result;
+        negotiation_status::type req_status;
+        negotiation_status::type resp_status;
+        negotiation_status::type nego_status;
+    } tests[] = {{"ERR_TIMEOUT",
+                  negotiation_status::type::SASL_CHALLENGE_RESP,
+                  negotiation_status::type::INVALID,
+                  negotiation_status::type::SASL_AUTH_FAIL},
+                 {"ERR_OK",
+                  negotiation_status::type::SASL_SELECT_MECHANISMS,
+                  negotiation_status::type::INVALID,
+                  negotiation_status::type::SASL_AUTH_FAIL},
+                 {"ERR_SASL_INCOMPLETE",
+                  negotiation_status::type::SASL_CHALLENGE_RESP,
+                  negotiation_status::type::SASL_CHALLENGE,
+                  negotiation_status::type::SASL_CHALLENGE},
+                 {"ERR_OK",
+                  negotiation_status::type::SASL_CHALLENGE_RESP,
+                  negotiation_status::type::SASL_SUCC,
+                  negotiation_status::type::SASL_SUCC}};
+
+    RPC_MOCKING(negotiation_rpc)
+    {
+        for (const auto &test : tests) {
+            fail::setup();
+            fail::cfg("sasl_server_wrapper_step", "return(" + test.sasl_step_result + ")");
+
+            auto rpc = create_negotiation_rpc(test.req_status, "");
+            on_challenge_resp(rpc);
+            ASSERT_EQ(rpc.response().status, test.resp_status);
+            ASSERT_EQ(get_negotiation_status(), test.nego_status);
+
+            fail::teardown();
+        }
+    }
+}
 } // namespace security
 } // namespace dsn