1 files changed, 78 insertions, 33 deletions

diff --git a/src/tmate/session.vala b/src/tmate/session.vala
index 5571654..356aa66 100644
--- a/src/tmate/session.vala
+++ b/src/tmate/session.vala
@@ -2,6 +2,8 @@ namespace Tmate
 {
 	public interface SessionInterface : Object
 	{
+		public abstract signal void address (SessionType session, string address);
+
 		public abstract string? @get(SessionType session);
 		public abstract bool	start(string? config = null);
 		public abstract bool	stop();
@@ -16,6 +18,7 @@ namespace Tmate
 			private string? _addr_http_ro = null;
 			private IOChannel? stdout     = null;
 			private Pid?       pid        = null;
+			private uint error_count = 0;
 
 			// Regex
 			private Regex r_connect = /^Connecting to .+\.\.\.$/;
@@ -26,9 +29,14 @@ namespace Tmate
 			private Regex r_ssh_ro  = /^ssh \w+ read \w+[:] .+ (.+)+$/;
 			private Regex r_joined  = /^.+ joined \(([^()]+)\) -- (\d+) clients? .+$/;
 			private Regex r_left    = /^.+ left \(([^()]+)\) -- (\d+) clients? .+$/;
+			private Regex r_closed  = /^Session closed$/;
+			private Regex r_note    = /^(Note: |Reconnecting)/;
+			private Regex r_restarted = /^Session shell restarted$/;
+			private Regex r_netfail= /^[^ ]+ lookup failure\..+ \((.+)\)$/;
 
 			public SessionType flags  { get; private set;}
 			public string?     socket { get; private set;}
+			public bool        terminate = true;
 
 			construct {	reset(); }
 
@@ -36,32 +44,47 @@ namespace Tmate
 			{
 				info (_("Session: unix://%s"), socket);
 			}
+
+			public virtual signal void network_error(string message)
+			{
+				if(error_count ++ > 5) stop();
+			}
+
 			public virtual signal void stopped ()
 			{
 				info (_("Session terminated."));
+				reset();
 			}
-			public override signal void address (SessionType session, string address)
+
+			private void address_set(SessionType session, string address)
 			{
 				info (_("%s address: %s"), session.to_string(), address);
-				switch(session) {
-				case SSH :
-					_addr_ssh  = address;
-					break;
-				case HTTP:
-					_addr_http = address;
-					break;
-				case HTTP_READONLY:
-					_addr_http_ro = address;
-					break;
-				case SSH_READONLY :
-					_addr_ssh_ro  = address;
-					break;
-				case DISCONNECTED :
+
+				if (session in SessionType.SSH|SessionType.SSH_READONLY|SessionType.HTTP|SessionType.HTTP_READONLY) {
+					switch(session) {
+					case SSH:
+						_addr_ssh  = address;
+						break;
+					case HTTP:
+						_addr_http = address;
+						break;
+					case HTTP_READONLY:
+						_addr_http_ro = address;
+						break;
+					case SSH_READONLY :
+						_addr_ssh_ro  = address;
+						break;
+					default:
+						warn_if_reached();
+						break;
+					}
+
+					flags |= session;
+					this.address(session, address);
+				} else if (session == DISCONNECTED) {
 					_addr_http = _addr_ssh = _addr_http_ro = _addr_ssh_ro = null;
-					break;
-				default:
-					warn_if_reached();
-					break;
+					flags = DISCONNECTED;
+					this.address(DISCONNECTED, "");
 				}
 			}
 			public virtual signal void joined (string mate, int mates)
@@ -83,6 +106,7 @@ namespace Tmate
 				pid           = null;
 				socket        = null;
 				flags         = DISCONNECTED;
+				error_count   = 0;
 			}
 
 			private bool send(string[] args)
@@ -129,10 +153,29 @@ namespace Tmate
 				return null;
 			}
 
+			public bool stop()
+			{
+				return send({"kill-server"});
+			}
+
+			private void restart()
+			{
+				if(terminate) {
+					stop();
+					return;
+				}
+
+				started(socket);
+				if((bool)_addr_ssh) address(SSH, _addr_ssh);
+				if((bool)_addr_ssh_ro) address(SSH_READONLY, _addr_ssh_ro);
+				if((bool)_addr_http) address(HTTP, _addr_http);
+				if((bool)_addr_http_ro) address(HTTP_READONLY, _addr_http_ro);
+			}
+
 			public bool start(string? config = null)
 			{
 				if(! (pid == null)) {
-					started.emit(socket);
+					started(socket);
 					return true;
 				}
 
@@ -164,15 +207,22 @@ namespace Tmate
 							if(r_socket.match(line, 0, out info))
 								socket = info.fetch(1);
 							else if (r_connect.match(line) && socket != null)
-								started.emit(socket);
+								started(socket);
 							else if(r_http.match(line, 0, out info))
-								address.emit(HTTP, info.fetch(1));
+								address_set(HTTP, info.fetch(1));
 							else if(r_http_ro.match(line, 0, out info))
-								address.emit(HTTP_READONLY, info.fetch(1));
+								address_set(HTTP_READONLY, info.fetch(1));
 							else if(r_ssh.match(line, 0, out info))
-								address.emit(SSH, info.fetch(1));
+								address_set(SSH, info.fetch(1));
 							else if(r_ssh_ro.match(line, 0, out info))
-								address.emit(SSH_READONLY, info.fetch(1));
+								address_set(SSH_READONLY, info.fetch(1));
+							else if(r_restarted.match(line))
+								restart();
+							else if(r_closed.match(line))
+								stopped();
+							else if(r_netfail.match(line, 0, out info))
+								network_error(info.fetch(1));
+							else if(r_note.match(line)) unlikely(false); // Ignore line
 							else
 								debug("Unprocessed line: \"%s\"", line);
 						} catch (IOChannelError e) {
@@ -187,21 +237,16 @@ namespace Tmate
 
 					ChildWatch.add(pid, (pid, status) => {
 						Process.close_pid(pid);
-						reset();
-						stopped.emit();
+						if(flags != DISCONNECTED)
+							stopped();
 					});
 
 					return true;
 				} catch (SpawnError e) {
 					warning("Error: %s", e.message);
-					reset();
+					stopped();
 					return false;
 				}
 			}
-
-			public bool stop()
-			{
-				return false;
-			}
 	}
 }