setenv and getenv

justanhduc · Nov 4, 2021 · b48a2d7 · b48a2d7
1 parent 80b6b36
commit b48a2d7
Show file tree

Hide file tree

Showing 5 changed files with 116 additions and 3 deletions.
diff --git a/client.c b/client.c
@@ -794,3 +794,48 @@ void c_show_cmd() {
             warning("Wrong internal message in show_cmd");
     }
 }
+
+void c_get_env() {
+    struct Msg m;
+    int res;
+    char *string = 0;
+
+    /* Send the request */
+    m.type = GET_ENV;
+    m.u.size = strlen(command_line.label) + 1;
+    send_msg(server_socket, &m);
+    send_bytes(server_socket, command_line.label, m.u.size);
+
+    /* Receive the answer */
+    res = recv_msg(server_socket, &m);
+    if (res != sizeof(m))
+        error("Error in get_env");
+
+    switch (m.type) {
+        case LIST_LINE:
+            if (m.u.size) {
+                string = (char *) malloc(m.u.size);
+                res = recv_bytes(server_socket, string, m.u.size);
+                if (res != m.u.size)
+                    error("Error in get_env - line size");
+
+                printf("%s\n", string);
+                free(string);
+            } else
+                printf("\n");
+
+            return;
+        default:
+            warning("Wrong internal message in get_env");
+    }
+}
+
+void c_set_env() {
+    struct Msg m;
+
+    /* Send the request */
+    m.type = SET_ENV;
+    m.u.size = strlen(command_line.label) + 1;
+    send_msg(server_socket, &m);
+    send_bytes(server_socket, command_line.label, m.u.size);
+}
diff --git a/jobs.c b/jobs.c
@@ -1589,3 +1589,33 @@ void joblist_dump(int fd) {
         p = p->next;
     }
 }
+
+void s_get_env(int s, int size) {
+    char *var = malloc(size);
+    int res = recv_bytes(s, var, size);
+    if (res != size)
+        error("Receiving environment variable name");
+
+    char *val = getenv(var);
+    struct Msg m;
+    m.type = LIST_LINE;
+    m.u.size = val ? sizeof(val) + 1 : 0;
+    send_msg(s, &m);
+    if (val)
+        send_bytes(s, val, m.u.size);
+
+    free(var);
+}
+
+void s_set_env(int s, int size) {
+    char *var = malloc(size);
+    int res = recv_bytes(s, var, size);
+    if (res != size)
+        error("Receiving environment variable name");
+
+    /* get the first token */
+    char *name = strtok(var, "=");
+    char *val = strtok(NULL, "=");
+    setenv(name, val, 1);
+    free(var);
+}
diff --git a/main.c b/main.c
@@ -97,6 +97,8 @@ static struct option longOptions[] = {
         {"gpus",          required_argument, NULL, 'G'},
         {"gpu_indices",   required_argument, NULL, 'g'},
         {"full_cmd",      optional_argument, NULL, 'F'},
+        {"getenv",        required_argument, NULL, 0},
+        {"setenv",        required_argument, NULL, 0},
         {NULL, 0,                            NULL, 0}
 };
 
@@ -114,6 +116,16 @@ void parse_opts(int argc, char **argv) {
             break;
 
         switch (c) {
+            case 0:
+                if (strcmp(longOptions[optionIdx].name, "getenv") == 0) {
+                    command_line.request = c_GET_ENV;
+                    command_line.label = optarg;  /* reuse this var */
+                } else if (strcmp(longOptions[optionIdx].name, "setenv") == 0) {
+                    command_line.request = c_SET_ENV;
+                    command_line.label = optarg;  /* reuse this var */
+                } else
+                    error("Wrong option %s.", longOptions[optionIdx].name);
+                break;
             case 'K':
                 command_line.request = c_KILL_SERVER;
                 command_line.should_go_background = 0;
@@ -622,13 +634,22 @@ int main(int argc, char **argv) {
             /* This will also print the state into stdout */
             c_get_state();
             break;
+        case c_GET_ENV:
+            if (!command_line.need_server)
+                error("The command %i needs the server", command_line.request);
+            c_get_env();
+            break;
+        case c_SET_ENV:
+            if (!command_line.need_server)
+                error("The command %i needs the server", command_line.request);
+            c_set_env();
+            break;
     }
 
     if (command_line.need_server) {
         close(server_socket);
     }
     free(command_line.depend_on);
-    free(command_line.label);
 
     return errorlevel;
 }
diff --git a/main.h b/main.h
@@ -47,7 +47,9 @@ enum MsgTypes {
     GET_LABEL,
     LAST_ID,
     KILL_ALL,
-    GET_CMD
+    GET_CMD,
+    GET_ENV,
+    SET_ENV
 };
 
 enum Request {
@@ -74,7 +76,9 @@ enum Request {
     c_GET_LABEL,
     c_LAST_ID,
     c_KILL_ALL,
-    c_SHOW_CMD
+    c_SHOW_CMD,
+    c_GET_ENV,
+    c_SET_ENV
 };
 
 struct CommandLine {
@@ -281,6 +285,10 @@ void c_kill_all_jobs();
 
 void c_show_cmd();
 
+void c_get_env();
+
+void c_set_env();
+
 /* jobs.c */
 void s_list(int s);
 
@@ -346,6 +354,9 @@ void s_send_label(int s, int jobid);
 
 void s_kill_all_jobs(int s);
 
+void s_get_env(int s, int size);
+
+void s_set_env(int s, int size);
 
 /* server.c */
 void server_main(int notify_fd, char *_path);

diff --git a/server.c b/server.c
@@ -471,6 +471,12 @@ client_read(int index) {
         case GET_STATE:
             s_send_state(s, m.u.jobid);
             break;
+        case GET_ENV:
+            s_get_env(s, m.u.size);
+            break;
+        case SET_ENV:
+            s_set_env(s, m.u.size);
+            break;
         case GET_VERSION:
             s_send_version(s);
             break;