From 0e0380917f6b3a1a03d33f536879eb63ca68d93b Mon Sep 17 00:00:00 2001
From: Zaid Abdulla <zaid@genieacs.com>
Date: Mon, 9 Nov 2020 17:22:48 -0800
Subject: [PATCH] XMPP fixes

---
 build/build.ts            |  2 ++
 lib/connection-request.ts | 11 +++----
 lib/xmpp-client.ts        | 68 +++++++++++++++++++++++++++++++++------
 3 files changed, 64 insertions(+), 17 deletions(-)

diff --git a/build/build.ts b/build/build.ts
index c4969bd..451e924 100644
--- a/build/build.ts
+++ b/build/build.ts
@@ -51,6 +51,8 @@ const externals = [
   "tls",
   "http",
   "https",
+  "net",
+  "events",
   "zlib",
   "crypto",
   "util",
diff --git a/lib/connection-request.ts b/lib/connection-request.ts
index 92a6a37..f474791 100644
--- a/lib/connection-request.ts
+++ b/lib/connection-request.ts
@@ -244,11 +244,8 @@ const XMPP_PASSWORD = config.get("XMPP_PASSWORD") as string;
 const XMPP_RESOURCE = crypto.randomBytes(8).toString("hex");
 
 let xmppClient: XmppClient;
-let xmppclientTimeout: NodeJS.Timeout;
 
 function xmppClientOnError(err: Error): void {
-  clearTimeout(xmppclientTimeout);
-  xmppClient.disconnect();
   xmppClient = null;
   logger.error({
     message: "XMPP exception",
@@ -257,8 +254,7 @@ function xmppClientOnError(err: Error): void {
   });
 }
 
-function xmppClientTimeoutCallback(): void {
-  xmppClient.disconnect();
+function xmppClientOnClose(): void {
   xmppClient = null;
 }
 
@@ -276,11 +272,12 @@ export async function xmppConnectionRequest(
       username,
       resource: XMPP_RESOURCE,
       password: XMPP_PASSWORD,
+      timeout: 120000,
     });
     xmppClient.on("error", xmppClientOnError);
+    xmppClient.on("close", xmppClientOnClose);
+    xmppClient.unref();
   }
-  clearTimeout(xmppclientTimeout);
-  setTimeout(xmppClientTimeoutCallback, 120000);
 
   let username: string;
   let password: string;
diff --git a/lib/xmpp-client.ts b/lib/xmpp-client.ts
index 81ca265..4d1d6de 100644
--- a/lib/xmpp-client.ts
+++ b/lib/xmpp-client.ts
@@ -131,7 +131,29 @@ function saltPassword(
   return hi;
 }
 
-function* login(
+function* loginPlain(
+  socket: net.Socket,
+  username: string,
+  password: string
+): Generator<void, void, Element> {
+  socket.write(
+    `<auth xmlns="urn:ietf:params:xml:ns:xmpp-sasl" mechanism="PLAIN">${encodeBase64(
+      `\x00${username}\x00${password}`
+    )}</auth>`
+  );
+  const res1 = yield;
+
+  if (
+    res1.name === "failure" &&
+    res1.children.some((c) => c.name === "not-authorized")
+  )
+    throw new Error("Not authorized");
+
+  if (res1.name !== "success")
+    throw new Error(`Unexpected response ${res1.name}`);
+}
+
+function* loginScram(
   socket: net.Socket,
   username: string,
   password: string
@@ -263,9 +285,14 @@ function* init(
       const mechanisms: Set<string> = new Set(
         feature.children.map((c) => c.text)
       );
-      if (mechanisms.has("SCRAM-SHA-1")) {
-        yield* login(socket, username, password);
+      if (mechanisms.has("PLAIN")) {
+        yield* loginPlain(socket, username, password);
         return STATUS_RESTART_STREAM;
+      } else if (mechanisms.has("SCRAM-SHA-1")) {
+        yield* loginScram(socket, username, password);
+        return STATUS_RESTART_STREAM;
+      } else {
+        throw new Error("No supported SASL method");
       }
     } else if (feature.name === "bind") {
       if (feature.children.some((c) => c.name === "required")) {
@@ -293,6 +320,7 @@ interface XmppClientOptions {
   username?: string;
   password?: string;
   resource?: string;
+  timeout?: number;
 }
 
 export default class XmppClient extends EventEmitter {
@@ -343,8 +371,10 @@ export default class XmppClient extends EventEmitter {
           client._host = opts.host;
           client._username = opts.username;
           client._resource = opts.resource;
-          socket.on("data", client._onData);
-          socket.on("error", client._onError);
+          socket.on("data", client._onData.bind(client));
+          socket.on("error", client._onError.bind(client));
+          if (opts.timeout)
+            socket.setTimeout(opts.timeout, client.close.bind(client));
           resolve(client);
         } catch (err) {
           socket.destroy();
@@ -354,10 +384,16 @@ export default class XmppClient extends EventEmitter {
     });
   }
 
-  disconnect(): void {
-    this._socket.removeListener("data", this._onData);
-    this._socket.removeListener("error", this._onError);
-    this._socket.end();
+  close(): void {
+    this._socket.end("</stream>");
+  }
+
+  ref(): void {
+    this._socket.ref();
+  }
+
+  unref(): void {
+    this._socket.unref();
   }
 
   get host(): string {
@@ -374,7 +410,12 @@ export default class XmppClient extends EventEmitter {
 
   private _onData(chunk: Buffer): void {
     try {
-      const str = chunk.toString("utf8");
+      let close = false;
+      let str = chunk.toString("utf8");
+      if (str.endsWith("</stream:stream>")) {
+        str = str.slice(0, -16);
+        close = true;
+      }
       const xml = parseXml(str);
       const idx = xml.children.map((c) => c.bodyIndex);
       for (const [i, c] of xml.children.entries()) {
@@ -389,7 +430,14 @@ export default class XmppClient extends EventEmitter {
         }
         this.emit("stanza", c, s);
       }
+      if (close) {
+        this._socket.removeAllListeners("data");
+        this._socket.removeAllListeners("error");
+        this.emit("close");
+      }
     } catch (err) {
+      this._socket.removeAllListeners("data");
+      this._socket.removeAllListeners("error");
       this.emit("error", err);
     }
   }