router.rb in vines-0.3.0

- old
+ new

@@ -4,11 +4,10 @@
   # The router tracks all stream connections to the server for all clients,
   # servers, and components. It sends stanzas to the correct stream based on
   # the 'to' attribute. Router is a singleton, shared by all streams, that must
   # be accessed with +Router.instance+, not +Router.new+.
   class Router
-    ROUTABLE_STANZAS = %w[message iq presence].freeze
 
     STREAM_TYPES = [:client, :server, :component].freeze
     STREAM_TYPES.each do |name|
       define_method "#{name}s" do
         @streams[name]
@@ -27,34 +26,36 @@
     end
 
     # Returns streams for all connected resources for this JID. A
     # resource is considered connected after it has completed authentication
     # and resource binding.
-    def connected_resources(jid)
-      jid = JID.new(jid)
+    def connected_resources(jid, from)
+      jid, from = JID.new(jid), JID.new(from)
       clients.select do |stream|
-        stream.connected? && jid == (jid.bare? ? stream.user.jid.bare : stream.user.jid)
+        stream.connected? &&
+          jid == (jid.bare? ? stream.user.jid.bare : stream.user.jid) &&
+          @config.allowed?(jid, from)
       end
     end
 
     # Returns streams for all available resources for this JID. A
     # resource is marked available after it sends initial presence.
     # This method accepts a single JID or a list of JIDs.
-    def available_resources(*jid)
-      ids = jid.flatten.map {|jid| JID.new(jid).bare }
+    def available_resources(*jids, from)
+      jids = filter_allowed(jids, from)
       clients.select do |stream|
-        stream.available? && ids.include?(stream.user.jid.bare)
+        stream.available? && jids.include?(stream.user.jid.bare)
       end
     end
 
     # Returns streams for all interested resources for this JID. A
     # resource is marked interested after it requests the roster.
     # This method accepts a single JID or a list of JIDs.
-    def interested_resources(*jid)
-      ids = jid.flatten.map {|jid| JID.new(jid).bare }
+    def interested_resources(*jids, from)
+      jids = filter_allowed(jids, from)
       clients.select do |stream|
-        stream.interested? && ids.include?(stream.user.jid.bare)
+        stream.interested? && jids.include?(stream.user.jid.bare)
       end
     end
 
     # Add the connection to the routing table. The connection must return
     # :client, :server, or :component from its +stream_type+ method so the
@@ -73,53 +74,84 @@
 
     # Send the stanza to the appropriate remote server-to-server stream
     # or an external component stream.
     def route(stanza)
       to, from = %w[to from].map {|attr| JID.new(stanza[attr]) }
-      if stream = connection_to(to.domain)
+      return unless @config.allowed?(to, from)
+      key = [to.domain, from.domain]
+
+      if stream = connection_to(to, from)
         stream.write(stanza)
-      elsif @pending.key?(to.domain)
-        @pending[to.domain] << stanza
+      elsif @pending.key?(key)
+        @pending[key] << stanza
       elsif @config.s2s?(to.domain)
-        @pending[to.domain] << stanza
+        @pending[key] << stanza
         Vines::Stream::Server.start(@config, to.domain, from.domain) do |stream|
-          if stream
-            @pending[to.domain].each {|s| stream.write(s) }
-          else
-            @pending[to.domain].each do |s|
-              xml = StanzaErrors::RemoteServerNotFound.new(s, 'cancel').to_xml
-              connected_resources(s['from']).each {|c| c.write(xml) }
-            end
-          end
-          @pending.delete(to.domain)
+          stream ? send_pending(key, stream) : return_pending(key)
+          @pending.delete(key)
         end
       else
         raise StanzaErrors::RemoteServerNotFound.new(stanza, 'cancel')
       end
     end
 
-    # Returns true if this stanza should be processed locally. Returns false
-    # if it's destined for a remote domain or external component.
-    def local?(stanza)
-      return true unless ROUTABLE_STANZAS.include?(stanza.name)
-      to = (stanza['to'] || '').strip
-      to.empty? || local_jid?(to)
-    end
-
-    def local_jid?(jid)
-      @config.vhost?(JID.new(jid).domain)
-    end
-
     # Returns the total number of streams connected to the server.
     def size
       @streams.values.inject(0) {|sum, arr| sum + arr.size }
     end
 
     private
 
-    def connection_to(domain)
-      (components + servers).find do |stream|
-        stream.ready? && stream.remote_domain == domain
+    # Write all pending stanzas for this domain to the stream. Called after a
+    # s2s stream has successfully connected and we need to dequeue all stanzas
+    # we received while waiting for the connection to finish.
+    def send_pending(key, stream)
+      @pending[key].each do |stanza|
+        stream.write(stanza)
+      end
+    end
+
+    # Return all pending stanzas to their senders as remote-server-not-found
+    # errors. Called after a s2s stream has failed to connect.
+    def return_pending(key)
+      @pending[key].each do |stanza|
+        to, from = JID.new(stanza['to']), JID.new(stanza['from'])
+        xml = StanzaErrors::RemoteServerNotFound.new(stanza, 'cancel').to_xml
+        if @config.component?(from)
+          connection_to(from, to).write(xml) rescue nil
+        else
+          connected_resources(from, to).each {|c| c.write(xml) }
+        end
+      end
+    end
+
+    # Return the bare JID's from the list that are allowed to talk to
+    # the +from+ JID. Store them in a Hash for fast +include?+ checks.
+    def filter_allowed(jids, from)
+      from = JID.new(from)
+      {}.tap do |ids|
+        jids.flatten.each do |jid|
+          jid = JID.new(jid).bare
+          ids[jid] = nil if @config.allowed?(jid, from)
+        end
+      end
+    end
+
+    def connection_to(to, from)
+      component_stream(to) || server_stream(to, from)
+    end
+
+    def component_stream(to)
+      components.find do |stream|
+        stream.ready? && stream.remote_domain == to.domain
+      end
+    end
+
+    def server_stream(to, from)
+      servers.find do |stream|
+        stream.ready? &&
+          stream.remote_domain == to.domain &&
+            stream.domain == from.domain
       end
     end
 
     def stream_type(connection)
       connection.stream_type.tap do |type|