some code cleanup and added tagging of output in SSH mode (not completly correct yet)

2 files changed, 140 insertions, 87 deletions

diff --git a/src/cssh.c b/src/cssh.c
index b9b1917..2281eea 100644
--- a/src/cssh.c
+++ b/src/cssh.c
@@ -16,6 +16,7 @@
 #include <inttypes.h>
 #include <unistd.h>
 #include <netdb.h>
+#include <assert.h>
 
 #include "port.h"
 #include "progressbar.h"
@@ -45,6 +46,16 @@ typedef enum execution_mode_e {
 	CSSH_EXECUTE_AS_SCP	
 } execution_mode_e;
 
+typedef struct ssh_data_t {
+	ssh_channel channel;
+	char *stdout_buf;
+	char *stdout_bptr;
+	char *stderr_buf;
+	char *stderr_bptr;
+	bool eof_sent;
+	bool is_eof;
+} ssh_data_t;
+
 typedef enum copy_direction_e {
 	CSSH_COPY_DIRECTION_UPLOAD,
 	CSSH_COPY_DIRECTION_DOWNLOAD
@@ -422,18 +433,29 @@ static int read_hosts_file( const char *hosts_file, unsigned short default_port,
 }
 
 // TODO: cleanup also progressbarpool
-static void cleanup_sessions( ssh_session **session, ssh_channel **channel, scp_data_t **scp_data, char **host, unsigned short *port, const int nof_sessions, bool verbose )
+static void cleanup_sessions( ssh_session **session, ssh_data_t **ssh_data, scp_data_t **scp_data, char **host, unsigned short *port, const int nof_sessions, bool verbose )
 {
 	for( unsigned int i = 0; i < nof_sessions; i++ ) {
 		if( ssh_is_connected( (*session)[i] ) ) {
 			if( verbose ) {
 				fprintf( stderr, "Disconnecting from '%s', port %d..\n", host[i], port[i] );
 			}
-			if( channel != NULL ) {	
-				if( ssh_channel_is_open( (*channel)[i] ) ) {
-					ssh_channel_send_eof( (*channel)[i] );
-					ssh_channel_close( (*channel)[i] );
-					ssh_channel_free( (*channel)[i] );	
+			if( ssh_data != NULL ) {
+				ssh_channel *channel = &(*ssh_data)[i].channel;
+				if( ssh_channel_is_open( *channel ) ) {
+					ssh_channel_send_eof( *channel );
+					ssh_channel_close( *channel );
+					ssh_channel_free( *channel );	
+				}
+				char *stdout_buf = (*ssh_data)[i].stdout_buf;
+				if( stdout_buf != NULL ) {
+					free( stdout_buf );
+					stdout_buf = NULL;
+				}
+				char *stderr_buf = (*ssh_data)[i].stderr_buf;
+				if( stderr_buf != NULL ) {
+					free( stderr_buf );
+					stderr_buf = NULL;
 				}
 			}
 			if( scp_data != NULL ) {
@@ -467,8 +489,8 @@ static void cleanup_sessions( ssh_session **session, ssh_channel **channel, scp_
 		}
 		ssh_free( (*session)[i] );
 	}
-	if( channel != NULL ) {
-		free( *channel );
+	if( ssh_data != NULL ) {
+		free( *ssh_data );
 	}
 	if( scp_data != NULL ) {
 		free( *scp_data );
@@ -478,6 +500,47 @@ static void cleanup_sessions( ssh_session **session, ssh_channel **channel, scp_
 	free( port );
 }
 
+static const char *buffer_contains_a_line( const char *buffer, size_t bufsize )
+{
+	for( const char *p = buffer; ( p < buffer + bufsize ) && *p != '\0'; p++ ) {
+		if( *p == '\n' ) {
+			return p;
+		}
+	}
+	return NULL;
+}
+
+static ssize_t output_buffer( FILE *f, const char *host, const char *buffer, const size_t bufsize, const bool tagging )
+{	
+	const char *b = buffer;
+	size_t bsize = bufsize;
+	ssize_t nwritten = 0;
+	const char *p;
+	while( ( p < buffer + bufsize ) && ( p = buffer_contains_a_line( b, bsize ) ) != NULL) {
+		if( tagging ) {
+			fprintf( f, "[%s]: ", host );
+		}
+		
+		ssize_t n = p - b;
+		ssize_t wrc = fwrite( b, 1, n, f );
+		if( wrc < 0 ) {
+			return wrc;
+		}
+		if( wrc != n ) {
+			fprintf( stderr, "ERROR: Buffer write mismatch on %s (%zu != %d)\n",
+				( f == stdout ) ? "stdout" : "stderr",
+				wrc, n );
+			return -1;
+		}
+		nwritten += n + 1;
+		b = p + 1;	
+		
+		fprintf( f, "\n" );
+	}
+	
+	return nwritten;
+}
+
 int main( int argc, char *argv[] )
 {
 	struct gengetopt_args_info args_info;
@@ -744,6 +807,8 @@ int main( int argc, char *argv[] )
 		
 		cssh_msleep( 10 );
 	}
+	
+	free( is_connected );
 
 	// authentication phase
 	
@@ -868,170 +933,154 @@ int main( int argc, char *argv[] )
 		case CSSH_EXECUTE_AS_SSH: {
 
 			// explicit low-level handling of channels in SSH mode
-			ssh_channel *channel = (ssh_channel *)malloc( ( nof_sessions + 1 ) * sizeof( ssh_channel ) );
-			memset( channel, 0, nof_sessions + 1 );
-			if( channel == NULL ) {
-				fprintf( stderr, "ERROR: Memory allocation failed for ssh_channels" );
+
+			ssh_data_t *ssh_data = (ssh_data_t *)malloc( ( nof_sessions + 1 ) * sizeof( ssh_data_t ) );
+			if( ssh_data == NULL ) {
+				fprintf( stderr, "ERROR: Memory allocation failed for ssh_data array" );
 				cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 				exit( EXIT_FAILURE );
 			}
+			memset( ssh_data, 0, ( nof_sessions + 1 ) * sizeof( ssh_data_t ) );
+			
 			for( unsigned int i = 0; i < nof_sessions; i++ ) {
-				channel[i] = ssh_channel_new( session[i] );
-				if( channel[i] == NULL ) {
+				ssh_data[i].channel = ssh_channel_new( session[i] );
+				if( ssh_data[i].channel == NULL ) {
 					fprintf( stderr, "ERROR: Unable to open SSH channel: %s\n",
 						ssh_get_error( session[i] ) );
-					cleanup_sessions( &session, &channel, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+					cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 					exit( EXIT_FAILURE );
 				}
 			}
 
 			for( unsigned int i = 0; i < nof_sessions; i++ ) {
-				rc = ssh_channel_open_session( channel[i] );
+				rc = ssh_channel_open_session( ssh_data[i].channel );
 				if( rc != SSH_OK ) {
-					cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+					cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 					exit( EXIT_FAILURE );
 				}
 			}
 
 			// execute command on all channels
 			for( unsigned int i = 0; i < nof_sessions; i++ ) {
-				rc = ssh_channel_request_exec( channel[i], cmd );
+				rc = ssh_channel_request_exec( ssh_data[i].channel, cmd );
 				if( rc != SSH_OK ) {
 					fprintf( stderr, "ERROR: Executing SSH command '%s' failed: %s\n",
 						cmd, ssh_get_error( session[i] ) );
-					cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+					cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 					exit( EXIT_FAILURE );
 				}
 			}
 
 			bool all_eof = false;
-			bool *eof_sent = (bool *)malloc( ( nof_sessions + 1 ) * sizeof( bool ) );
-			if( eof_sent == NULL ) {
-				fprintf( stderr, "ERROR: Memory allocation failed for 'eof_sent'" );
-				cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
-				exit( EXIT_FAILURE );
-			}
-			memset( eof_sent, false, ( nof_sessions + 1 ) * sizeof( bool ) );
-			bool *is_eof = (bool *)malloc( ( nof_sessions + 1 ) * sizeof( bool ) );
-			if( is_eof == NULL ) {
-				fprintf( stderr, "ERROR: Memory allocation failed for 'is_eof'" );
-				free( eof_sent );
-				cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
-				exit( EXIT_FAILURE );
+
+			for( unsigned int i = 0; i < nof_sessions; i++ ) {
+				ssh_data[i].eof_sent = false;
+				ssh_data[i].is_eof = false;
+				ssh_data[i].stdout_buf = (char *)malloc( BUFSIZE );
+				if( ssh_data[i].stdout_buf == NULL ) {
+					fprintf( stderr, "ERROR: Memory allocation failed for receiving buffers of ssh_channels" );
+					cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+					exit( EXIT_FAILURE );
+				}
+				ssh_data[i].stderr_buf = (char *)malloc( BUFSIZE );
+				if( ssh_data[i].stderr_buf == NULL ) {
+					fprintf( stderr, "ERROR: Memory allocation failed for receiving buffers of ssh_channels" );
+					cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+					exit( EXIT_FAILURE );
+				}
+				ssh_data[i].stdout_bptr = ssh_data[i].stdout_buf;
+				ssh_data[i].stderr_bptr = ssh_data[i].stderr_buf;
 			}
-			memset( is_eof, false, ( nof_sessions + 1 ) * sizeof( bool ) );
 			
 			while( !all_eof ) {
 
 				all_eof = true;
 				for( unsigned int i = 0; i < nof_sessions; i++ ) {
-					if( !ssh_channel_is_closed( channel[i] ) && !ssh_channel_is_eof( channel[i] ) ) {
+					if( !ssh_channel_is_closed( ssh_data[i].channel ) && !ssh_channel_is_eof( ssh_data[i].channel ) ) {
 						all_eof = false;
 					}
 				}
 				
 				// no stdin sent to commands on remote machines for now
 				for( unsigned int i = 0; i < nof_sessions; i++ ) {
-					if( is_eof[i] && !eof_sent[i] ) {
-						ssh_channel_send_eof( channel[i] );
-						eof_sent[i] = true;
+					if( ssh_data[i].is_eof && !ssh_data[i].eof_sent ) {
+						ssh_channel_send_eof( ssh_data[i].channel );
+						ssh_data[i].eof_sent = true;
 						continue;
 					}
 				}
 				
 				for( unsigned int i = 0; i < nof_sessions; i++ ) {
-					char buffer[BUFSIZE];
-					rc = ssh_channel_poll( channel[i], 0 );
+					
+					// handle stdout
+					rc = ssh_channel_poll( ssh_data[i].channel, 0 );
 					if( rc == SSH_ERROR ) {
 						fprintf( stderr, "ERROR: ssh_channel_poll on stdout failed: %s\n",
 							ssh_get_error( session[i] ) );
-						free( eof_sent );
-						free( is_eof );
-						cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+						cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 						exit( EXIT_FAILURE );
 					}
 										
 					if( rc > 0 ) {
-						unsigned int nread = ssh_channel_read_nonblocking( channel[i], buffer, sizeof( buffer ), 0 );
+						unsigned int nread = ssh_channel_read_nonblocking( ssh_data[i].channel, ssh_data[i].stdout_bptr, BUFSIZE - ( ssh_data[i].stdout_bptr - ssh_data[i].stdout_buf ), 0 );
 						if( nread == SSH_ERROR ) {
 							fprintf( stderr, "ERROR: ssh_channel_read_nonblocking on stdout failed: %s\n",
 								ssh_get_error( session[i] ) );
-							free( eof_sent );
-							free( is_eof );
-							cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+							cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 							exit( EXIT_FAILURE );
 						}
 						
 						if( nread == 0 ) {
-							is_eof[i] = true;
+							ssh_data[i].is_eof = true;
 						}
 
 						if( nread > 0 ) {
-							size_t wrc = fwrite( buffer, 1, nread, stdout );
+							ssize_t wrc = output_buffer( stdout, host[i], ssh_data[i].stdout_buf, nread + ( ssh_data[i].stdout_bptr - ssh_data[i].stdout_buf ), args_info.tagging_given );
 							if( wrc < 0 ) {
 								fprintf( stderr, "ERROR: while writing to stdout: %s\n",
 									strerror( errno ) );
-								free( eof_sent );
-								free( is_eof );
-								cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
-								exit( EXIT_FAILURE );
-							}
-
-							if( wrc != nread ) {
-								fprintf( stderr, "ERROR: Write mismatch on stdout (%zu != %d)\n",
-									wrc, nread );
-								free( eof_sent );
-								free( is_eof );
-								cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+								cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 								exit( EXIT_FAILURE );
+							} else if( wrc < nread ) {
+								memmove( ssh_data[i].stdout_buf, ssh_data[i].stdout_bptr + wrc, nread - wrc );
+								ssh_data[i].stdout_bptr = ssh_data[i].stdout_buf + nread - wrc + 1;
 							}
 						}
 					}
-
-					rc = ssh_channel_poll( channel[i], 1 );
+					
+					// handle stderr
+					rc = ssh_channel_poll( ssh_data[i].channel, 1 );
 					if( rc == SSH_ERROR ) {
 						fprintf( stderr, "ERROR: ssh_channel_poll on stderr failed: %s\n",
 							ssh_get_error( session[i] ) );
-						free( eof_sent );
-						free( is_eof );
-						cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+						cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 						exit( EXIT_FAILURE );
 					}
 							
 					if( rc > 0 ) {
-						unsigned int nread = ssh_channel_read_nonblocking( channel[i], buffer, sizeof( buffer ), 1 );
+						unsigned int nread = ssh_channel_read_nonblocking( ssh_data[i].channel, ssh_data[i].stderr_bptr, BUFSIZE - ( ssh_data[i].stdout_bptr - ssh_data[i].stderr_buf ), 1 );
 						if( nread == SSH_ERROR ) {
 							fprintf( stderr, "ERROR: ssh_channel_read_nonblocking on stderr failed: %s\n",
 								ssh_get_error( session[i] ) );
-							free( eof_sent );
-							free( is_eof );
-							cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+							cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 							exit( EXIT_FAILURE );
 						}
 
 						if( nread > 0 ) {
-							size_t wrc = fwrite( buffer, 1, nread, stderr );
+							ssize_t wrc = output_buffer( stderr, host[i], ssh_data[i].stderr_buf, nread + ( ssh_data[i].stdout_bptr - ssh_data[i].stderr_buf ) , args_info.tagging_given );
 							if( wrc < 0 ) {
 								fprintf( stderr, "ERROR: while writting to stderr: %s\n",
 									strerror( errno ) );
-								free( eof_sent );
-								free( is_eof );
-								cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
-								exit( EXIT_FAILURE );
-							}
-
-							if( wrc != nread ) {
-								fprintf( stderr, "ERROR: Write mismatch on stderr (%zu != %d)\n",
-									wrc, nread );
-								free( eof_sent );
-								free( is_eof );
-								cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+								cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 								exit( EXIT_FAILURE );
+							} else if( wrc < nread ) {
+								memmove( ssh_data[i].stderr_buf, ssh_data[i].stderr_bptr + wrc, nread - wrc );
+								ssh_data[i].stderr_bptr += nread - wrc + 1;
 							}
 						}
 
 						if( nread == 0 ) {
-							is_eof[i] = true;
+							ssh_data[i].is_eof = true;
 						}
 					}
 				}
@@ -1039,9 +1088,7 @@ int main( int argc, char *argv[] )
 				cssh_msleep( 1 );
 			}
 			
-			free( eof_sent );
-			free( is_eof );
-			cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
+			cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 
 			} break;
 
@@ -1049,7 +1096,7 @@ int main( int argc, char *argv[] )
 			
 			scp_data_t *scp_data = (scp_data_t *)malloc( ( nof_sessions + 1 ) * sizeof( scp_data_t ) );
 			if( scp_data == NULL ) {
-				fprintf( stderr, "ERROR: Memory allocation failed for ssh_data array" );
+				fprintf( stderr, "ERROR: Memory allocation failed for scp_data array" );
 				cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 );
 				exit( EXIT_FAILURE );
 			}
diff --git a/src/options.ggo.in b/src/options.ggo.in
index 8a90a72..0f003ea 100644
--- a/src/options.ggo.in
+++ b/src/options.ggo.in
@@ -29,6 +29,12 @@ you can specify a port immediatelly after a colon \
 after the hostname. Example 'user@host:port'"
 		string typestr="hosts-file"
 		optional
+
+section "SSH options"
+
+	option "tagging" -
+		"Enable tagging of each line with the name of the host"
+		optional
 		
 section "SCP options"