Move torture_setup_ssh_agent() and torture_cleanup_ssh_agent() to torture.c

Signed-off-by: salonidabgar <salonidabgar@gmail.com>
Reviewed-by: Jakub Jelen <jjelen@redhat.com>

tests/client/torture_auth.c

@@ -108,11 +108,8 @@ static int pubkey_setup(void **state)
 static int agent_setup(void **state)
 {
     struct torture_state *s = *state;
-    char ssh_agent_cmd[4096];
-    char ssh_agent_sock[1024];
-    char ssh_agent_pidfile[1024];
-    char ssh_key_add[1024];
     struct passwd *pwd;
+    char ssh_key_path[1024];
     int rc;
 
     rc = pubkey_setup(state);
@@ -123,45 +120,18 @@ static int agent_setup(void **state)
     pwd = getpwnam("bob");
     assert_non_null(pwd);
 
-    snprintf(ssh_agent_sock,
+    /* Use the common function to set up the SSH agent with Bob's key */
-             sizeof(ssh_agent_sock),
+    snprintf(ssh_key_path, sizeof(ssh_key_path), "%s/.ssh/id_rsa", pwd->pw_dir);
-             "%s/agent.sock",
+    rc = torture_setup_ssh_agent(s, ssh_key_path);
-             s->socket_dir);
+    if (rc != 0) {
-
+        return rc;
-    snprintf(ssh_agent_pidfile,
+    }
-             sizeof(ssh_agent_pidfile),
-             "%s/agent.pid",
-             s->socket_dir);
-
-    /* Production ready code!!! */
-    snprintf(ssh_agent_cmd,
-             sizeof(ssh_agent_cmd),
-             "eval `ssh-agent -a %s`; echo $SSH_AGENT_PID > %s",
-             ssh_agent_sock, ssh_agent_pidfile);
-
-    /* run ssh-agent and ssh-add as the normal user */
-    unsetenv("UID_WRAPPER_ROOT");
-
-    rc = system(ssh_agent_cmd);
-    assert_return_code(rc, errno);
-
-    setenv("SSH_AUTH_SOCK", ssh_agent_sock, 1);
-    setenv("TORTURE_SSH_AGENT_PIDFILE", ssh_agent_pidfile, 1);
-
-    snprintf(ssh_key_add,
-             sizeof(ssh_key_add),
-             "ssh-add %s/.ssh/id_rsa",
-             pwd->pw_dir);
-
-    rc = system(ssh_key_add);
-    assert_return_code(rc, errno);
 
     return 0;
 }
 
 static int agent_teardown(void **state)
 {
-    const char *ssh_agent_pidfile;
     int rc;
 
     rc = session_teardown(state);
@@ -169,17 +139,11 @@ static int agent_teardown(void **state)
         return rc;
     }
 
-    ssh_agent_pidfile = getenv("TORTURE_SSH_AGENT_PIDFILE");
+    /* Use the common function to clean up the SSH agent */
-    assert_non_null(ssh_agent_pidfile);
+    rc = torture_cleanup_ssh_agent();
-
+    if (rc != 0) {
-    /* kill agent pid */
+        return rc;
-    rc = torture_terminate_process(ssh_agent_pidfile);
+    }
-    assert_return_code(rc, errno);
-
-    unlink(ssh_agent_pidfile);
-
-    unsetenv("TORTURE_SSH_AGENT_PIDFILE");
-    unsetenv("SSH_AUTH_SOCK");
 
     return 0;
 }

tests/client/torture_auth_cert.c

