Further SSH tunnel improvements:

 - Redesigned SSH key authentication dialog
 - Added ability to add SSH key passphrases to keychain (sharing details with system SSH)
 - SSH tunnels with keys which fail are now correctly restarted, interacting with the GUI as necessary
 - GUI interaction now performed on the main thread for increased stability

6 files changed, 91 insertions, 14 deletions

diff --git a/Source/CMMCPConnection.m b/Source/CMMCPConnection.m
index 478d5684..8fbf948d 100644
--- a/Source/CMMCPConnection.m
+++ b/Source/CMMCPConnection.m
@@ -355,7 +355,7 @@ static void forcePingTimeout(int signalNumber);
 		[connectionTunnel setConnectionStateChangeSelector:nil delegate:nil];
 		if ([connectionTunnel state] != SPSSH_STATE_IDLE) [connectionTunnel disconnect];
 		[connectionTunnel connect];
-		NSDate *tunnelStartDate = [NSDate date];
+		NSDate *tunnelStartDate = [NSDate date], *interfaceInteractionTimer;
 		
 		// Allow the tunnel to attempt to connect in a loop
 		while (1) {
@@ -367,7 +367,11 @@ static void forcePingTimeout(int signalNumber);
 				[connectionTunnel disconnect];
 				break;
 			}
-			[NSThread sleepForTimeInterval:0.25];
+			
+			// Process events for a short time, allowing dialogs to be shown but waiting for the tunnel
+			interfaceInteractionTimer = [NSDate date];
+			[[NSRunLoop currentRunLoop] runMode:NSModalPanelRunLoopMode beforeDate:[NSDate dateWithTimeIntervalSinceNow:0.25]];
+			tunnelStartDate = [tunnelStartDate addTimeInterval:([[NSDate date] timeIntervalSinceDate:interfaceInteractionTimer] - 0.25)];
 		}
 		currentSSHTunnelState = [connectionTunnel state];
 		[connectionTunnel setConnectionStateChangeSelector:@selector(sshTunnelStateChange:) delegate:self];
diff --git a/Source/KeyChain.h b/Source/KeyChain.h
index 966a8c04..ca95d09f 100644
--- a/Source/KeyChain.h
+++ b/Source/KeyChain.h
@@ -29,6 +29,7 @@
 @interface KeyChain : NSObject
 
 - (void)addPassword:(NSString *)password forName:(NSString *)name account:(NSString *)account;
+- (void)addPassword:(NSString *)password forName:(NSString *)name account:(NSString *)account withLabel:(NSString *)label;
 - (NSString *)getPasswordForName:(NSString *)name account:(NSString *)account;
 - (void)deletePasswordForName:(NSString *)name account:(NSString *)account;
 - (BOOL)passwordExistsForName:(NSString *)name account:(NSString *)account;
