From e985ead7f9377e4d61d2e47c0b964f917d8054b4 Mon Sep 17 00:00:00 2001
From: zhenwei pi <pizhenwei@bytedance.com>
Date: Fri, 1 Nov 2024 20:28:09 +0800
Subject: [PATCH] RDMA: Prevent IO for child process (#1244)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

RDMA MR (memory region) is not forkable, the VMA (virtual memory area)
of a MR gets empty in a child process. Prevent IO for child process to
avoid server crash.

In the check for whether read and write is allowed in an RDMA
connection, a check that if we're in a child process is added. If we
are, the function returns an error, which will cause the RDMA client to
be disconnected.

Suggested-by: Viktor Söderqvist <viktor.soderqvist@est.tech>
Signed-off-by: zhenwei pi <pizhenwei@bytedance.com>
---
 src/rdma.c | 36 +++++++++++++++++++++++++++++++-----
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/rdma.c b/src/rdma.c
index 9329fd3ab5..bb38baa0f1 100644
--- a/src/rdma.c
+++ b/src/rdma.c
@@ -143,12 +143,34 @@ static void serverRdmaError(char *err, const char *fmt, ...) {
     va_end(ap);
 }
 
+static inline int connRdmaAllowCommand(void) {
+    /* RDMA MR is not accessible in a child process, avoid segment fault due to
+     * invalid MR access, close it rather than server random crash */
+    if (server.in_fork_child != CHILD_TYPE_NONE) {
+        return C_ERR;
+    }
+
+    return C_OK;
+}
+
+static inline int connRdmaAllowRW(connection *conn) {
+    if (conn->state == CONN_STATE_ERROR || conn->state == CONN_STATE_CLOSED) {
+        return C_ERR;
+    }
+
+    return connRdmaAllowCommand();
+}
+
 static int rdmaPostRecv(RdmaContext *ctx, struct rdma_cm_id *cm_id, ValkeyRdmaCmd *cmd) {
     struct ibv_sge sge;
     size_t length = sizeof(ValkeyRdmaCmd);
     struct ibv_recv_wr recv_wr, *bad_wr;
     int ret;
 
+    if (connRdmaAllowCommand()) {
+        return C_ERR;
+    }
+
     sge.addr = (uint64_t)cmd;
     sge.length = length;
     sge.lkey = ctx->cmd_mr->lkey;
@@ -1214,6 +1236,10 @@ static size_t connRdmaSend(connection *conn, const void *data, size_t data_len)
     char *remote_addr = ctx->tx_addr + ctx->tx.offset;
     int ret;
 
+    if (connRdmaAllowCommand()) {
+        return C_ERR;
+    }
+
     memcpy(addr, data, data_len);
 
     sge.addr = (uint64_t)addr;
@@ -1247,7 +1273,7 @@ static int connRdmaWrite(connection *conn, const void *data, size_t data_len) {
     RdmaContext *ctx = cm_id->context;
     uint32_t towrite;
 
-    if (conn->state == CONN_STATE_ERROR || conn->state == CONN_STATE_CLOSED) {
+    if (connRdmaAllowRW(conn)) {
         return C_ERR;
     }
 
@@ -1290,7 +1316,7 @@ static int connRdmaRead(connection *conn, void *buf, size_t buf_len) {
     struct rdma_cm_id *cm_id = rdma_conn->cm_id;
     RdmaContext *ctx = cm_id->context;
 
-    if (conn->state == CONN_STATE_ERROR || conn->state == CONN_STATE_CLOSED) {
+    if (connRdmaAllowRW(conn)) {
         return C_ERR;
     }
 
@@ -1312,7 +1338,7 @@ static ssize_t connRdmaSyncWrite(connection *conn, char *ptr, ssize_t size, long
     long long start = mstime();
     uint32_t towrite;
 
-    if (conn->state == CONN_STATE_ERROR || conn->state == CONN_STATE_CLOSED) {
+    if (connRdmaAllowRW(conn)) {
         return C_ERR;
     }
 
@@ -1355,7 +1381,7 @@ static ssize_t connRdmaSyncRead(connection *conn, char *ptr, ssize_t size, long
     long long start = mstime();
     uint32_t toread;
 
-    if (conn->state == CONN_STATE_ERROR || conn->state == CONN_STATE_CLOSED) {
+    if (connRdmaAllowRW(conn)) {
         return C_ERR;
     }
 
@@ -1390,7 +1416,7 @@ static ssize_t connRdmaSyncReadLine(connection *conn, char *ptr, ssize_t size, l
     char *c;
     char nl = 0;
 
-    if (conn->state == CONN_STATE_ERROR || conn->state == CONN_STATE_CLOSED) {
+    if (connRdmaAllowRW(conn)) {
         return C_ERR;
     }