Add master node session recovery support

5 files changed, 133 insertions, 30 deletions

diff --git a/raptorsmiface/libraptorsmiface.c b/raptorsmiface/libraptorsmiface.c
index 7ea3de09..76220dfe 100644
--- a/raptorsmiface/libraptorsmiface.c
+++ b/raptorsmiface/libraptorsmiface.c
@@ -103,11 +103,11 @@ char raptor_sm_deallocate_session(char* username) {
 		return 1;
 	}
 
-	pid_t serverpid = raptor_sm_get_pid_for_username(username);
+	char* hostname = raptor_sm_get_hostname_for_username(username, false);
+	pid_t serverpid = raptor_sm_get_pid_for_username(username, RAPTOR_SM_SERVER_PID_FIELD);
 	if (serverpid >= 0) {
-		// Verify existence of PID on remote server
-		dprint("Verifying process %d on %s...\n\r", serverpid, row[1]);
-		char* ip = raptor_sm_get_ip_for_hostname(row[1], 0);
+		// Verify non-existence of PID on remote server before removing session information from the database
+		char* ip = raptor_sm_get_ip_for_hostname(hostname, 0);
 		char* command_string;
 		asprintf(&command_string, "ssh root@%s \'ps -p %d | grep %d\'", ip, serverpid, serverpid);
 		FILE *fp;
@@ -124,7 +124,6 @@ char raptor_sm_deallocate_session(char* username) {
 		pclose(fp);
 		free(command_string);
 		free(ip);
-		dprint("...result was %s\n\r", output);
 		if (strcmp(output, "") != 0) {
 			mysql_free_result(res);
 			mysql_close(conn);
@@ -433,7 +432,7 @@ bool raptor_sm_sesslimit_reached(char* username) {
 	return true;
 }
 
-pid_t raptor_sm_run_remote_server(char* username, char *const argv[]) {
+pid_t raptor_sm_run_remote_server(char* username, char *const argv[], char* dbfield, int display) {
 	MYSQL_RES *res;
 	MYSQL_ROW row;
 	char* query;
@@ -443,16 +442,23 @@ pid_t raptor_sm_run_remote_server(char* username, char *const argv[]) {
 		return -1;
 	}
 
-	// Respect maximum session number for the group for this user
-	if (raptor_sm_sesslimit_reached(username)) {
-		mysql_close(conn);
-		return -5;
+	if (strcmp(dbfield, RAPTOR_SM_SERVER_PID_FIELD) == 0) {
+		// Respect maximum session number for the group for this user
+		if (raptor_sm_sesslimit_reached(username)) {
+			mysql_close(conn);
+			return -5;
+		}
 	}
 
 	// Make sure a server is not already running for this user
 	// Return the existing PID if it is
 	char* safe_username = get_mysql_escaped_string(conn, username);
-	asprintf(&query, "SELECT pid,servername FROM sessions WHERE username='%s' AND state<>'%d'", safe_username, SM_STATUS_ALLOCATED);
+	if (strcmp(dbfield, RAPTOR_SM_SERVER_PID_FIELD) == 0) {
+		asprintf(&query, "SELECT %s,servername FROM sessions WHERE username='%s' AND state<>'%d'", dbfield, safe_username, SM_STATUS_ALLOCATED);
+	}
+	else {
+		asprintf(&query, "SELECT %s,servername FROM sessions WHERE username='%s'", dbfield, safe_username);
+	}
 	free(safe_username);
 	if (mysql_query_internal(conn, query)) {
 		// Server error
@@ -516,7 +522,14 @@ pid_t raptor_sm_run_remote_server(char* username, char *const argv[]) {
 		free(origstr);
 	}
 	char* origstr = command_string;
-	asprintf(&command_string, "ssh root@%s \'%s & echo $! &\'", ipaddr, origstr);
+
+	if (strcmp(dbfield, RAPTOR_SM_SERVER_PID_FIELD) == 0) {
+		asprintf(&command_string, "ssh root@%s \'%s & echo $! &\'", ipaddr, origstr);
+	}
+	else {
+		asprintf(&command_string, "ssh root@%s \"su %s -c \'export DISPLAY=:%d && %s\' &> /dev/null & echo \\$!\"", ipaddr, username, display, origstr);
+	}
+dprint("Running command %s...\n\r", command_string);
 	free(origstr);
 	
 	FILE *fp;
@@ -541,7 +554,7 @@ pid_t raptor_sm_run_remote_server(char* username, char *const argv[]) {
 	return atoi(output);
 }
 
-pid_t raptor_sm_get_pid_for_username(char* username) {
+pid_t raptor_sm_get_pid_for_username(char* username, char* dbfield) {
 	MYSQL_RES *res;
 	MYSQL_ROW row;
 	char* query;
@@ -554,7 +567,7 @@ pid_t raptor_sm_get_pid_for_username(char* username) {
 	// Make sure a server is not already running for this user
 	// Return the existing PID if it is
 	char* safe_username = get_mysql_escaped_string(conn, username);
-	asprintf(&query, "SELECT pid FROM sessions WHERE username='%s'", safe_username);
+	asprintf(&query, "SELECT %s FROM sessions WHERE username='%s'", dbfield, safe_username);
 	free(safe_username);
 	if (mysql_query_internal(conn, query)) {
 		// Server error
@@ -582,7 +595,7 @@ pid_t raptor_sm_get_pid_for_username(char* username) {
 	return -3;
 }
 
-char* raptor_sm_server_started(char* username, pid_t pid, int display) {
+char* raptor_sm_server_started(char* username, pid_t pid, int display, char* dbfield) {
 	MYSQL_RES *res;
 	MYSQL_ROW row;
 	char* query;
@@ -596,7 +609,34 @@ char* raptor_sm_server_started(char* username, pid_t pid, int display) {
 
 	// Update new information into the sessions database and set status to ALLOCATED
 	char* safe_username = get_mysql_escaped_string(conn, username);
-	asprintf(&query, "UPDATE sessions SET pid='%d', stamp_start='%lld', state='%d', display='%d', stamp_statechange='%lld' WHERE username='%s' AND state='%d'", pid, timestamp, SM_STATUS_RUNNING, display, timestamp, safe_username, SM_STATUS_ALLOCATED);
+	asprintf(&query, "UPDATE sessions SET %s='%d', stamp_start='%lld', state='%d', display='%d', stamp_statechange='%lld' WHERE username='%s' AND state='%d'", dbfield, pid, timestamp, SM_STATUS_RUNNING, display, timestamp, safe_username, SM_STATUS_ALLOCATED);
+	free(safe_username);
+	if (mysql_query_internal(conn, query)) {
+		// Server error
+		free(query);
+		mysql_close(conn);
+		return -2;
+	}
+	else {
+		free(query);
+		mysql_close(conn);
+		return 0;
+	}
+}
+
+char* raptor_sm_wm_started(char* username, pid_t pid, char* dbfield) {
+	MYSQL_RES *res;
+	MYSQL_ROW row;
+	char* query;
+
+	MYSQL *conn = connect_if_needed();
+	if (!conn) {
+		return -1;
+	}
+
+	// Update new information into the sessions database and set status to ALLOCATED
+	char* safe_username = get_mysql_escaped_string(conn, username);
+	asprintf(&query, "UPDATE sessions SET %s='%d' WHERE username='%s'", dbfield, pid, safe_username);
 	free(safe_username);
 	if (mysql_query_internal(conn, query)) {
 		// Server error
@@ -666,6 +706,35 @@ void raptor_sm_session_terminated(char* username) {
 	raptor_sm_deallocate_session(username);
 }
 
+void raptor_sm_wm_terminated(char* username) {
+	MYSQL_RES *res;
+	MYSQL_ROW row;
+	char* query;
+
+	long long timestamp = time(NULL);
+
+	MYSQL *conn = connect_if_needed();
+	if (!conn) {
+		return -1;
+	}
+
+	// Update new information into the sessions database
+	char* safe_username = get_mysql_escaped_string(conn, username);
+	asprintf(&query, "UPDATE sessions SET %s=NULL WHERE username='%s'", RAPTOR_SM_WM_PID_FIELD, safe_username);
+	free(safe_username);
+	if (mysql_query_internal(conn, query)) {
+		// Server error
+		free(query);
+		mysql_close(conn);
+		return -2;
+	}
+	else {
+		free(query);
+		mysql_close(conn);
+		return 0;
+	}
+}
+
 int raptor_sm_get_new_unique_display(int mindisplay, int maxdisplay) {
 	MYSQL_RES *res;
 	MYSQL_ROW row;
@@ -740,9 +809,14 @@ void raptor_sm_run_remote_desktop(char* username, int display, char* executable)
 	asprintf(&command_string, "ssh root@%s \"su %s -c \'export DISPLAY=:%d && %s && exit\' &> /dev/null\"", ipaddr, username, display, executable);
 	system(command_string);
 	free(command_string);
+}
+
+void raptor_sm_terminate_server(char* username) {
+	char* ipaddr = raptor_sm_get_ip_for_username(username, true);
+	char* command_string;
 
 	// Terminate remote X server
-	pid_t pid = raptor_sm_get_pid_for_username(username);
+	pid_t pid = raptor_sm_get_pid_for_username(username, RAPTOR_SM_SERVER_PID_FIELD);
 	if (pid > 0) {
 		asprintf(&command_string, "ssh root@%s \'kill -9 %ld\'", ipaddr, pid);
 		system(command_string);
diff --git a/raptorsmiface/libraptorsmiface.h b/raptorsmiface/libraptorsmiface.h
index 1d56a021..232dfe33 100644
--- a/raptorsmiface/libraptorsmiface.h
+++ b/raptorsmiface/libraptorsmiface.h
@@ -20,17 +20,23 @@ enum raptor_sm_status {
 	SM_STATUS_FORCEKILL
 };
 
+#define RAPTOR_SM_SERVER_PID_FIELD "server_pid"
+#define RAPTOR_SM_WM_PID_FIELD "wm_pid"
+
 char* raptor_sm_get_ip_for_hostname(char* hostname, char* err);
 char* raptor_sm_get_hostname_for_username(char* username, bool create);
 
 char* raptor_sm_get_ip_for_username(char* username, bool create);
-pid_t raptor_sm_run_remote_server(char* username, char *const argv[]);
-pid_t raptor_sm_get_pid_for_username(char* username);
-char* raptor_sm_server_started(char* username, pid_t pid, int display);
+pid_t raptor_sm_run_remote_server(char* username, char *const argv[], char* dbfield, int display);
+pid_t raptor_sm_get_pid_for_username(char* username, char* dbfield);
+char* raptor_sm_server_started(char* username, pid_t pid, int display, char* dbfield);
+char* raptor_sm_wm_started(char* username, pid_t pid, char* dbfield);
 int   raptor_sm_get_display_for_username(char* username);
 void  raptor_sm_wait_for_pid_exit(char* username, pid_t pid);
 void  raptor_sm_session_terminated(char* username);
+void  raptor_sm_wm_terminated(char* username);
 int   raptor_sm_get_new_unique_display(int mindisplay, int maxdisplay);
 bool  raptor_sm_sesslimit_reached(char* username);
 char  raptor_sm_set_session_state(int display, int state);
-void  raptor_sm_run_remote_desktop(char* username, int display, char* executable);
\ No newline at end of file
+void  raptor_sm_run_remote_desktop(char* username, int display, char* executable);
+void  raptor_sm_terminate_server(char* username);
\ No newline at end of file
diff --git a/sesman/scp_v0.c b/sesman/scp_v0.c
index 6d6cab33..bd80762d 100644
--- a/sesman/scp_v0.c
+++ b/sesman/scp_v0.c
@@ -79,7 +79,7 @@ scp_v0_process(struct SCP_CONNECTION *c, struct SCP_SESSION *s)
                                     s->bpp, s->type, s->client_ip);
 
         // RAPTOR session management
-        pid_t serverpid = raptor_sm_get_pid_for_username(s->username);
+        pid_t serverpid = raptor_sm_get_pid_for_username(s->username, RAPTOR_SM_SERVER_PID_FIELD);
         if (serverpid < 0) {
             // Session NOT already running
             if (s_item != 0) {
diff --git a/sesman/scp_v1.c b/sesman/scp_v1.c
index c766a2cd..32a7f115 100644
--- a/sesman/scp_v1.c
+++ b/sesman/scp_v1.c
@@ -110,7 +110,7 @@ scp_v1_process(struct SCP_CONNECTION *c, struct SCP_SESSION *s)
     slist = session_get_byuser(s->username, &scount, SESMAN_SESSION_STATUS_DISCONNECTED);
 
     // RAPTOR session management
-    pid_t serverpid = raptor_sm_get_pid_for_username(s->username);
+    pid_t serverpid = raptor_sm_get_pid_for_username(s->username, RAPTOR_SM_SERVER_PID_FIELD);
     if (serverpid < 0) {
         // Session NOT already running
         scount = 0;
diff --git a/sesman/session.c b/sesman/session.c
index f9879b59..5a7b945c 100644
--- a/sesman/session.c
+++ b/sesman/session.c
@@ -654,9 +654,6 @@ session_start_fork(int width, int height, int bpp, char *username,
             }
             else if (pampid == 0) /* child: X11/client */
             {
-                if (session_was_already_running) {
-                    g_exit(0);
-                }
                 char* remote_server = wait_for_remote_hostname(username);
                 wait_for_remote_xserver(remote_server, display);
                 env_set_user(username, 0, display,
@@ -668,8 +665,34 @@ session_start_fork(int width, int height, int bpp, char *username,
                 {
                     g_free(remote_server);
 
-                    // RAPTOR session management
-                    raptor_sm_run_remote_desktop(username, display, "/opt/trinity/bin/starttde");
+//                    // RAPTOR session management
+//                    raptor_sm_run_remote_desktop(username, display, "/opt/trinity/bin/starttde");
+//                    g_exit(0);
+
+                    struct list * wm_params = (struct list *)NULL;
+                    char ** pp2 = (char **)NULL;
+                    wm_params = list_create();
+                    wm_params->auto_free = 1;
+                    /* these are the must have parameters */
+                    list_add_item(wm_params, (long)g_strdup("/opt/trinity/bin/starttde"));
+
+                    /* make sure it ends with a zero */
+                    list_add_item(wm_params, 0);
+                    pp2 = (char**)wm_params->items;
+
+                    pid_t wmpid;
+                    wmpid = raptor_sm_run_remote_server(username, pp2, RAPTOR_SM_WM_PID_FIELD, display);
+                    log_message(&(g_cfg->log), LOG_LEVEL_ALWAYS, "new window manager pid code was %d for user %s", wmpid, username);
+
+                    if (wmpid >= 0) {
+                        raptor_sm_wm_started(username, wmpid, RAPTOR_SM_WM_PID_FIELD);
+
+                        // Wait for PID exit and remove information from the session database
+                        raptor_sm_wait_for_pid_exit(username, wmpid);
+                        raptor_sm_wm_terminated(username);
+                        raptor_sm_terminate_server(username);
+                    }
+
                     g_exit(0);
 
                     auth_set_env(data);
@@ -856,14 +879,14 @@ session_start_fork(int width, int height, int bpp, char *username,
                     log_message(LOG_LEVEL_INFO, "X11rdp start:%s", dumpItemsToString(xserver_params, execvpparams, 2048));
 
                     pid_t serverpid;
-                    serverpid = raptor_sm_run_remote_server(username, pp1);
+                    serverpid = raptor_sm_run_remote_server(username, pp1, RAPTOR_SM_SERVER_PID_FIELD, -1);
                     log_message(&(g_cfg->log), LOG_LEVEL_ALWAYS, "new server pid code was %d during login for user %s", serverpid, username);
 
                     if (serverpid >= 0) {
                         if (!session_was_already_running) {
                             char *friendlyscreen = g_strdup(screen);
                             friendlyscreen[0] = ' ';
-                            raptor_sm_server_started(username, serverpid, atoi(friendlyscreen));
+                            raptor_sm_server_started(username, serverpid, atoi(friendlyscreen), RAPTOR_SM_SERVER_PID_FIELD);
                             g_free(friendlyscreen);
 
                             // Wait for PID exit and remove information from the session database