From eaafaacef1e3d6bb66cf36ec10b3090b3584a8e9 Mon Sep 17 00:00:00 2001
From: wongwong <luckboystudent@gmail.com>
Date: Fri, 22 Feb 2019 14:26:07 +0800
Subject: [PATCH] add setting unicast bit

---
 index.js | 15 +++++++++++++--
 test.js  | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/index.js b/index.js
index d55b688..513b99e 100644
--- a/index.js
+++ b/index.js
@@ -1608,8 +1608,12 @@ question.encode = function (q, buf, offset) {
 
   buf.writeUInt16BE(types.toType(q.type), offset)
   offset += 2
+  if (q.qu === undefined || q.qu !== true) {
+    buf.writeUInt16BE(classes.toClass(q.class === undefined ? 'IN' : q.class), offset)
+  } else {
+    buf.writeUInt16BE(classes.toClass(q.class === undefined ? 'IN' : q.class) + QU_MASK, offset)
+  }
 
-  buf.writeUInt16BE(classes.toClass(q.class === undefined ? 'IN' : q.class), offset)
   offset += 2
 
   question.encode.bytes = offset - oldOffset
@@ -1630,7 +1634,14 @@ question.decode = function (buf, offset) {
   q.type = types.toString(buf.readUInt16BE(offset))
   offset += 2
 
-  q.class = classes.toString(buf.readUInt16BE(offset))
+  const classValue = buf.readUInt16BE(offset)
+  if (classValue >= QU_MASK) {
+    q.qu = true
+    q.class = classes.toString(classValue - QU_MASK)
+  } else {
+    q.class = classes.toString(classValue)
+  }
+
   offset += 2
 
   const qu = !!(q.class & QU_MASK)
diff --git a/test.js b/test.js
index e4d8669..24273d7 100644
--- a/test.js
+++ b/test.js
@@ -163,6 +163,45 @@ tape('query', function (t) {
       name: 'hello.srv.com'
     }]
   })
+  testEncoder(t, packet, {
+    type: 'query',
+    questions: [{
+      type: 'A',
+      qu: true,
+      name: 'hello.a.com'
+    }, {
+      type: 'SRV',
+      name: 'hello.srv.com'
+    }]
+  })
+
+  testEncoder(t, packet, {
+    type: 'query',
+    id: 42,
+    questions: [{
+      type: 'A',
+      qu: true,
+      class: 'IN',
+      name: 'hello.a.com'
+    }, {
+      type: 'SRV',
+      name: 'hello.srv.com'
+    }]
+  })
+
+  testEncoder(t, packet, {
+    type: 'query',
+    id: 42,
+    questions: [{
+      type: 'A',
+      qu: true,
+      class: 'CH',
+      name: 'hello.a.com'
+    }, {
+      type: 'SRV',
+      name: 'hello.srv.com'
+    }]
+  })
 
   t.end()
 })