Merge pull request #5703 from grondo/pty-fix

mergify[bot] · web-flow · commit 0719ae6b1cb6 · 2024-01-24T02:10:49.000Z
enable encode of pty data as base64 and make `flux alloc vi` test more reliable
diff --git a/src/common/libterminus/client.c b/src/common/libterminus/client.c
@@ -26,6 +26,7 @@
 #include "src/common/libczmqcontainers/czmq_containers.h"
 #include "src/common/libutil/fdutils.h"
 #include "src/common/libutil/llog.h"
+#include "src/common/libutil/errno_safe.h"
 #include "ccan/str/str.h"
 
 #include "pty.h"
@@ -257,11 +258,17 @@ static void pty_client_attached (struct flux_pty_client *c)
 
 static void pty_client_data (struct flux_pty_client *c, flux_future_t *f)
 {
-    const char *data;
+    const flux_msg_t *msg;
+    char *data;
     size_t len;
+    flux_error_t error;
 
-    if (flux_rpc_get_unpack (f, "{s:s%}", "data", &data, &len) < 0) {
-        llog_error (c, "unpack: %s", future_strerror (f, errno));
+    if (flux_future_get (f, (const void **)&msg) < 0) {
+        llog_error (c, "data response: %s", future_strerror (f, errno));
+        return;
+    }
+    if (pty_data_unpack (msg, &error, &data, &len) < 0) {
+        llog_error (c, "unpack: %s", error.text);
         return;
     }
     if (write (STDIN_FILENO, data, len) < 0) {
@@ -407,14 +414,18 @@ flux_future_t *flux_pty_client_write (struct flux_pty_client *c,
                                       const void *buf,
                                       ssize_t len)
 {
+    flux_future_t *f;
+    json_t *o = NULL;
+
     if (!c || !buf || len < 0) {
         errno = EINVAL;
         return NULL;
     }
-    return flux_rpc_pack (c->h, c->service, c->rank, 0,
-                          "{s:s s:s%}",
-                          "type", "data",
-                          "data", buf, len);
+    if (!(o = pty_data_encode (buf, len)) )
+        return NULL;
+    f = flux_rpc_pack (c->h, c->service, c->rank, 0, "O", o);
+    ERRNO_SAFE_WRAP (json_decref, o);
+    return f;
 }
 
 static void pty_read_cb (flux_reactor_t *r,
diff --git a/src/common/libterminus/pty.c b/src/common/libterminus/pty.c
@@ -56,7 +56,9 @@
 #include "src/common/libutil/fdutils.h"
 #include "src/common/libutil/errno_safe.h"
 #include "src/common/libutil/aux.h"
+#include "src/common/libutil/errprintf.h"
 #include "ccan/str/str.h"
+#include "ccan/base64/base64.h"
 
 #define LLOG_SUBSYSTEM "pty"
 #include "src/common/libutil/llog.h"
@@ -405,6 +407,55 @@ int flux_pty_attach (struct flux_pty *pty)
     return 0;
 }
 
+static int encode_base64 (const char *data, int len, char **destp, int *lenp)
+{
+    ssize_t n;
+    char *dest = NULL;
+    size_t destlen = base64_encoded_length (len) + 1;
+
+    if (!(dest = malloc (destlen))
+        || ((n = base64_encode (dest, destlen, data, len)) < 0))
+        return -1;
+    *destp = dest;
+    *lenp = n;
+    return 0;
+}
+
+static json_t *pty_data_encode_base64 (const void *data, int len)
+{
+    json_t *o = NULL;
+    char *b64data = NULL;
+    int b64len = 0;
+
+    if (encode_base64 (data, len, &b64data, &b64len) < 0)
+        return NULL;
+
+    o = json_pack ("{s:s s:s s:s#}",
+                   "type", "data",
+                   "encoding", "base64",
+                   "data", b64data, b64len);
+
+    ERRNO_SAFE_WRAP (free, b64data);
+    return o;
+}
+
+json_t *pty_data_encode (const void *data, int len)
+{
+    json_t *o;
+
+    if (!(o = json_pack ("{s:s s:s#}",
+                         "type", "data",
+                         "data", (char *) data, len))) {
+        /*
+         *  json_pack() may fail if there are bytes that cannot be encoded
+         *  as UTF-8, e.g. U+0000 is specifically not allowed in jansson.
+         *  Try encoding as base64 instead.
+         */
+        o = pty_data_encode_base64 (data, len);
+    }
+    return o;
+}
+
 void pty_client_send_data (struct flux_pty *pty, void *data, int len)
 {
     struct pty_client *c = zlist_first (pty->clients);
@@ -414,12 +465,10 @@ void pty_client_send_data (struct flux_pty *pty, void *data, int len)
 
     while (c) {
         if (c->read_enabled) {
-            if (flux_respond_pack (pty->h,
-                                   c->req,
-                                   "{s:s s:s#}",
-                                   "type", "data",
-                                   "data", (char *) data, len) < 0)
+            json_t *o = pty_data_encode (data, len);
+            if (!o || flux_respond_pack (pty->h, c->req, "O", o) < 0)
                 llog_error (pty, "send data: %s", strerror (errno));
+            json_decref (o);
         }
         c = zlist_next (pty->clients);
     }
@@ -570,12 +619,82 @@ static int pty_attach (struct flux_pty *pty, const flux_msg_t *msg)
     return -1;
 }
 
+static int decode_base64 (char *src,
+                          size_t srclen,
+                          char **datap,
+                          size_t *lenp)
+{
+    ssize_t rc;
+    char *data;
+    size_t size = base64_decoded_length (srclen) + 1;
+
+    if (!(data = malloc (size)))
+        return -1;
+    if ((rc = base64_decode (data, size, src, srclen)) < 0) {
+        ERRNO_SAFE_WRAP (free, data);
+        return -1;
+    }
+    *lenp = rc;
+    *datap = data;
+    return 0;
+}
+
+int pty_data_unpack (const flux_msg_t *msg,
+                     flux_error_t *errp,
+                     char **datap,
+                     size_t *lenp)
+{
+    const char *encoding = NULL;
+    void *data;
+    int len;
+
+    if (flux_msg_unpack (msg,
+                         "{s:s s?s}",
+                         "data", &data,
+                         "encoding", &encoding) < 0)
+        return errprintf (errp,
+                          "failed to unpack data msg: %s",
+                          strerror (errno));
+    len = strlen (data);
+
+    if (encoding) {
+        if (streq (encoding, "base64")) {
+            char *b64data = NULL;
+            size_t b64len;
+            if (decode_base64 (data, len, &b64data, &b64len) < 0) {
+                return errprintf (errp,
+                                  "failed to decode %d bytes of base64",
+                                  len);
+            }
+            if (flux_msg_aux_set (msg, NULL, b64data, free) < 0) {
+                ERRNO_SAFE_WRAP (free, b64data);
+                return errprintf (errp,
+                                  "flux_msg_aux_set: %s",
+                                  strerror (errno));
+            }
+            data = b64data;
+            len = b64len;
+        }
+        else {
+            errno = EPROTO;
+            return errprintf (errp,
+                              "invalid pty data encoding: %s",
+                              encoding);
+        }
+    }
+    *datap = data;
+    *lenp = len;
+    return 0;
+}
+
 static int pty_write (struct flux_pty *pty, const flux_msg_t *msg)
 {
-    const char *data;
+    char *data;
     size_t len;
-    if (flux_msg_unpack (msg, "{s:s%}", "data", &data, &len) < 0) {
-        llog_error (pty, "msg_unpack failed");
+    flux_error_t error;
+
+    if (pty_data_unpack (msg, &error, &data, &len) < 0) {
+        llog_error (pty, "%s", error.text);
         return -1;
     }
     if (write (pty->leader, data, len) < 0) {
diff --git a/src/common/libterminus/pty.h b/src/common/libterminus/pty.h
@@ -13,6 +13,7 @@
 #ifndef FLUX_PTY_H
 #define FLUX_PTY_H
 
+#include <jansson.h>
 #include <flux/core.h>
 
 typedef void (*pty_log_f) (void *arg,
@@ -201,10 +202,23 @@ int flux_pty_aux_set (struct flux_pty *pty,
 
 void * flux_pty_aux_get (struct flux_pty *pty, const char *name);
 
+
 /*  Function exported for testing only
  */
 void pty_client_send_data (struct flux_pty *pty, void *data, int len);
 
+
+/*  Common function for unpacking pty data msg.
+ */
+int pty_data_unpack (const flux_msg_t *msg,
+                     flux_error_t *errp,
+                     char **datap,
+                     size_t *lenp);
+
+/*  Common function for encoding pty data msg.
+ */
+json_t *pty_data_encode (const void *data, int len);
+
 #endif /* !FLUX_PTY_H */
 
 /*
diff --git a/src/common/libterminus/test/pty.c b/src/common/libterminus/test/pty.c
@@ -433,6 +433,14 @@ static void test_client (void)
         rc == 0 ? "Success" : strerror (errno));
     flux_future_destroy (f);
 
+    f = flux_pty_client_write (c, "bar\0\r\n", 6);
+    ok (f != NULL,
+        "flux_pty_client_write with U+0000");
+    ok (rc == 0,
+        "flux_pty_client_write: %s",
+        rc == 0 ? "Success" : strerror (errno));
+    flux_future_destroy (f);
+
     ok (flux_pty_client_detach (c) == 0,
         "flux_pty_client_detach");
 
diff --git a/t/t2712-python-cli-alloc.t b/t/t2712-python-cli-alloc.t
@@ -173,13 +173,12 @@ test_expect_success 'flux alloc: resource.norestrict works in subinstance' '
 '
 test_expect_success 'flux alloc: flux alloc vi works' '
 	cat <<-'EOF' >input.json &&
-	[{"expect":"test\\\.txt", "send":":q!\n", "timeout":30}]
+	[{"expect":"test text file", "send":":q!\n", "timeout":30}]
 	EOF
 	cat <<-EOF >test.txt &&
 	test text file
 	EOF
-	$runpty -o vi.out --expect=input.json flux alloc -n1 vi test.txt &&
-	grep "test text file" vi.out
+	$runpty -o vi.out --expect=input.json flux alloc -n1 vi test.txt
 '
 test_expect_success 'flux alloc: flux alloc flux alloc works' '
 	cat <<-'EOF' >input2.json &&