@@ -119,11 +119,8 @@ static int session_teardown(void **state)
 static int agent_setup(void **state)
 {
     struct torture_state *s = *state;
-    char ssh_agent_cmd[4096];
-    char ssh_agent_sock[1024];
-    char ssh_agent_pidfile[1024];
-    char ssh_key_add[1024];
     struct passwd *pwd;
+    char key_path[1024];
     int rc;
 
     rc = session_setup(state);
@@ -134,45 +131,18 @@ static int agent_setup(void **state)
     pwd = getpwnam("doe");
     assert_non_null(pwd);
 
-    snprintf(ssh_agent_sock,
+    snprintf(key_path, sizeof(key_path), "%s/.ssh/id_rsa", pwd->pw_dir);
-             sizeof(ssh_agent_sock),
-             "%s/agent.sock",
-             s->socket_dir);
 
-    snprintf(ssh_agent_pidfile,
+    /* run ssh-agent and add the key */
-             sizeof(ssh_agent_pidfile),
+    rc = torture_setup_ssh_agent(s, key_path);
-             "%s/agent.pid",
+    assert_int_equal(rc, 0);
-             s->socket_dir);
-
-    /* Production ready code!!! */
-    snprintf(ssh_agent_cmd,
-             sizeof(ssh_agent_cmd),
-             "eval `ssh-agent -a %s`; echo $SSH_AGENT_PID > %s",
-             ssh_agent_sock, ssh_agent_pidfile);
-
-    /* run ssh-agent and ssh-add as the normal user */
-    unsetenv("UID_WRAPPER_ROOT");
-
-    rc = system(ssh_agent_cmd);
-    assert_return_code(rc, errno);
-
-    setenv("SSH_AUTH_SOCK", ssh_agent_sock, 1);
-    setenv("TORTURE_SSH_AGENT_PIDFILE", ssh_agent_pidfile, 1);
-
-    snprintf(ssh_key_add,
-             sizeof(ssh_key_add),
-             "ssh-add %s/.ssh/id_rsa",
-             pwd->pw_dir);
-
-    rc = system(ssh_key_add);
-    assert_return_code(rc, errno);
 
     return 0;
 }
 
 static int agent_cert_setup(void **state)
 {
-    char doe_alt_ssh_key[1024];
+    char ssh_key_cmd[1024];
     struct passwd *pwd;
     int rc;
 
@@ -185,20 +155,56 @@ static int agent_cert_setup(void **state)
     assert_non_null(pwd);
 
     /* remove all keys, load alternative key + cert */
-    snprintf(doe_alt_ssh_key,
+    snprintf(ssh_key_cmd,
-             sizeof(doe_alt_ssh_key),
+             sizeof(ssh_key_cmd),
              "ssh-add -D && ssh-add %s/.ssh/id_rsa",
              pwd->pw_dir);
 
-    rc = system(doe_alt_ssh_key);
+    rc = system(ssh_key_cmd);
     assert_return_code(rc, errno);
 
     return 0;
 }
 
+static int agent_cert_setup_explicit(void **state)
+{
+    char orig_doe_ssh_key[1024];
+    char doe_ssh_key[1024];
+    char keydata[2048];
+    struct passwd *pwd = NULL;
+    int fd;
+    int rc;
+
+    rc = agent_cert_setup(state);
+    if (rc != 0) {
+        return rc;
+    }
+
+    pwd = getpwnam("doe");
+    assert_non_null(pwd);
+
+    snprintf(orig_doe_ssh_key,
+             sizeof(orig_doe_ssh_key),
+             "%s/.ssh/id_rsa",
+             pwd->pw_dir);
+
+    snprintf(doe_ssh_key, sizeof(doe_ssh_key), "%s/.ssh/my_rsa", pwd->pw_dir);
+
+    /* move the private key away from the default location the certificate can
+     * not be loaded automatically */
+    fd = open(orig_doe_ssh_key, O_RDONLY);
+    assert_true(fd > 0);
+    rc = read(fd, keydata, sizeof(keydata));
+    assert_true(rc > 0);
+    keydata[rc] = '\0';
+    close(fd);
+    torture_write_file(doe_ssh_key, keydata);
+
+    return 0;
+}
+
 static int agent_teardown(void **state)
 {
-    const char *ssh_agent_pidfile;
     int rc;
 
     rc = session_teardown(state);
@@ -206,17 +212,8 @@ static int agent_teardown(void **state)
         return rc;
     }
 
-    ssh_agent_pidfile = getenv("TORTURE_SSH_AGENT_PIDFILE");
+    rc = torture_cleanup_ssh_agent();
-    assert_non_null(ssh_agent_pidfile);
+    assert_int_equal(rc, 0);
-
-    /* kill agent pid */
-    rc = torture_terminate_process(ssh_agent_pidfile);
-    assert_return_code(rc, errno);
-
-    unlink(ssh_agent_pidfile);
-
-    unsetenv("TORTURE_SSH_AGENT_PIDFILE");
-    unsetenv("SSH_AUTH_SOCK");
 
     return 0;
 }
@@ -705,45 +702,7 @@ torture_auth_agent_cert_identities_only_nonblocking(void **state)
     assert_ssh_return_code(session, rc);
 }
 
