nonblocking support in dh_handshake() from client.c

some packet nonblocking fixes.
reenable sftp from the sample client.


git-svn-id: svn+ssh://svn.berlios.de/svnroot/repos/libssh/trunk@52 7dcaeef0-15fb-0310-b436-a5af3365683c

2 files changed, 123 insertions, 69 deletions

diff --git a/libssh/client.c b/libssh/client.c
index 1ca8952d..437edfd2 100644
--- a/libssh/client.c
+++ b/libssh/client.c
@@ -104,62 +104,100 @@ int ssh_send_banner(SSH_SESSION *session,int server){
     return 0;
 }
 
-
+#define DH_STATE_INIT 0
+#define DH_STATE_INIT_TO_SEND 1
+#define DH_STATE_INIT_SENT 2
+#define DH_STATE_NEWKEYS_TO_SEND 3
+#define DH_STATE_NEWKEYS_SENT 4
+#define DH_STATE_FINISHED 5
 static int dh_handshake(SSH_SESSION *session){
     STRING *e,*f,*pubkey,*signature;
-    packet_clear_out(session);
-    buffer_add_u8(session->out_buffer,SSH2_MSG_KEXDH_INIT);
-    dh_generate_x(session);
-    dh_generate_e(session);
-    e=dh_get_e(session);
-    buffer_add_ssh_string(session->out_buffer,e);
-    packet_send(session);
-    free(e);
-    if(packet_wait(session,SSH2_MSG_KEXDH_REPLY,1))
-        return -1;
-    pubkey=buffer_get_ssh_string(session->in_buffer);
-    if(!pubkey){
-        ssh_set_error(session,SSH_FATAL,"No public key in packet");
-        return -1;
-    }
-    dh_import_pubkey(session,pubkey);
-    f=buffer_get_ssh_string(session->in_buffer);
-    if(!f){
-        ssh_set_error(session,SSH_FATAL,"No F number in packet");
-        return -1;
-    }
-    dh_import_f(session,f);
-    free(f);
-    if(!(signature=buffer_get_ssh_string(session->in_buffer))){
-        ssh_set_error(session,SSH_FATAL,"No signature in packet");
-        return -1;
-    }
-
-    dh_build_k(session);
-    packet_wait(session,SSH2_MSG_NEWKEYS,1);
-    ssh_say(2,"Got SSH_MSG_NEWKEYS\n");
-    packet_clear_out(session);
-    buffer_add_u8(session->out_buffer,SSH2_MSG_NEWKEYS);
-    packet_send(session);
-    ssh_say(2,"SSH_MSG_NEWKEYS sent\n");
-    make_sessionid(session);
-    /* set the cryptographic functions for the next crypto (it is needed for generate_session_keys for key lenghts) */
-    if(crypt_set_algorithms(session))
-        return -1;
-    generate_session_keys(session);
-    /* verify the host's signature. XXX do it sooner */
-    if(signature_verify(session,signature)){
-        free(signature);
-        return -1;
+    int ret;
+    switch(session->dh_handshake_state){
+        case DH_STATE_INIT:
+            packet_clear_out(session);
+            buffer_add_u8(session->out_buffer,SSH2_MSG_KEXDH_INIT);
+            dh_generate_x(session);
+            dh_generate_e(session);
+            e=dh_get_e(session);
+            buffer_add_ssh_string(session->out_buffer,e);
+            ret=packet_send(session);
+            free(e);
+            session->dh_handshake_state=DH_STATE_INIT_TO_SEND;
+            if(ret==SSH_ERROR)
+                return ret;
+        case DH_STATE_INIT_TO_SEND:
+            ret=packet_flush(session,0);
+            if(ret!=SSH_OK)
+                return ret; // SSH_ERROR or SSH_AGAIN
+            session->dh_handshake_state=DH_STATE_INIT_SENT;
+        case DH_STATE_INIT_SENT:
+            ret=packet_wait(session,SSH2_MSG_KEXDH_REPLY,1);
+            if(ret != SSH_OK)
+                return ret;
+            pubkey=buffer_get_ssh_string(session->in_buffer);
+            if(!pubkey){
+                ssh_set_error(session,SSH_FATAL,"No public key in packet");
+                return SSH_ERROR;
+            }
+            dh_import_pubkey(session,pubkey);
+            f=buffer_get_ssh_string(session->in_buffer);
+            if(!f){
+                ssh_set_error(session,SSH_FATAL,"No F number in packet");
+                return SSH_ERROR;
+            }
+            dh_import_f(session,f);
+            free(f);
+            if(!(signature=buffer_get_ssh_string(session->in_buffer))){
+                ssh_set_error(session,SSH_FATAL,"No signature in packet");
+                return SSH_ERROR;
+            }
+            session->dh_server_signature=signature;
+            dh_build_k(session);
+            // send the MSG_NEWKEYS
+            packet_clear_out(session);
+            buffer_add_u8(session->out_buffer,SSH2_MSG_NEWKEYS);
+            packet_send(session);
+            session->dh_handshake_state=DH_STATE_NEWKEYS_TO_SEND;
+        case DH_STATE_NEWKEYS_TO_SEND:
+            ret=packet_flush(session,0);
+            if(ret != SSH_OK)
+                return ret;
+            ssh_say(2,"SSH_MSG_NEWKEYS sent\n");
+            session->dh_handshake_state=DH_STATE_NEWKEYS_SENT;
+        case DH_STATE_NEWKEYS_SENT:
+            ret=packet_wait(session,SSH2_MSG_NEWKEYS,1);
+            if(ret != SSH_OK)
+                return ret;
+            ssh_say(2,"Got SSH_MSG_NEWKEYS\n");
+            make_sessionid(session);
+            /* set the cryptographic functions for the next crypto */
+            /* (it is needed for generate_session_keys for key lenghts) */
+            if(crypt_set_algorithms(session))
+                return SSH_ERROR;
+            generate_session_keys(session);
+            /* verify the host's signature. XXX do it sooner */
+            signature=session->dh_server_signature;
+            session->dh_server_signature=NULL;
+            if(signature_verify(session,signature)){
+                free(signature);
+                return SSH_ERROR;
+            }
+            free(signature);	/* forget it for now ... */
+            /* once we got SSH2_MSG_NEWKEYS we can switch next_crypto and current_crypto */
+            if(session->current_crypto)
+                crypto_free(session->current_crypto);
+                /* XXX later, include a function to change keys */
+            session->current_crypto=session->next_crypto;
+            session->next_crypto=crypto_new();
+            session->dh_handshake_state=DH_STATE_FINISHED;
+            return SSH_OK;
+        default:
+            ssh_set_error(session,SSH_FATAL,"Invalid state in dh_handshake():%d",session->dh_handshake_state);
+            return SSH_ERROR;
     }
-    free(signature);	/* forget it for now ... */
-    /* once we got SSH2_MSG_NEWKEYS we can switch next_crypto and current_crypto */
-    if(session->current_crypto)
-        crypto_free(session->current_crypto);
-    /* XXX later, include a function to change keys */
-    session->current_crypto=session->next_crypto;
-    session->next_crypto=crypto_new();
-    return 0;
+    /* not reached */
+    return SSH_ERROR;
 }
 
 int ssh_service_request(SSH_SESSION *session,char *service){
@@ -185,13 +223,13 @@ int ssh_connect(SSH_SESSION *session){
   SSH_OPTIONS *options=session->options;
   if(!session->options){
       ssh_set_error(session,SSH_FATAL,"Must set options before connect");
-      return -1;
+      return SSH_ERROR;
   }
   session->client=1;
   ssh_crypto_init();
   if(options->fd==-1 && !options->host){
       ssh_set_error(session,SSH_FATAL,"Hostname required");
-      return -1;
+      return SSH_ERROR;
   } 
   if(options->fd != -1)
       fd=options->fd;
diff --git a/libssh/packet.c b/libssh/packet.c
index 330b6908..b653bb63 100644
--- a/libssh/packet.c
+++ b/libssh/packet.c
@@ -335,19 +335,19 @@ int packet_translate(SSH_SESSION *session){
 
 static int atomic_write(int fd, void *buffer, int len){
     int written;
-    do {
+    while(len >0) {
         written=write(fd,buffer,len);
         if(written==0 || written==-1)
-            return -1;
+            return SSH_ERROR;
         len-=written;
         buffer+=written;
-    } while (len > 0);
+    }
     return SSH_OK;
 }
 
 /* when doing a nonblocking write, you should issue the packet_write only once, then 
  * do packet_nonblocking_flush() until you get a SSH_OK or a SSH_ERROR */
-int packet_nonblocking_flush(SSH_SESSION *session){
+static int packet_nonblocking_flush(SSH_SESSION *session){
     int except, can_write;
     int w;
     ssh_fd_poll(session,&can_write,&except); /* internally sets data_to_write */
@@ -375,7 +375,11 @@ int packet_nonblocking_flush(SSH_SESSION *session){
 }
 
 /* blocking packet flush */
-int packet_blocking_flush(SSH_SESSION *session){
+static int packet_blocking_flush(SSH_SESSION *session){
+    if(session->data_except)
+        return SSH_ERROR;
+    if(buffer_get_rest(session->out_socket_buffer)==0)
+        return SSH_OK;
     if(atomic_write(session->fd,buffer_get_rest(session->out_socket_buffer),
        buffer_get_rest_len(session->out_socket_buffer))){
         session->data_to_write=0;
@@ -392,6 +396,16 @@ int packet_blocking_flush(SSH_SESSION *session){
     return SSH_OK; // no data pending
 }
 
+/* Write the the bufferized output. If the session is blocking, or enforce_blocking 
+ * is set, the call may block. Otherwise, it won't block.
+ * return SSH°OK if everything has been sent, SSH_AGAIN if there are still things 
+ * to send on buffer, SSH_ERROR if there is an error. */
+int packet_flush(SSH_SESSION *session, int enforce_blocking){
+    if(enforce_blocking || session->blocking)
+        return packet_blocking_flush(session);
+    return packet_nonblocking_flush(session);
+}
+
 /* this function places the outgoing packet buffer into an outgoing socket buffer */
 static int socket_write(SSH_SESSION *session){
     if(!session->out_socket_buffer){
@@ -596,15 +610,17 @@ static int packet_wait1(SSH_SESSION *session,int type,int blocking){
 }
 #endif /* HAVE_SSH1 */
 static int packet_wait2(SSH_SESSION *session,int type,int blocking){
+    int ret;
     while(1){
-        if(packet_read2(session))
-            return -1;
+        ret=packet_read2(session);
+        if(ret != SSH_OK)
+            return ret;
         if(packet_translate(session))
-            return -1;
+            return SSH_ERROR;
         switch(session->in_packet.type){
            case SSH2_MSG_DISCONNECT:
                packet_parse(session);
-                return -1;
+                return SSH_ERROR;
            case SSH2_MSG_CHANNEL_WINDOW_ADJUST:
            case SSH2_MSG_CHANNEL_DATA:
            case SSH2_MSG_CHANNEL_EXTENDED_DATA:
@@ -612,20 +628,20 @@ static int packet_wait2(SSH_SESSION *session,int type,int blocking){
            case SSH2_MSG_CHANNEL_EOF:
            case SSH2_MSG_CHANNEL_CLOSE:
                packet_parse(session);
-                break;;
+               break;
            case SSH2_MSG_IGNORE:
                break;
            default:
                if(type && (type != session->in_packet.type)){
                    ssh_set_error(session,SSH_FATAL,"waitpacket(): Received a %d type packet, was waiting for a %d\n",session->in_packet.type,type);
-                   return -1;
+                   return SSH_ERROR;
                }
-               return 0;
+               return SSH_OK;
            }
         if(blocking==0)
-            return 0;
+            return SSH_OK; //shouldn't it return SSH_AGAIN here ?
     }
-    return 0;
+    return SSH_OK;
 }
 int packet_wait(SSH_SESSION *session, int type, int block){
 #ifdef HAVE_SSH1