storage/mod_ejabberd: Reorganized some code.

[prosody.git] / net / server_select.lua

diff --git a/net/server_select.lua b/net/server_select.lua
index e3a389a878864aa959d40fd3338170aa5be1c5c4..51ae4e66a1e1ccbbf990ef6356c2cb9b58303dda 100644 (file)
--- a/net/server_select.lua
+++ b/net/server_select.lua
@@ -167,7 +167,7 @@ wrapserver = function( listeners, socket, ip, serverport, pattern, sslctx, maxco
 
        local connections = 0
 
-       local dispatch, disconnect = listeners.onincoming, listeners.ondisconnect
+       local dispatch, disconnect = listeners.onconnect or listeners.onincoming, listeners.ondisconnect
 
        local accept = socket.accept
 
@@ -343,9 +343,9 @@ wrapconnection = function( server, listeners, socket, ip, serverport, clientport
                        _closelist[ handler ] = nil
                        handler = nil
                end
-       if server then
-               server.remove( )
-       end
+               if server then
+                       server.remove( )
+               end
                out_put "server.lua: closed client handler and removed socket from list"
                return true
        end
@@ -479,11 +479,11 @@ wrapconnection = function( server, listeners, socket, ip, serverport, clientport
                        bufferqueuelen = 0
                        bufferlen = 0
                        _sendlistlen = removesocket( _sendlist, socket, _sendlistlen ) -- delete socket from writelist
-                       _ = needtls and handler:starttls(nil, true)
                        _writetimes[ handler ] = nil
                        if drain then
                                drain(handler)
                        end
+                       _ = needtls and handler:starttls(nil)
                        _ = toclose and handler:close( )
                        return true
                elseif byte and ( err == "timeout" or err == "wantwrite" ) then -- want write
@@ -564,13 +564,13 @@ wrapconnection = function( server, listeners, socket, ip, serverport, clientport
                        end
                else
                        local sslctx;
-                       handler.starttls = function( self, _sslctx, now )
+                       handler.starttls = function( self, _sslctx)
                                if _sslctx then
                                        sslctx = _sslctx;
                                        handler:set_sslctx(sslctx);
                                end
-                               if not now then
-                                       out_put "server.lua: we need to do tls, but delaying until later"
+                               if bufferqueuelen > 0 then
+                                       out_put "server.lua: we need to do tls, but delaying until send buffer empty"
                                        needtls = true
                                        return
                                end
@@ -623,16 +623,6 @@ wrapconnection = function( server, listeners, socket, ip, serverport, clientport
 
        _socketlist[ socket ] = handler
        _readlistlen = addsocket(_readlist, socket, _readlistlen)
-       if listeners.onconnect then
-               _sendlistlen = addsocket(_sendlist, socket, _sendlistlen)
-               handler.sendbuffer = function ()
-                       listeners.onconnect(handler);
-                       handler.sendbuffer = _sendbuffer;
-                       if bufferqueuelen > 0 then
-                               return _sendbuffer();
-                       end
-               end
-       end
        return handler, socket
 end
 
@@ -854,6 +844,18 @@ local wrapclient = function( socket, ip, serverport, listeners, pattern, sslctx
        local handler = wrapconnection( nil, listeners, socket, ip, serverport, "clientport", pattern, sslctx )
        _socketlist[ socket ] = handler
        _sendlistlen = addsocket(_sendlist, socket, _sendlistlen)
+       if listeners.onconnect then
+               -- When socket is writeable, call onconnect
+               local _sendbuffer = handler.sendbuffer;
+               handler.sendbuffer = function ()
+                       listeners.onconnect(handler);
+                       handler.sendbuffer = _sendbuffer;
+                       -- If there was data with the incoming packet, handle it now.
+                       if #handler:bufferqueue() > 0 then
+                               return _sendbuffer();
+                       end
+               end
+       end
        return handler, socket
 end