Add server/group mapping

2 files changed, 97 insertions, 48 deletions

diff --git a/raptorsmiface/libraptorsmiface.c b/raptorsmiface/libraptorsmiface.c
index 933ea98e..ead1fd06 100644
--- a/raptorsmiface/libraptorsmiface.c
+++ b/raptorsmiface/libraptorsmiface.c
@@ -125,13 +125,15 @@ char* raptor_sm_allocate_session(char* username) {
 	MYSQL_ROW row;
 	MYSQL_RES *svr_res;
 	MYSQL_ROW svr_row;
+	MYSQL_RES *per_res;
+	MYSQL_ROW per_row;
 	MYSQL_RES *cnt_res;
 	MYSQL_ROW cnt_row;
 	char* query;
 
 	MYSQL *conn = connect_if_needed();
 	if (!conn) {
-		return strdup("SQLERR001");
+		return strdup("ERROR");
 	}
 
 	// Verify that this user is not already on the system
@@ -142,7 +144,7 @@ char* raptor_sm_allocate_session(char* username) {
 		// Server error
 		free(query);
 		mysql_close(conn);
-		return strdup("SQLERR002");
+		return strdup("ERROR");
 	}
 	else {
 		free(query);
@@ -154,59 +156,95 @@ char* raptor_sm_allocate_session(char* username) {
 				// Server error
 				mysql_free_result(res);
 				mysql_close(conn);
-				return strdup("SQLERR003");
+				return strdup("ERROR");
 			}
 			else {
 				svr_res = mysql_store_result(conn);
-				char* bestserver = strdup("");
-				int bestusage = INT_MAX;
-				while ((svr_row = mysql_fetch_row(svr_res)) != NULL) {
-					char* safe_servername = get_mysql_escaped_string(conn, svr_row[0]);
-					asprintf(&query, "SELECT username FROM sessions WHERE servername='%s'", safe_servername);
-					free(safe_servername);
-					if (mysql_query_internal(conn, query)) {
-						// Server error
-						free(query);
-						free(bestserver);
-						mysql_free_result(res);
-						mysql_free_result(svr_res);
-						mysql_close(conn);
-						return strdup("SQLERR004");
-					}
-					else {
-						free(query);
-						cnt_res = mysql_store_result(conn);
-						int usagecount = 0;
-						while ((cnt_row = mysql_fetch_row(cnt_res)) != NULL) {
-							usagecount++;
-						}
-						mysql_free_result(cnt_res);
-						if (usagecount < bestusage) {
-							free(bestserver);
-							bestserver = strdup(svr_row[0]);
-							bestusage = usagecount;
-						}
-					}
-				}
-				mysql_free_result(res);
-				mysql_free_result(svr_res);
-
-				// Insert new information into the sessions database and set status to ALLOCATED
-				char* safe_servername = get_mysql_escaped_string(conn, bestserver);
-				char* safe_username = get_mysql_escaped_string(conn, username);
-				asprintf(&query, "INSERT INTO sessions (username, servername, state) VALUES ('%s', '%s', '%d')", safe_username, safe_servername, SM_STATUS_ALLOCATED);
-				free(safe_servername);
-				free(safe_username);
+
+				// Get group for user
+				char* groupname = get_group_for_user(username);
+				char* safe_groupname = get_mysql_escaped_string(conn, groupname);
+				free(groupname);
+				// Get the list of allowed nodes for this group
+				asprintf(&query, "SELECT server FROM allowed_servers WHERE groupname='%s'", safe_groupname);
+				free(safe_groupname);
 				if (mysql_query_internal(conn, query)) {
 					// Server error
-					free(query);
+					mysql_free_result(res);
+					mysql_free_result(svr_res);
 					mysql_close(conn);
-					return strdup("SQLERR005");
+					return strdup("ERROR");
 				}
 				else {
-					free(query);
-					mysql_close(conn);
-					return strdup(bestserver);
+					per_res = mysql_store_result(conn);
+					char* bestserver = strdup("");
+					int bestusage = INT_MAX;
+					while ((svr_row = mysql_fetch_row(svr_res)) != NULL) {
+						// Am I allowed to use this server?
+						bool can_use_server = false;
+						while ((per_row = mysql_fetch_row(per_res)) != NULL) {
+							if (strcmp(per_row[0], svr_row[0]) == 0) {
+								can_use_server = true;
+							}
+						}
+						mysql_data_seek(per_res, 0);
+						if (can_use_server) {
+							char* safe_servername = get_mysql_escaped_string(conn, svr_row[0]);
+							asprintf(&query, "SELECT username FROM sessions WHERE servername='%s'", safe_servername);
+							free(safe_servername);
+							if (mysql_query_internal(conn, query)) {
+								// Server error
+								free(query);
+								free(bestserver);
+								mysql_free_result(res);
+								mysql_free_result(svr_res);
+								mysql_close(conn);
+								return strdup("ERROR");
+							}
+							else {
+								free(query);
+								cnt_res = mysql_store_result(conn);
+								int usagecount = 0;
+								while ((cnt_row = mysql_fetch_row(cnt_res)) != NULL) {
+									usagecount++;
+								}
+								mysql_free_result(cnt_res);
+								if (usagecount < bestusage) {
+									free(bestserver);
+									bestserver = strdup(svr_row[0]);
+									bestusage = usagecount;
+								}
+							}
+						}
+					}
+					mysql_free_result(res);
+					mysql_free_result(svr_res);
+					mysql_free_result(per_res);
+	
+					if (strcmp(bestserver, "") != 0) {
+						// Insert new information into the sessions database and set status to ALLOCATED
+						char* safe_servername = get_mysql_escaped_string(conn, bestserver);
+						char* safe_username = get_mysql_escaped_string(conn, username);
+						asprintf(&query, "INSERT INTO sessions (username, servername, state) VALUES ('%s', '%s', '%d')", safe_username, safe_servername, SM_STATUS_ALLOCATED);
+						free(safe_servername);
+						free(safe_username);
+						if (mysql_query_internal(conn, query)) {
+							// Server error
+							free(query);
+							mysql_close(conn);
+							return strdup("ERROR");
+						}
+						else {
+							free(query);
+							mysql_close(conn);
+							return strdup(bestserver);
+						}
+					}
+					else {
+						// No usable server found!
+						mysql_close(conn);
+						return strdup("ERROR");
+					}
 				}
 			}
 		}
@@ -287,10 +325,11 @@ char* raptor_sm_get_ip_for_username(char* username, bool create) {
 	char* hostname = raptor_sm_get_hostname_for_username(username, create);
 	char err;
 	char* ip = raptor_sm_get_ip_for_hostname(hostname, &err);
+	free(hostname);
 	if (err) {
 		raptor_sm_deallocate_session(username);
+		return strdup("ERROR");
 	}
-	free(hostname);
 	return ip;
 }
 
diff --git a/xrdp/xrdp_mm.c b/xrdp/xrdp_mm.c
index f19eaada..1ffa197a 100644
--- a/xrdp/xrdp_mm.c
+++ b/xrdp/xrdp_mm.c
@@ -522,6 +522,15 @@ xrdp_mm_setup_mod2(struct xrdp_mm *self)
             else if (self->code == 10 || self->code == 20) /* X11rdp/Xorg */
             {
                  char* rsmip = raptor_sm_get_ip_for_username(self->login_username, true);
+                 if (strcmp(rsmip, "ERROR") == 0) {
+                     g_snprintf(raptortext, 255, "[LICENSE] Instantaneous limit exceeded.");
+                     xrdp_wm_log_msg(self->wm, raptortext);
+                     g_snprintf(raptortext, 255, "[LICENSE] Login for user %s denied.", self->login_username);
+                     xrdp_wm_log_msg(self->wm, raptortext);
+                     g_free(rsmip);
+                     raptor_sm_session_terminated(self->login_username);
+                     return 1;
+                 }
                  int allocdisplay = raptor_sm_get_display_for_username(self->login_username);
                  if ((raptor_sm_sesslimit_reached(self->login_username)) && (allocdisplay < 0)) {
                      g_snprintf(raptortext, 255, "[LICENSE] Maximum concurrent session");
@@ -530,6 +539,7 @@ xrdp_mm_setup_mod2(struct xrdp_mm *self)
                      xrdp_wm_log_msg(self->wm, raptortext);
                      g_snprintf(raptortext, 255, "[LICENSE] Login for user %s denied.", self->login_username);
                      xrdp_wm_log_msg(self->wm, raptortext);
+                     g_free(rsmip);
                      raptor_sm_session_terminated(self->login_username);
                      return 1;
                  }