diff --git a/Source/KeyChain.m b/Source/KeyChain.m
index ec5ee03f..ad0251f7 100644
--- a/Source/KeyChain.m
+++ b/Source/KeyChain.m
@@ -34,6 +34,14 @@
  */
 - (void)addPassword:(NSString *)password forName:(NSString *)name account:(NSString *)account
 {
+	[self addPassword:password forName:name account:account withLabel:name];
+}
+
+/**
+ * Add the supplied password to the user's Keychain using the supplied name, account, and label.
+ */
+- (void)addPassword:(NSString *)password forName:(NSString *)name account:(NSString *)account withLabel:(NSString *)label;
+{
 	OSStatus status;
 	SecTrustedApplicationRef sequelProRef, sequelProHelperRef;
 	SecAccessRef passwordAccessRef;
@@ -61,8 +69,8 @@
 		attributes[0].data = "application password";
 		attributes[0].length = 20;
 		attributes[1].tag = kSecLabelItemAttr;
-		attributes[1].data = (unichar *)[name UTF8String];
-		attributes[1].length = strlen([name UTF8String]);
+		attributes[1].data = (unichar *)[label UTF8String];
+		attributes[1].length = strlen([label UTF8String]);
 		attributes[2].tag = kSecAccountItemAttr;
 		attributes[2].data = (unichar *)[account UTF8String];
 		attributes[2].length = strlen([account UTF8String]);
@@ -173,7 +181,7 @@
 	attributes[0].data   = (void *)[account UTF8String];
 	attributes[0].length = [account length];
 	
-	attributes[1].tag    = kSecLabelItemAttr;
+	attributes[1].tag    = kSecServiceItemAttr;
     attributes[1].data   = (void *)[name UTF8String];
     attributes[1].length = [name length];
 	
diff --git a/Source/SPSSHTunnel.h b/Source/SPSSHTunnel.h
index 4486685a..0c34f4f9 100644
--- a/Source/SPSSHTunnel.h
+++ b/Source/SPSSHTunnel.h
@@ -19,6 +19,7 @@ enum spsshtunnel_password_modes
 {
 	IBOutlet NSWindow *sshQuestionDialog;
 	IBOutlet NSTextField *sshQuestionText;
+	IBOutlet NSButton *sshPasswordKeychainCheckbox;
 	IBOutlet NSWindow *sshPasswordDialog;
 	IBOutlet NSTextField *sshPasswordText;
 	IBOutlet NSSecureTextField *sshPasswordField;
@@ -38,6 +39,8 @@ enum spsshtunnel_password_modes
 	NSString *password;
 	NSString *keychainName;
 	NSString *keychainAccount;
+	NSString *requestedPassphrase;
+	BOOL requestedResponse;
 	BOOL passwordInKeychain;
 	int sshPort;
 	int remotePort;
@@ -59,7 +62,9 @@ enum spsshtunnel_password_modes
 - (void) standardErrorHandler:(NSNotification*)aNotification;
 - (NSString *) getPasswordWithVerificationHash:(NSString *)theHash;
 - (BOOL) getResponseForQuestion:(NSString *)theQuestion;
+- (void) workerGetResponseForQuestion:(NSString *)theQuestion;
 - (NSString *) getPasswordForQuery:(NSString *)theQuery verificationHash:(NSString *)theHash;
+- (void) workerGetPasswordForQuery:(NSString *)theQuery;
 - (IBAction) closeSheet:(id)sender;
 
 @end
diff --git a/Source/SPSSHTunnel.m b/Source/SPSSHTunnel.m
index d8332c14..6df20d2e 100644
--- a/Source/SPSSHTunnel.m
+++ b/Source/SPSSHTunnel.m
@@ -22,6 +22,8 @@
 //  More info at <http://code.google.com/p/sequel-pro/>
 
 #import "SPSSHTunnel.h"
+#import "RegexKitLite.h"
+#import "KeyChain.h"
 #import <netinet/in.h>
 
 
@@ -69,6 +71,8 @@
 	keychainName = nil;
 	keychainAccount = nil;
 	passwordInKeychain = NO;
+	requestedPassphrase = nil;
+	requestedResponse = NO;
 	task = nil;
 	localPort = 0;
 	connectionState = SPSSH_STATE_IDLE;
@@ -226,7 +230,7 @@
 //	[taskArguments addObject:@"-C"]; // TODO: compression?
 	[taskArguments addObject:@"-o ExitOnForwardFailure=yes"];
 	[taskArguments addObject:[NSString stringWithFormat:@"-o ConnectTimeout=%i", connectionTimeout]];
-	[taskArguments addObject:@"-o NumberOfPasswordPrompts=1"];
+	[taskArguments addObject:@"-o NumberOfPasswordPrompts=3"];
 	if (useKeepAlive && keepAliveInterval) {
 		[taskArguments addObject:@"-o TCPKeepAlive=no"];		
 		[taskArguments addObject:[NSString stringWithFormat:@"-o ServerAliveInterval=%i", (int)ceil(keepAliveInterval)]];		
@@ -403,6 +407,13 @@
  */
 - (BOOL) getResponseForQuestion:(NSString *)theQuestion
 {
+	[self performSelectorOnMainThread:@selector(workerGetResponseForQuestion:) withObject:theQuestion waitUntilDone:YES];
+	
+	return requestedResponse;
+}
+- (void) workerGetResponseForQuestion:(NSString *)theQuestion
+{	
+
 	NSSize questionTextSize;
 	NSRect windowFrameRect;
 
@@ -410,7 +421,7 @@
 	[sshQuestionText setStringValue:theQuestion];
 	questionTextSize = [[sshQuestionText cell] cellSizeForBounds:NSMakeRect(0, 0, [sshQuestionText bounds].size.width, 500)];
 	windowFrameRect = [sshQuestionDialog frame];
-	windowFrameRect.size.height = ((questionTextSize.height < 100)?100:questionTextSize.height) + 90;
+	windowFrameRect.size.height = ((questionTextSize.height < 100)?100:questionTextSize.height) + 70 + ([sshPasswordDialog isSheet]?0:22);
 	[sshQuestionDialog setFrame:windowFrameRect display:NO];
 	[NSApp beginSheet:sshQuestionDialog modalForWindow:parentWindow modalDelegate:self didEndSelector:nil contextInfo:nil];
 	int sshQueryResponseCode = [NSApp runModalForWindow:sshQuestionDialog];
@@ -421,11 +432,13 @@
 
 		// Yes
 		case 1:
-			return YES;
+			requestedResponse = YES;
+			return;
 
 		// No
 		default:
-			return NO;
+			requestedResponse = NO;
+			return;
 	}
 }
 
@@ -437,15 +450,36 @@
 {
 	if (![theHash isEqualToString:tunnelConnectionVerifyHash]) return nil;
 
+	NSString *thePassword;
+	
+	[self performSelectorOnMainThread:@selector(workerGetPasswordForQuery:) withObject:theQuery waitUntilDone:YES];
+
+	if (!requestedPassphrase) return nil;
+	thePassword = [NSString stringWithString:requestedPassphrase];
+	[requestedPassphrase release], requestedPassphrase = nil;
+	return thePassword;
+}
+- (void) workerGetPasswordForQuery:(NSString *)theQuery
+{
 	NSSize queryTextSize;
 	NSRect windowFrameRect;
 	NSString *thePassword;
+	KeyChain *keychain;
+
+	// Work out whether a passphrase is being requested, extracting the key name
+	NSString *keyName = [theQuery stringByMatching:@"^\\s*Enter passphrase for key \\'(.*)\\':\\s*$" capture:1L];
+	if (keyName) {
+		[sshPasswordText setStringValue:[NSString stringWithFormat:@"Enter your password for the SSH key\n\"%@\"", keyName]];
+		[sshPasswordKeychainCheckbox setHidden:NO];
+	} else {
+		[sshPasswordText setStringValue:theQuery];
+		[sshPasswordKeychainCheckbox setHidden:YES];	
+	}
 
 	// Request the password, sizing the window appropriately to fit the query
-	[sshPasswordText setStringValue:theQuery];
 	queryTextSize = [[sshPasswordText cell] cellSizeForBounds:NSMakeRect(0, 0, [sshPasswordText bounds].size.width, 500)];
 	windowFrameRect = [sshPasswordDialog frame];
-	windowFrameRect.size.height = ((queryTextSize.height < 40)?40:queryTextSize.height) + 143;
+	windowFrameRect.size.height = ((queryTextSize.height < 40)?40:queryTextSize.height) + 140 + ([sshPasswordDialog isSheet]?0:22);
 	[sshPasswordDialog setFrame:windowFrameRect display:NO];
 	[NSApp beginSheet:sshPasswordDialog modalForWindow:parentWindow modalDelegate:self didEndSelector:nil contextInfo:nil];
 	int sshQueryResponseCode = [NSApp runModalForWindow:sshPasswordDialog];
@@ -459,11 +493,19 @@
 			thePassword = [NSString stringWithString:[sshPasswordField stringValue]];
 			[sshPasswordField setStringValue:@""];
 			[[delegate undoManager] removeAllActionsWithTarget:sshPasswordField];
-			return thePassword;
+			requestedPassphrase = [[NSString alloc] initWithString:thePassword];
+			
+			// Add to keychain if appropriate
+			if (keyName && [sshPasswordKeychainCheckbox state] == NSOnState) {
+				keychain = [[KeyChain alloc] init];
+				[keychain addPassword:thePassword forName:@"SSH" account:keyName withLabel:[NSString stringWithFormat:@"SSH: %@", keyName]];
+				[keychain release];
+			}
+			return;
 
 		// Cancel
 		default:
-			return nil;
+			return;
 	}
 }
  
diff --git a/Source/TunnelPassphraseRequester.m b/Source/TunnelPassphraseRequester.m
index 31c4b54a..11665449 100644
--- a/Source/TunnelPassphraseRequester.m
+++ b/Source/TunnelPassphraseRequester.m
@@ -23,6 +23,7 @@
 #import <Cocoa/Cocoa.h>
 #import "KeyChain.h"
 #import "SPSSHTunnel.h"
+#import "RegexKitLite.h"
 
 int main(int argc, const char *argv[])
 {
@@ -120,9 +121,25 @@ int main(int argc, const char *argv[])
 		}
 	}
 
-	// Check whether we're being asked for a SSH key passphrase, forward requests to the GUI
+	// Check whether we're being asked for a SSH key passphrase
 	if (argument && [[argument lowercaseString] rangeOfString:@"enter passphrase for"].location != NSNotFound ) {
 		NSString *passphrase;
+		NSString *keyName = [argument stringByMatching:@"^\\s*Enter passphrase for key \\'(.*)\\':\\s*$" capture:1L];
+
+		if (keyName) {
+		
+			// Check whether the passphrase is in the keychain, using standard OS X sshagent name and account
+			KeyChain *keychain = [[KeyChain alloc] init];
+			if ([keychain passwordExistsForName:@"SSH" account:keyName]) {
+				printf("%s\n", [[keychain getPasswordForName:@"SSH" account:keyName] UTF8String]);
+				[keychain release];
+				[pool release];
+				return 0;
+			}
+			[keychain release];
+		}
+		
+		// Not found in the keychain - we need to ask the GUI.
 
 		if (!verificationHash) {
 			NSLog(@"SSH Tunnel: key passphrase authentication required but insufficient details supplied to connect to GUI");