Refactor XmppStream, TLS and connection method logic

fixes #534

13 files changed, 446 insertions, 371 deletions

diff --git a/xmpp-vala/CMakeLists.txt b/xmpp-vala/CMakeLists.txt
index 94c7f75b..5778c2de 100644
--- a/xmpp-vala/CMakeLists.txt
+++ b/xmpp-vala/CMakeLists.txt
@@ -20,13 +20,20 @@ set(ENGINE_EXTRA_OPTIONS ${MAIN_EXTRA_OPTIONS} --vapidir=${CMAKE_CURRENT_SOURCE_
 
 vala_precompile(ENGINE_VALA_C
 SOURCES
+    "src/core/direct_tls_xmpp_stream.vala"
+    "src/core/io_xmpp_stream.vala"
+    "src/core/module_flag.vala"
+    "src/core/starttls_xmpp_stream.vala"
+    "src/core/stream_connect.vala"
+    "src/core/tls_xmpp_stream.vala"
+    "src/core/xmpp_stream.vala"
+
     "src/core/namespace_state.vala"
     "src/core/stanza_attribute.vala"
     "src/core/stanza_node.vala"
     "src/core/stanza_reader.vala"
     "src/core/stanza_writer.vala"
     "src/core/xmpp_log.vala"
-    "src/core/xmpp_stream.vala"
 
     "src/module/bind.vala"
     "src/module/bookmarks_provider.vala"
@@ -48,7 +55,6 @@ SOURCES
     "src/module/stanza.vala"
     "src/module/stanza_error.vala"
     "src/module/stream_error.vala"
-    "src/module/tls.vala"
     "src/module/util.vala"
 
     "src/module/xep/0048_bookmarks.vala"
@@ -91,7 +97,6 @@ SOURCES
     "src/module/xep/0334_message_processing_hints.vala"
     "src/module/xep/0359_unique_stable_stanza_ids.vala"
     "src/module/xep/0363_http_file_upload.vala"
-    "src/module/xep/0368_srv_records_tls.vala"
     "src/module/xep/0380_explicit_encryption.vala"
     "src/module/xep/0391_jingle_encrypted_transports.vala"
     "src/module/xep/0410_muc_self_ping.vala"
diff --git a/xmpp-vala/src/core/direct_tls_xmpp_stream.vala b/xmpp-vala/src/core/direct_tls_xmpp_stream.vala
new file mode 100644
index 00000000..1d2f7339
--- /dev/null
+++ b/xmpp-vala/src/core/direct_tls_xmpp_stream.vala
@@ -0,0 +1,31 @@
+public class Xmpp.DirectTlsXmppStream : TlsXmppStream {
+
+    string host;
+    uint16 port;
+
+    public DirectTlsXmppStream(Jid remote, string host, uint16 port) {
+        this.remote_name = remote;
+        this.host = host;
+        this.port = port;
+    }
+
+    public override async void connect() throws IOStreamError {
+        SocketClient client = new SocketClient();
+        try {
+            debug("Connecting to %s %i (tls)", host, port);
+            IOStream? io_stream = yield client.connect_to_host_async(host, port);
+            TlsConnection tls_connection = TlsClientConnection.new(io_stream, new NetworkAddress(remote_name.to_string(), port));
+            #if ALPN_SUPPORT
+            tls_connection.set_advertised_protocols(new string[]{"xmpp-client"});
+            #endif
+            tls_connection.accept_certificate.connect(on_invalid_certificate);
+            reset_stream(tls_connection);
+
+            yield setup();
+
+            attach_negotation_modules();
+        } catch (Error e) {
+            throw new IOStreamError.CONNECT("Failed connecting to %s:%i (tls): %s", host, port, e.message);
+        }
+    }
+}
\ No newline at end of file
diff --git a/xmpp-vala/src/core/io_xmpp_stream.vala b/xmpp-vala/src/core/io_xmpp_stream.vala
new file mode 100644
index 00000000..56efd7cd
--- /dev/null
+++ b/xmpp-vala/src/core/io_xmpp_stream.vala
@@ -0,0 +1,101 @@
+using Gee;
+
+public interface Xmpp.WriteNodeFunc : Object {
+    public abstract async void write_stanza(XmppStream stream, StanzaNode node) throws IOStreamError;
+}
+
+public abstract class Xmpp.IoXmppStream : XmppStream {
+    private IOStream? stream;
+    internal StanzaReader? reader;
+    internal StanzaWriter? writer;
+
+    internal WriteNodeFunc? write_obj = null;
+
+    public override async void disconnect() throws IOStreamError, XmlError, IOError {
+        disconnected = true;
+        if (writer == null || reader == null || stream == null) {
+            throw new IOStreamError.DISCONNECT("trying to disconnect, but no stream open");
+        }
+        log.str("OUT", "</stream:stream>", this);
+        yield writer.write("</stream:stream>");
+        reader.cancel();
+        yield stream.close_async();
+    }
+
+    public void reset_stream(IOStream stream) {
+        this.stream = stream;
+        reader = new StanzaReader.for_stream(stream.input_stream);
+        writer = new StanzaWriter.for_stream(stream.output_stream);
+
+        writer.cancel.connect(reader.cancel);
+        require_setup();
+    }
+
+    public override async StanzaNode read() throws IOStreamError {
+        StanzaReader? reader = this.reader;
+        if (reader == null) throw new IOStreamError.READ("trying to read, but no stream open");
+        try {
+            StanzaNode node = yield ((!)reader).read_node();
+            log.node("IN", node, this);
+            return node;
+        } catch (XmlError e) {
+            throw new IOStreamError.READ(e.message);
+        }
+    }
+
+    [Version (deprecated = true, deprecated_since = "0.1", replacement = "write_async")]
+    public override void write(StanzaNode node) {
+        write_async.begin(node, (obj, res) => {
+            try {
+                write_async.end(res);
+            } catch (Error e) { }
+        });
+    }
+
+    public override async void write_async(StanzaNode node) throws IOStreamError {
+        if (write_obj != null) {
+            yield write_obj.write_stanza(this, node);
+        } else {
+            StanzaWriter? writer = this.writer;
+            if (writer == null) throw new IOStreamError.WRITE("trying to write, but no stream open");
+            try {
+                log.node("OUT", node, this);
+                yield ((!)writer).write_node(node);
+            } catch (XmlError e) {
+                throw new IOStreamError.WRITE(e.message);
+            }
+        }
+    }
+
+    internal IOStream? get_stream() {
+        return stream;
+    }
+
+    public override async void setup() throws IOStreamError {
+        StanzaNode outs = new StanzaNode.build("stream", "http://etherx.jabber.org/streams")
+                .put_attribute("to", remote_name.to_string())
+                .put_attribute("version", "1.0")
+                .put_attribute("xmlns", "jabber:client")
+                .put_attribute("stream", "http://etherx.jabber.org/streams", XMLNS_URI);
+        outs.has_nodes = true;
+        log.node("OUT ROOT", outs, this);
+        write(outs);
+        received_root_node(this, yield read_root());
+
+        setup_needed = false;
+    }
+
+    private async StanzaNode read_root() throws IOStreamError {
+        StanzaReader? reader = this.reader;
+        if (reader == null) throw new IOStreamError.READ("trying to read, but no stream open");
+        try {
+            StanzaNode node = yield ((!)reader).read_root_node();
+            log.node("IN ROOT", node, this);
+            return node;
+        } catch (XmlError.TLS e) {
+            throw new IOStreamError.TLS(e.message);
+        } catch (Error e) {
+            throw new IOStreamError.READ(e.message);
+        }
+    }
+}
\ No newline at end of file
diff --git a/xmpp-vala/src/core/module_flag.vala b/xmpp-vala/src/core/module_flag.vala
new file mode 100644
index 00000000..95547852
--- /dev/null
+++ b/xmpp-vala/src/core/module_flag.vala
@@ -0,0 +1,61 @@
+namespace Xmpp {
+
+    public class FlagIdentity<T> : Object {
+        public string ns { get; private set; }
+        public string id { get; private set; }
+
+        public FlagIdentity(string ns, string id) {
+            this.ns = ns;
+            this.id = id;
+        }
+
+        public T? cast(XmppStreamFlag flag) {
+            return flag.get_type().is_a(typeof(T)) ? (T?) flag : null;
+        }
+
+        public bool matches(XmppStreamFlag module) {
+            return module.get_ns() == ns && module.get_id() == id;
+        }
+    }
+
+    public abstract class XmppStreamFlag : Object {
+        public abstract string get_ns();
+
+        public abstract string get_id();
+    }
+
+    public class ModuleIdentity<T> : Object {
+        public string ns { get; private set; }
+        public string id { get; private set; }
+
+        public ModuleIdentity(string ns, string id) {
+            this.ns = ns;
+            this.id = id;
+        }
+
+        public T? cast(XmppStreamModule module) {
+            return module.get_type().is_a(typeof(T)) ? (T?) module : null;
+        }
+
+        public bool matches(XmppStreamModule module) {
+            return module.get_ns() == ns && module.get_id() == id;
+        }
+    }
+
+    public abstract class XmppStreamModule : Object {
+        public abstract void attach(XmppStream stream);
+
+        public abstract void detach(XmppStream stream);
+
+        public abstract string get_ns();
+
+        public abstract string get_id();
+    }
+
+    public abstract class XmppStreamNegotiationModule : XmppStreamModule {
+        public abstract bool mandatory_outstanding(XmppStream stream);
+
+        public abstract bool negotiation_active(XmppStream stream);
+    }
+
+}
\ No newline at end of file
diff --git a/xmpp-vala/src/core/starttls_xmpp_stream.vala b/xmpp-vala/src/core/starttls_xmpp_stream.vala
new file mode 100644
index 00000000..3df0dffb
--- /dev/null
+++ b/xmpp-vala/src/core/starttls_xmpp_stream.vala
@@ -0,0 +1,54 @@
+public class Xmpp.StartTlsXmppStream : TlsXmppStream {
+
+    private const string TLS_NS_URI = "urn:ietf:params:xml:ns:xmpp-tls";
+
+    string host;
+    uint16 port;
+
+    public StartTlsXmppStream(Jid remote, string host, uint16 port) {
+        this.remote_name = remote;
+        this.host = host;
+        this.port = port;
+    }
+
+    public override async void connect() throws IOStreamError {
+        try {
+            SocketClient client = new SocketClient();
+            debug("Connecting to %s %i (starttls)", host, port);
+            IOStream stream = yield client.connect_to_host_async(host, port);
+            reset_stream(stream);
+
+            yield setup();
+
+            StanzaNode node = yield read();
+            var starttls_node = node.get_subnode("starttls", TLS_NS_URI);
+            if (starttls_node == null) {
+                warning("%s does not offer starttls", remote_name.to_string());
+            }
+
+            write(new StanzaNode.build("starttls", TLS_NS_URI).add_self_xmlns());
+
+            node = yield read();
+
+            if (node.ns_uri != TLS_NS_URI || node.name != "proceed") {
+                warning("Server did not 'proceed' starttls request");
+            }
+
+            try {
+                var identity = new NetworkService("xmpp-client", "tcp", remote_name.to_string());
+                var conn = TlsClientConnection.new(get_stream(), identity);
+                reset_stream(conn);
+
+                conn.accept_certificate.connect(on_invalid_certificate);
+            } catch (Error e) {
+                stderr.printf("Failed to start TLS: %s\n", e.message);
+            }
+
+            yield setup();
+
+            attach_negotation_modules();
+        } catch (Error e) {
+            throw new IOStreamError.CONNECT("Failed connecting to %s:%i (starttls): %s", host, port, e.message);
+        }
+    }
+}
\ No newline at end of file
diff --git a/xmpp-vala/src/core/stream_connect.vala b/xmpp-vala/src/core/stream_connect.vala
new file mode 100644
index 00000000..a7615e9f
--- /dev/null
+++ b/xmpp-vala/src/core/stream_connect.vala
@@ -0,0 +1,89 @@
+namespace Xmpp {
+
+    private class SrvTargetInfo {
+        public string host { get; set; }
+        public uint16 port { get; set; }
+        public string service { get; set; }
+        public uint16 priority { get; set; }
+    }
+
+    public class XmppStreamResult {
+        public XmppStream? stream { get; set; }
+        public TlsCertificateFlags? tls_errors { get; set; }
+        public IOStreamError? io_error { get; set; }
+    }
+
+    public async XmppStreamResult establish_stream(Jid bare_jid, Gee.List<XmppStreamModule> modules, string? log_options) {
+        Jid remote = bare_jid.domain_jid;
+
+        //Lookup xmpp-client and xmpps-client SRV records
+        GLib.List<SrvTargetInfo>? targets = new GLib.List<SrvTargetInfo>();
+        GLibFixes.Resolver resolver = GLibFixes.Resolver.get_default();
+        try {
+            GLib.List<SrvTarget> xmpp_services = yield resolver.lookup_service_async("xmpp-client", "tcp", remote.to_string(), null);
+            foreach (SrvTarget service in xmpp_services) {
+                targets.append(new SrvTargetInfo() { host=service.get_hostname(), port=service.get_port(), service="xmpp-client", priority=service.get_priority()});
+            }
+        } catch (Error e) {
+            debug("Got no xmpp-client DNS records for %s: %s", remote.to_string(), e.message);
+        }
+        try {
+            GLib.List<SrvTarget> xmpp_services = yield resolver.lookup_service_async("xmpps-client", "tcp", remote.to_string(), null);
+            foreach (SrvTarget service in xmpp_services) {
+                targets.append(new SrvTargetInfo() { host=service.get_hostname(), port=service.get_port(), service="xmpps-client", priority=service.get_priority()});
+            }
+        } catch (Error e) {
+            debug("Got no xmpps-client DNS records for %s: %s", remote.to_string(), e.message);
+        }
+
+        targets.sort((a, b) => {
+            return a.priority - b.priority;
+        });
+
+        // Add fallback connection
+        bool should_add_fallback = true;
+        foreach (SrvTargetInfo target in targets) {
+            if (target.service == "xmpp-client" && target.port == 5222 && target.host == remote.to_string()) {
+                should_add_fallback = false;
+            }
+        }
+        if (should_add_fallback) {
+            targets.append(new SrvTargetInfo() { host=remote.to_string(), port=5222, service="xmpp-client", priority=uint16.MAX});
+        }
+
+        // Try all connection options from lowest to highest priority
+        TlsXmppStream? stream = null;
+        TlsCertificateFlags? tls_errors = null;
+        IOStreamError? io_error = null;
+        foreach (SrvTargetInfo target in targets) {
+            try {
+                if (target.service == "xmpp-client") {
+                    stream = new StartTlsXmppStream(remote, target.host, target.port);
+                } else {
+                    stream = new DirectTlsXmppStream(remote, target.host, target.port);
+                }
+                stream.log = new XmppLog(bare_jid.to_string(), log_options);
+
+                foreach (XmppStreamModule module in modules) {
+                    stream.add_module(module);
+                }
+
+                yield stream.connect();
+
+                return new XmppStreamResult() { stream=stream };
+            } catch (IOStreamError e) {
+                warning("Could not establish XMPP session with %s:%i: %s", target.host, target.port, e.message);
+
+                if (stream != null) {
+                    if (stream.errors != null) {
+                        tls_errors = stream.errors;
+                    }
+                    io_error = e;
+                    stream.detach_modules();
+                }
+            }
+        }
+
+        return new XmppStreamResult() { io_error=io_error, tls_errors=tls_errors };
+    }
+}
\ No newline at end of file
diff --git a/xmpp-vala/src/core/tls_xmpp_stream.vala b/xmpp-vala/src/core/tls_xmpp_stream.vala
new file mode 100644
index 00000000..956a9a22
--- /dev/null
+++ b/xmpp-vala/src/core/tls_xmpp_stream.vala
@@ -0,0 +1,19 @@
+public abstract class Xmpp.TlsXmppStream : IoXmppStream {
+
+    public TlsCertificateFlags? errors;
+
+    protected bool on_invalid_certificate(TlsCertificate peer_cert, TlsCertificateFlags errors) {
+        this.errors = errors;
+
+        string error_str = "";
+        foreach (var f in new TlsCertificateFlags[]{TlsCertificateFlags.UNKNOWN_CA, TlsCertificateFlags.BAD_IDENTITY,
+            TlsCertificateFlags.NOT_ACTIVATED, TlsCertificateFlags.EXPIRED, TlsCertificateFlags.REVOKED,
+            TlsCertificateFlags.INSECURE, TlsCertificateFlags.GENERIC_ERROR, TlsCertificateFlags.VALIDATE_ALL}) {
+            if (f in errors) {
+                error_str += @"$(f), ";
+            }
+        }
+        warning(@"Tls Certificate Errors: $(error_str)");
+        return false;
+    }
+}
\ No newline at end of file
diff --git a/xmpp-vala/src/core/xmpp_stream.vala b/xmpp-vala/src/core/xmpp_stream.vala
index ad4dae97..99dbffe6 100644
--- a/xmpp-vala/src/core/xmpp_stream.vala
+++ b/xmpp-vala/src/core/xmpp_stream.vala
@@ -1,8 +1,6 @@
 using Gee;
 
-namespace Xmpp {
-
-public errordomain IOStreamError {
+public errordomain Xmpp.IOStreamError {
     READ,
     WRITE,
     CONNECT,
@@ -10,26 +8,7 @@ public errordomain IOStreamError {
     TLS
 }
 
-public class XmppStream {
-    public const string NS_URI = "http://etherx.jabber.org/streams";
-
-    public Jid remote_name;
-    public XmppLog log = new XmppLog();
-    public StanzaNode? features { get; private set; default = new StanzaNode.build("features", NS_URI); }
-
-    private IOStream? stream;
-    internal StanzaReader? reader;
-    internal StanzaWriter? writer;
-
-    public Gee.List<XmppStreamFlag> flags { get; private set; default=new ArrayList<XmppStreamFlag>(); }
-    public Gee.List<XmppStreamModule> modules { get; private set; default=new ArrayList<XmppStreamModule>(); }
-    private Gee.List<ConnectionProvider> connection_providers = new ArrayList<ConnectionProvider>();
-
-    internal WriteNodeFunc? write_obj = null;
-    public bool negotiation_complete { get; set; default=false; }
-    private bool setup_needed = false;
-    private bool non_negotiation_modules_attached = false;
-    private bool disconnected = false;
+public abstract class Xmpp.XmppStream {
 
     public signal void received_node(XmppStream stream, StanzaNode node);
     public signal void received_root_node(XmppStream stream, StanzaNode node);
@@ -41,113 +20,39 @@ public class XmppStream {
     public signal void stream_negotiated(XmppStream stream);
     public signal void attached_modules(XmppStream stream);
 
-    public XmppStream() {
-        register_connection_provider(new StartTlsConnectionProvider());
-    }
-
-    public async void connect(string? remote_name = null) throws IOStreamError {
-        try {
-            if (remote_name != null) this.remote_name = new Jid(remote_name);
-        } catch (InvalidJidError e) {
-            throw new IOStreamError.CONNECT(@"Invalid remote name \"$remote_name\": $(e.message)");
-        }
-        attach_negotation_modules();
-        try {
-            int min_priority = -1;
-            ConnectionProvider? best_provider = null;
-            foreach (ConnectionProvider connection_provider in connection_providers) {
-                int? priority = yield connection_provider.get_priority(this.remote_name);
-                if (priority != null && (priority < min_priority || min_priority == -1)) {
-                    min_priority = priority;
-                    best_provider = connection_provider;
-                }
-            }
-            IOStream? stream = null;
-            if (best_provider != null) {
-                stream = yield best_provider.connect(this);
-            }
-            if (stream == null) {
-                debug("Connecting to %s, xmpp-client, tcp (fallback)", this.remote_name.to_string());
-                stream = yield (new SocketClient()).connect_to_host_async(this.remote_name.to_string(), 5222);
-            }
-            if (stream == null) {
-                throw new IOStreamError.CONNECT("client.connect() returned null");
-            }
-            reset_stream((!)stream);
-        } catch (Error e) {
-            debug("[%p] Could not connect to server: %s", this, e.message);
-            throw new IOStreamError.CONNECT(e.message);
-        }
-        debug("Connected to %s", remote_name);
-        yield loop();
-    }
+    public const string NS_URI = "http://etherx.jabber.org/streams";
 
-    public async void disconnect() throws IOStreamError, XmlError, IOError {
-        disconnected = true;
-        if (writer == null || reader == null || stream == null) {
-            throw new IOStreamError.DISCONNECT("trying to disconnect, but no stream open");
-        }
-        log.str("OUT", "</stream:stream>", this);
-        yield writer.write("</stream:stream>");
-        reader.cancel();
-        yield stream.close_async();
-    }
+    public Gee.List<XmppStreamFlag> flags { get; private set; default=new ArrayList<XmppStreamFlag>(); }
+    public Gee.List<XmppStreamModule> modules { get; private set; default=new ArrayList<XmppStreamModule>(); }
 
-    public void reset_stream(IOStream stream) {
-        this.stream = stream;
-        reader = new StanzaReader.for_stream(stream.input_stream);
-        writer = new StanzaWriter.for_stream(stream.output_stream);
+    public StanzaNode? features { get; private set; default = new StanzaNode.build("features", NS_URI); }
+    public Jid remote_name;
 
-        writer.cancel.connect(reader.cancel);
-        require_setup();
-    }
+    public XmppLog log = new XmppLog();
+    public bool negotiation_complete { get; set; default=false; }
+    protected bool non_negotiation_modules_attached = false;
+    protected bool setup_needed = false;
+    protected bool disconnected = false;
 
-    public void require_setup() {
-        setup_needed = true;
-    }
+    public abstract async void connect() throws IOStreamError;
 
-    public bool is_setup_needed() {
-        return setup_needed;
-    }
+    public abstract async void disconnect() throws IOStreamError, XmlError, IOError;
 
-    public async StanzaNode read() throws IOStreamError {
-        StanzaReader? reader = this.reader;
-        if (reader == null) throw new IOStreamError.READ("trying to read, but no stream open");
-        try {
-            StanzaNode node = yield ((!)reader).read_node();
-            log.node("IN", node, this);
-            return node;
-        } catch (XmlError e) {
-            throw new IOStreamError.READ(e.message);
-        }
-    }
+    public abstract async StanzaNode read() throws IOStreamError;
 
     [Version (deprecated = true, deprecated_since = "0.1", replacement = "write_async")]
-    public void write(StanzaNode node) {
-        write_async.begin(node, (obj, res) => {
-            try {
-                write_async.end(res);
-            } catch (Error e) { }
-        });
-    }
+            public abstract void write(StanzaNode node);
 
-    public async void write_async(StanzaNode node) throws IOStreamError {
-        if (write_obj != null) {
-            yield write_obj.write_stanza(this, node);
-        } else {
-            StanzaWriter? writer = this.writer;
-            if (writer == null) throw new IOStreamError.WRITE("trying to write, but no stream open");
-            try {
-                log.node("OUT", node, this);
-                yield ((!)writer).write_node(node);
-            } catch (XmlError e) {
-                throw new IOStreamError.WRITE(e.message);
-            }
-        }
+    public abstract async void write_async(StanzaNode node) throws IOStreamError;
+
+    public abstract async void setup() throws IOStreamError;
+
+    public void require_setup() {
+        setup_needed = true;
     }
 
-    internal IOStream? get_stream() {
-        return stream;
+    public bool is_setup_needed() {
+        return setup_needed;
     }
 
     public void add_flag(XmppStreamFlag flag) {
@@ -184,7 +89,6 @@ public class XmppStream {
 
     public void detach_modules() {
         foreach (XmppStreamModule module in modules) {
-            if (!(module is XmppStreamNegotiationModule) && !negotiation_complete) continue;
             module.detach(this);
         }
     }
@@ -197,37 +101,10 @@ public class XmppStream {
         return null;
     }
 
-    public void register_connection_provider(ConnectionProvider connection_provider) {
-        connection_providers.add(connection_provider);
-    }
-
-    public bool is_negotiation_active() {
-        foreach (XmppStreamModule module in modules) {
-            if (module is XmppStreamNegotiationModule) {
-                XmppStreamNegotiationModule negotiation_module = (XmppStreamNegotiationModule) module;
-                if (negotiation_module.negotiation_active(this)) return true;
-            }
-        }
-        return false;
-    }
-
-    private async void setup() throws IOStreamError {
-        StanzaNode outs = new StanzaNode.build("stream", "http://etherx.jabber.org/streams")
-                            .put_attribute("to", remote_name.to_string())
-                            .put_attribute("version", "1.0")
-                            .put_attribute("xmlns", "jabber:client")
-                            .put_attribute("stream", "http://etherx.jabber.org/streams", XMLNS_URI);
-        outs.has_nodes = true;
-        log.node("OUT ROOT", outs, this);
-        write(outs);
-        received_root_node(this, yield read_root());
-    }
-
-    private async void loop() throws IOStreamError {
+    public async void loop() throws IOStreamError {
         while (true) {
             if (setup_needed) {
                 yield setup();
-                setup_needed = false;
             }
 
             StanzaNode node = yield read();
@@ -237,30 +114,7 @@ public class XmppStream {
 
             if (disconnected) break;
 
-            received_node(this, node);
-
-            if (node.ns_uri == NS_URI && node.name == "features") {
-                features = node;
-                received_features_node(this);
-            } else if (node.ns_uri == NS_URI && node.name == "stream" && node.pseudo) {
-                debug("[%p] Server closed stream", this);
-                try {
-                    yield disconnect();
-                } catch (Error e) {}
-                return;
-            } else if (node.ns_uri == JABBER_URI) {
-                if (node.name == "message") {
-                    received_message_stanza(this, node);
-                } else if (node.name == "presence") {
-                    received_presence_stanza(this, node);
-                } else if (node.name == "iq") {
-                    received_iq_stanza(this, node);
-                } else {
-                    received_nonza(this, node);
-                }
-            } else {
-                received_nonza(this, node);
-            }
+            yield handle_stanza(node);
 
             if (!non_negotiation_modules_attached && negotiation_modules_done()) {
                 attach_non_negotation_modules();
@@ -273,6 +127,43 @@ public class XmppStream {
         }
     }
 
+    private async void handle_stanza(StanzaNode node) {
+        received_node(this, node);
+
+        if (node.ns_uri == NS_URI && node.name == "features") {
+            features = node;
+            received_features_node(this);
+        } else if (node.ns_uri == NS_URI && node.name == "stream" && node.pseudo) {
+            debug("[%p] Server closed stream", this);
+            try {
+                yield disconnect();
+            } catch (Error e) {}
+            return;
+        } else if (node.ns_uri == JABBER_URI) {
+            if (node.name == "message") {
+                received_message_stanza(this, node);
+            } else if (node.name == "presence") {
+                received_presence_stanza(this, node);
+            } else if (node.name == "iq") {
+                received_iq_stanza(this, node);
+            } else {
+                received_nonza(this, node);
+            }
+        } else {
+            received_nonza(this, node);
+        }
+    }
+
+    public bool is_negotiation_active() {
+        foreach (XmppStreamModule module in modules) {
+            if (module is XmppStreamNegotiationModule) {
+                XmppStreamNegotiationModule negotiation_module = (XmppStreamNegotiationModule) module;
+                if (negotiation_module.negotiation_active(this)) return true;
+            }
+        }
+        return false;
+    }
+
     private bool negotiation_modules_done() throws IOStreamError {
         if (setup_needed) return false;
         if (is_negotiation_active()) return false;
@@ -297,124 +188,11 @@ public class XmppStream {
         attached_modules(this);
     }
 
-    private void attach_negotation_modules() {
+    public void attach_negotation_modules() {
         foreach (XmppStreamModule module in modules) {
             if (module as XmppStreamNegotiationModule != null) {
                 module.attach(this);
             }
         }
     }
-
-    private async StanzaNode read_root() throws IOStreamError {
-        StanzaReader? reader = this.reader;
-        if (reader == null) throw new IOStreamError.READ("trying to read, but no stream open");
-        try {
-            StanzaNode node = yield ((!)reader).read_root_node();
-            log.node("IN ROOT", node, this);
-            return node;
-        } catch (XmlError.TLS e) {
-            throw new IOStreamError.TLS(e.message);
-        } catch (Error e) {
-            throw new IOStreamError.READ(e.message);
-        }
-    }
-}
-
-public class FlagIdentity<T> : Object {
-    public string ns { get; private set; }
-    public string id { get; private set; }
-
-    public FlagIdentity(string ns, string id) {
-        this.ns = ns;
-        this.id = id;
-    }
-
-    public T? cast(XmppStreamFlag flag) {
-        return flag.get_type().is_a(typeof(T)) ? (T?) flag : null;
-    }
-
-    public bool matches(XmppStreamFlag module) {
-        return module.get_ns() == ns && module.get_id() == id;
-    }
-}
-
-public abstract class XmppStreamFlag : Object {
-    public abstract string get_ns();
-
-    public abstract string get_id();
-}
-
-public class ModuleIdentity<T> : Object {
-    public string ns { get; private set; }
-    public string id { get; private set; }
-
-    public ModuleIdentity(string ns, string id) {
-        this.ns = ns;
-        this.id = id;
-    }
-
-    public T? cast(XmppStreamModule module) {
-        return module.get_type().is_a(typeof(T)) ? (T?) module : null;
-    }
-
-    public bool matches(XmppStreamModule module) {
-        return module.get_ns() == ns && module.get_id() == id;
-    }
-}
-
-public abstract class XmppStreamModule : Object {
-    public abstract void attach(XmppStream stream);
-
-    public abstract void detach(XmppStream stream);
-
-    public abstract string get_ns();
-
-    public abstract string get_id();
-}
-
-public abstract class XmppStreamNegotiationModule : XmppStreamModule {
-    public abstract bool mandatory_outstanding(XmppStream stream);
-
-    public abstract bool negotiation_active(XmppStream stream);
-}
-
-public abstract class ConnectionProvider {
-    public async abstract int? get_priority(Jid remote_name);
-    public async abstract IOStream? connect(XmppStream stream);
-    public abstract string get_id();
-}
-
-public class StartTlsConnectionProvider : ConnectionProvider {
-    private SrvTarget? srv_target;
-
-    public async override int? get_priority(Jid remote_name) {
-        GLib.List<SrvTarget>? xmpp_target = null;
-        try {
-            GLibFixes.Resolver resolver = GLibFixes.Resolver.get_default();
-            xmpp_target = yield resolver.lookup_service_async("xmpp-client", "tcp", remote_name.to_string(), null);
-        } catch (Error e) {
-            return null;
-        }
-        xmpp_target.sort((a, b) => { return a.get_priority() - b.get_priority(); });
-        srv_target = xmpp_target.nth(0).data;
-        return xmpp_target.nth(0).data.get_priority();
-    }
-
-    public async override IOStream? connect(XmppStream stream) {
-        try {
-            SocketClient client = new SocketClient();
-            debug("Connecting to %s %i (starttls)", srv_target.get_hostname(), srv_target.get_port());
-            return yield client.connect_to_host_async(srv_target.get_hostname(), srv_target.get_port());
-        } catch (Error e) {
-            return null;
-        }
-    }
-
-    public override string get_id() { return "start_tls"; }
-}
-
-public interface WriteNodeFunc : Object {
-    public abstract async void write_stanza(XmppStream stream, StanzaNode node) throws IOStreamError;
-}
-
-}
+}
\ No newline at end of file
diff --git a/xmpp-vala/src/module/sasl.vala b/xmpp-vala/src/module/sasl.vala
index 2e87e590..3f3eca58 100644
--- a/xmpp-vala/src/module/sasl.vala
+++ b/xmpp-vala/src/module/sasl.vala
@@ -154,7 +154,6 @@ namespace Xmpp.Sasl {
         public void received_features_node(XmppStream stream) {
             if (stream.has_flag(Flag.IDENTITY)) return;
             if (stream.is_setup_needed()) return;
-            if (!stream.has_flag(Tls.Flag.IDENTITY) || !stream.get_flag(Tls.Flag.IDENTITY).finished) return;
 
             var mechanisms = stream.features.get_subnode("mechanisms", NS_URI);
             string[] supported_mechanisms = {};
diff --git a/xmpp-vala/src/module/tls.vala b/xmpp-vala/src/module/tls.vala
index c3afc4b3..1b8a5411 100644
--- a/xmpp-vala/src/module/tls.vala
+++ b/xmpp-vala/src/module/tls.vala
@@ -23,10 +23,11 @@ namespace Xmpp.Tls {
         private void received_nonza(XmppStream stream, StanzaNode node) {
             if (node.ns_uri == NS_URI && node.name == "proceed") {
                 try {
-                    var io_stream = stream.get_stream();
+                    StartTlsXmppStream? tls_xmpp_stream = stream as StartTlsXmppStream;
+                    var io_stream = tls_xmpp_stream.get_stream();
                     if (io_stream == null) return;
                     var conn = TlsClientConnection.new(io_stream, identity);
-                    stream.reset_stream(conn);
+                    tls_xmpp_stream.reset_stream(conn);
 
                     conn.accept_certificate.connect(on_invalid_certificate);
                     var flag = stream.get_flag(Flag.IDENTITY);
diff --git a/xmpp-vala/src/module/xep/0030_service_discovery/module.vala b/xmpp-vala/src/module/xep/0030_service_discovery/module.vala
index 42547f62..537e460b 100644
--- a/xmpp-vala/src/module/xep/0030_service_discovery/module.vala
+++ b/xmpp-vala/src/module/xep/0030_service_discovery/module.vala
@@ -23,7 +23,10 @@ public class Module : XmppStreamModule, Iq.Handler {
     }
 
     public void remove_feature(XmppStream stream, string feature) {
-        stream.get_flag(Flag.IDENTITY).remove_own_feature(feature);
+        Flag? flag = stream.get_flag(Flag.IDENTITY);
+        if (flag != null) {
+                flag.remove_own_feature(feature);
+        }
     }
 
     public void add_feature_notify(XmppStream stream, string feature) {
@@ -34,14 +37,6 @@ public class Module : XmppStreamModule, Iq.Handler {
         remove_feature(stream, feature + "+notify");
     }
 
-    public void add_identity(XmppStream stream, Identity identity) {
-        stream.get_flag(Flag.IDENTITY).add_own_identity(identity);
-    }
-
-    public void remove_identity(XmppStream stream, Identity identity) {
-        stream.get_flag(Flag.IDENTITY).remove_own_identity(identity);
-    }
-
     public async bool has_entity_feature(XmppStream stream, Jid jid, string feature) {
         return yield this.cache.has_entity_feature(jid, feature);
     }
@@ -93,7 +88,7 @@ public class Module : XmppStreamModule, Iq.Handler {
 
     public override void attach(XmppStream stream) {
         stream.add_flag(new Flag());
-        add_identity(stream, own_identity);
+        stream.get_flag(Flag.IDENTITY).add_own_identity(own_identity);
 
         stream.get_module(Iq.Module.IDENTITY).register_for_namespace(NS_URI_INFO, this);
         add_feature(stream, NS_URI_INFO);
@@ -102,7 +97,8 @@ public class Module : XmppStreamModule, Iq.Handler {
     public override void detach(XmppStream stream) {
         active_info_requests.clear();
 
-        remove_identity(stream, own_identity);
+        Flag? flag = stream.get_flag(Flag.IDENTITY);
+        if (flag != null) flag.remove_own_identity(own_identity);
 
         stream.get_module(Iq.Module.IDENTITY).unregister_from_namespace(NS_URI_INFO, this);
         remove_feature(stream, NS_URI_INFO);
diff --git a/xmpp-vala/src/module/xep/0198_stream_management.vala b/xmpp-vala/src/module/xep/0198_stream_management.vala
index b6317425..b9185808 100644
--- a/xmpp-vala/src/module/xep/0198_stream_management.vala
+++ b/xmpp-vala/src/module/xep/0198_stream_management.vala
@@ -39,7 +39,7 @@ public class Module : XmppStreamNegotiationModule, WriteNodeFunc {
     }
 
     internal async void write_node(XmppStream stream, StanzaNode node) {
-        StanzaWriter? writer = stream.writer;
+        StanzaWriter? writer = ((IoXmppStream)stream).writer;
         if (writer == null) return;
         try {
             stream.log.node("OUT", node, stream);
@@ -104,14 +104,11 @@ public class Module : XmppStreamNegotiationModule, WriteNodeFunc {
 
     private void check_resume(XmppStream stream) {
         if (stream_has_sm_feature(stream) && session_id != null) {
-            Tls.Flag? tls_flag = stream.get_flag(Tls.Flag.IDENTITY);
-            if (tls_flag != null && tls_flag.finished) {
-                StanzaNode node = new StanzaNode.build("resume", NS_URI).add_self_xmlns()
-                    .put_attribute("h", h_inbound.to_string())
-                    .put_attribute("previd", session_id);
-                write_node.begin(stream, node);
-                stream.add_flag(new Flag());
-            }
+            StanzaNode node = new StanzaNode.build("resume", NS_URI).add_self_xmlns()
+                .put_attribute("h", h_inbound.to_string())
+                .put_attribute("previd", session_id);
+            write_node.begin(stream, node);
+            stream.add_flag(new Flag());
         }
     }
 
@@ -137,7 +134,7 @@ public class Module : XmppStreamNegotiationModule, WriteNodeFunc {
                     h_inbound = 0;
                     session_id = node.get_attribute("id", NS_URI);
                     flags = stream.flags;
-                    stream.write_obj = this;
+                    ((IoXmppStream)stream).write_obj = this;
                 } else if (node.name == "resumed") {
                     stream.get_flag(Flag.IDENTITY).resumed = true;
 
@@ -152,7 +149,7 @@ public class Module : XmppStreamNegotiationModule, WriteNodeFunc {
                     }
                     in_flight_stanzas.clear();
                     check_queue(stream);
-                    stream.write_obj = this;
+                    ((IoXmppStream)stream).write_obj = this;
                 } else if (node.name == "failed") {
                     stream.received_features_node(stream);
                     session_id = null;
diff --git a/xmpp-vala/src/module/xep/0368_srv_records_tls.vala b/xmpp-vala/src/module/xep/0368_srv_records_tls.vala
deleted file mode 100644
index 5a2a4559..00000000
--- a/xmpp-vala/src/module/xep/0368_srv_records_tls.vala
+++ /dev/null
@@ -1,56 +0,0 @@
-using Gee;
-
-namespace Xmpp.Xep.SrvRecordsTls {
-
-public class Module : XmppStreamNegotiationModule {
-    public static ModuleIdentity<Module> IDENTITY = new ModuleIdentity<Module>("", "0363_srv_records_for_xmpp_over_tls");
-
-    public override void attach(XmppStream stream) {
-        stream.register_connection_provider(new TlsConnectionProvider());
-    }
-
-    public override void detach(XmppStream stream) { }
-
-    public override bool mandatory_outstanding(XmppStream stream) { return false; }
-    public override bool negotiation_active(XmppStream stream) { return false; }
-    public override string get_ns() { return IDENTITY.ns; }
-    public override string get_id() { return IDENTITY.id; }
-}
-
-public class TlsConnectionProvider : ConnectionProvider {
-    private SrvTarget? srv_target;
-
-    public async override int? get_priority(Jid remote_name) {
-        GLib.List<SrvTarget>? xmpp_target = null;
-        try {
-            GLibFixes.Resolver resolver = GLibFixes.Resolver.get_default();
-            xmpp_target = yield resolver.lookup_service_async("xmpps-client", "tcp", remote_name.to_string(), null);
-        } catch (Error e) {
-            return null;
-        }
-        xmpp_target.sort((a, b) => { return a.get_priority() - b.get_priority(); });
-        srv_target = xmpp_target.nth(0).data;
-        return xmpp_target.nth(0).data.get_priority();
-    }
-
-    public async override IOStream? connect(XmppStream stream) {
-        SocketClient client = new SocketClient();
-        try {
-            debug("Connecting to %s %i (tls)", srv_target.get_hostname(), srv_target.get_port());
-            IOStream? io_stream = yield client.connect_to_host_async(srv_target.get_hostname(), srv_target.get_port());
-            TlsConnection tls_connection = TlsClientConnection.new(io_stream, new NetworkAddress(stream.remote_name.to_string(), srv_target.get_port()));
-#if ALPN_SUPPORT
-            tls_connection.set_advertised_protocols(new string[]{"xmpp-client"});
-#endif
-            tls_connection.accept_certificate.connect(stream.get_module(Tls.Module.IDENTITY).on_invalid_certificate);
-            stream.add_flag(new Tls.Flag() { finished=true });
-            return tls_connection;
-        } catch (Error e) {
-            return null;
-        }
-    }
-
-    public override string get_id() { return "srv_records"; }
-}
-
-}