-static int agent_cert_setup_explicit(void **state)
+static void torture_auth_agent_cert_identities_only_explicit(void **state)
-{
-    char orig_doe_ssh_key[1024];
-    char doe_ssh_key[1024];
-    char keydata[2048];
-    struct passwd *pwd = NULL;
-    int fd ;
-    int rc;
-
-    agent_cert_setup(state);
-
-    pwd = getpwnam("doe");
-    assert_non_null(pwd);
-
-    snprintf(orig_doe_ssh_key,
-             sizeof(orig_doe_ssh_key),
-             "%s/.ssh/id_rsa",
-             pwd->pw_dir);
-
-    snprintf(doe_ssh_key,
-             sizeof(doe_ssh_key),
-             "%s/.ssh/my_rsa",
-             pwd->pw_dir);
-
-    /* move the private key away from the default location the certificate can
-     * not be loaded automatically */
-    fd = open(orig_doe_ssh_key, O_RDONLY);
-    assert_true(fd > 0);
-    rc = read(fd, keydata, sizeof(keydata));
-    assert_true(rc > 0);
-    keydata[rc] = '\0';
-    close(fd);
-    torture_write_file(doe_ssh_key, keydata);
-
-    return 0;
-}
-
-static void
-torture_auth_agent_cert_identities_only_explicit(void **state)
 {
     struct torture_state *s = *state;
     ssh_session session = s->ssh.session;

tests/torture.c

@@ -1983,3 +1983,108 @@ int main(int argc, char **argv) {
 
     return torture_run_tests();
 }
+
+/**
+ * @brief Setup an SSH agent for testing
+ *
+ * This function starts an SSH agent, exports the environment variables,
+ * and optionally adds an SSH key to the agent.
+ *
+ * @param s             The torture state
+ * @param add_key       Path to the key to add to the agent, or NULL to skip
+ *
+ * @return              0 on success, -1 on error
+ */
+int torture_setup_ssh_agent(struct torture_state *s, const char *add_key)
+{
+#ifndef WIN32
+    int rc;
+    char ssh_agent_cmd[4096];
+    char ssh_agent_sock[1024];
+    char ssh_agent_pidfile[1024];
+    char long_cmd[2048];
+
+    /* Setup SSH agent */
+    snprintf(ssh_agent_sock,
+             sizeof(ssh_agent_sock),
+             "%s/agent.sock",
+             s->socket_dir);
+
+    snprintf(ssh_agent_pidfile,
+             sizeof(ssh_agent_pidfile),
+             "%s/agent.pid",
+             s->socket_dir);
+
+    /* Create command to start SSH agent with our custom socket */
+    snprintf(ssh_agent_cmd,
+             sizeof(ssh_agent_cmd),
+             "eval `ssh-agent -a %s`; echo $SSH_AGENT_PID > %s",
+             ssh_agent_sock,
+             ssh_agent_pidfile);
+
+    /* Run ssh-agent as the normal user */
+    torture_unsetenv("UID_WRAPPER_ROOT");
+
+    rc = system(ssh_agent_cmd);
+    if (rc != 0) {
+        return -1;
+    }
+
+    /* Set environment variables for SSH agent */
+    torture_setenv("SSH_AUTH_SOCK", ssh_agent_sock);
+    torture_setenv("TORTURE_SSH_AGENT_PIDFILE", ssh_agent_pidfile);
+
+    /* Add key to the agent if specified */
+    if (add_key != NULL) {
+        snprintf(long_cmd, sizeof(long_cmd), "ssh-add %s", add_key);
+        rc = system(long_cmd);
+        if (rc != 0) {
+            return -1;
+        }
+    }
+
+    return 0;
+#else
+    /* On Windows, we don't set up an SSH agent */
+    (void)s;
+    (void)add_key;
+
+    /* Return failure to make it explicit that agent forwarding isn't supported
+     * on Windows */
+    return -1;
+#endif
+}
+
+/**
+ * @brief Teardown an SSH agent
+ *
+ * This function kills the SSH agent process and cleans up environment
+ * variables.
+ *
+ * @return              0 on success, -1 on error
+ */
+int torture_cleanup_ssh_agent(void)
+{
+#ifndef WIN32
+    const char *ssh_agent_pidfile;
+    int rc;
+
+    ssh_agent_pidfile = getenv("TORTURE_SSH_AGENT_PIDFILE");
+    if (ssh_agent_pidfile == NULL) {
+        return 0;
+    }
+
+    rc = torture_terminate_process(ssh_agent_pidfile);
+    if (rc != 0) {
+        return -1;
+    }
+
+    torture_unsetenv("TORTURE_SSH_AGENT_PIDFILE");
+    torture_unsetenv("SSH_AUTH_SOCK");
+
+    return 0;
+#else
+    /* On Windows, we don't start an SSH agent, so nothing to clean up */
+    return -1;
+#endif
+}

tests/torture.h

@@ -182,5 +182,7 @@ int torture_change_dir(char *path);
 
 void torture_setenv(char const* variable, char const* value);
 void torture_unsetenv(char const* variable);
+int torture_setup_ssh_agent(struct torture_state *s, const char *add_key);
+int torture_cleanup_ssh_agent(void);
 
 #endif /* _TORTURE_H */