Add support for custom TXT records in NSD

Change-Id: I8e6dc9852ad4d273c71ad6a63a7fbd28a206806d

4 files changed, 352 insertions, 28 deletions

diff --git a/api/current.txt b/api/current.txt
index f93ea06af1b2..ff84764312c6 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -14415,10 +14415,13 @@ package android.net.nsd {
   public final class NsdServiceInfo implements android.os.Parcelable {
     ctor public NsdServiceInfo();
     method public int describeContents();
+    method public java.util.Map<java.lang.String, byte[]> getAttributes();
     method public java.net.InetAddress getHost();
     method public int getPort();
     method public java.lang.String getServiceName();
     method public java.lang.String getServiceType();
+    method public void removeAttribute(java.lang.String);
+    method public void setAttribute(java.lang.String, java.lang.String);
     method public void setHost(java.net.InetAddress);
     method public void setPort(int);
     method public void setServiceName(java.lang.String);
diff --git a/core/java/android/net/nsd/NsdServiceInfo.java b/core/java/android/net/nsd/NsdServiceInfo.java
index 205a21d95b25..6fdb0d0f8462 100644
--- a/core/java/android/net/nsd/NsdServiceInfo.java
+++ b/core/java/android/net/nsd/NsdServiceInfo.java
@@ -18,8 +18,15 @@ package android.net.nsd;
 
 import android.os.Parcelable;
 import android.os.Parcel;
+import android.util.Log;
+import android.util.ArrayMap;
 
+import java.io.UnsupportedEncodingException;
 import java.net.InetAddress;
+import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+import java.util.Map;
+
 
 /**
  * A class representing service information for network service discovery
@@ -27,11 +34,13 @@ import java.net.InetAddress;
  */
 public final class NsdServiceInfo implements Parcelable {
 
+    private static final String TAG = "NsdServiceInfo";
+
     private String mServiceName;
 
     private String mServiceType;
 
-    private DnsSdTxtRecord mTxtRecord;
+    private final ArrayMap<String, byte[]> mTxtRecord = new ArrayMap<String, byte[]>();
 
     private InetAddress mHost;
 
@@ -41,10 +50,9 @@ public final class NsdServiceInfo implements Parcelable {
     }
 
     /** @hide */
-    public NsdServiceInfo(String sn, String rt, DnsSdTxtRecord tr) {
+    public NsdServiceInfo(String sn, String rt) {
         mServiceName = sn;
         mServiceType = rt;
-        mTxtRecord = tr;
     }
 
     /** Get the service name */
@@ -67,16 +75,6 @@ public final class NsdServiceInfo implements Parcelable {
         mServiceType = s;
     }
 
-    /** @hide */
-    public DnsSdTxtRecord getTxtRecord() {
-        return mTxtRecord;
-    }
-
-    /** @hide */
-    public void setTxtRecord(DnsSdTxtRecord t) {
-        mTxtRecord = new DnsSdTxtRecord(t);
-    }
-
     /** Get the host address. The host address is valid for a resolved service. */
     public InetAddress getHost() {
         return mHost;
@@ -97,14 +95,134 @@ public final class NsdServiceInfo implements Parcelable {
         mPort = p;
     }
 
+    /** @hide */
+    public void setAttribute(String key, byte[] value) {
+        // Key must be printable US-ASCII, excluding =.
+        for (int i = 0; i < key.length(); ++i) {
+            char character = key.charAt(i);
+            if (character < 0x20 || character > 0x7E) {
+                throw new IllegalArgumentException("Key strings must be printable US-ASCII");
+            } else if (character == 0x3D) {
+                throw new IllegalArgumentException("Key strings must not include '='");
+            }
+        }
+
+        // Key length + value length must be < 255.
+        if (key.length() + (value == null ? 0 : value.length) >= 255) {
+            throw new IllegalArgumentException("Key length + value length must be < 255 bytes");
+        }
+
+        // Warn if key is > 9 characters, as recommended by RFC 6763 section 6.4.
+        if (key.length() > 9) {
+            Log.w(TAG, "Key lengths > 9 are discouraged: " + key);
+        }
+
+        // Check against total TXT record size limits.
+        // Arbitrary 400 / 1300 byte limits taken from RFC 6763 section 6.2.
+        int txtRecordSize = getTxtRecordSize();
+        int futureSize = txtRecordSize + key.length() + (value == null ? 0 : value.length) + 2;
+        if (futureSize > 1300) {
+            throw new IllegalArgumentException("Total length of attributes must be < 1300 bytes");
+        } else if (futureSize > 400) {
+            Log.w(TAG, "Total length of all attributes exceeds 400 bytes; truncation may occur");
+        }
+
+        mTxtRecord.put(key, value);
+    }
+
+    /**
+     * Add a service attribute as a key/value pair.
+     *
+     * <p> Service attributes are included as DNS-SD TXT record pairs.
+     *
+     * <p> The key must be US-ASCII printable characters, excluding the '=' character.  Values may
+     * be UTF-8 strings or null.  The total length of key + value must be less than 255 bytes.
+     *
+     * <p> Keys should be short, ideally no more than 9 characters, and unique per instance of
+     * {@link NsdServiceInfo}.  Calling {@link #setAttribute} twice with the same key will overwrite
+     * first value.
+     */
+    public void setAttribute(String key, String value) {
+        try {
+            setAttribute(key, value == null ? (byte []) null : value.getBytes("UTF-8"));
+        } catch (UnsupportedEncodingException e) {
+            throw new IllegalArgumentException("Value must be UTF-8");
+        }
+    }
+
+    /** Remove an attribute by key */
+    public void removeAttribute(String key) {
+        mTxtRecord.remove(key);
+    }
+
+    /**
+     * Retrive attributes as a map of String keys to byte[] values.
+     *
+     * <p> The returned map is unmodifiable; changes must be made through {@link #setAttribute} and
+     * {@link #removeAttribute}.
+     */
+    public Map<String, byte[]> getAttributes() {
+        return Collections.unmodifiableMap(mTxtRecord);
+    }
+
+    private int getTxtRecordSize() {
+        int txtRecordSize = 0;
+        for (Map.Entry<String, byte[]> entry : mTxtRecord.entrySet()) {
+            txtRecordSize += 2;  // One for the length byte, one for the = between key and value.
+            txtRecordSize += entry.getKey().length();
+            byte[] value = entry.getValue();
+            txtRecordSize += value == null ? 0 : value.length;
+        }
+        return txtRecordSize;
+    }
+
+    /** @hide */
+    public byte[] getTxtRecord() {
+        int txtRecordSize = getTxtRecordSize();
+        if (txtRecordSize == 0) {
+            return null;
+        }
+
+        byte[] txtRecord = new byte[txtRecordSize];
+        int ptr = 0;
+        for (Map.Entry<String, byte[]> entry : mTxtRecord.entrySet()) {
+            String key = entry.getKey();
+            byte[] value = entry.getValue();
+
+            // One byte to record the length of this key/value pair.
+            txtRecord[ptr++] = (byte) (key.length() + (value == null ? 0 : value.length) + 1);
+
+            // The key, in US-ASCII.
+            // Note: use the StandardCharsets const here because it doesn't raise exceptions and we
+            // already know the key is ASCII at this point.
+            System.arraycopy(key.getBytes(StandardCharsets.US_ASCII), 0, txtRecord, ptr,
+                    key.length());
+            ptr += key.length();
+
+            // US-ASCII '=' character.
+            txtRecord[ptr++] = (byte)'=';
+
+            // The value, as any raw bytes.
+            if (value != null) {
+                System.arraycopy(value, 0, txtRecord, ptr, value.length);
+                ptr += value.length;
+            }
+        }
+        return txtRecord;
+    }
+
     public String toString() {
         StringBuffer sb = new StringBuffer();
 
-        sb.append("name: ").append(mServiceName).
-            append("type: ").append(mServiceType).
-            append("host: ").append(mHost).
-            append("port: ").append(mPort).
-            append("txtRecord: ").append(mTxtRecord);
+        sb.append("name: ").append(mServiceName)
+                .append(", type: ").append(mServiceType)
+                .append(", host: ").append(mHost)
+                .append(", port: ").append(mPort);
+
+        byte[] txtRecord = getTxtRecord();
+        if (txtRecord != null) {
+            sb.append(", txtRecord: ").append(new String(txtRecord, StandardCharsets.UTF_8));
+        }
         return sb.toString();
     }
 
@@ -117,14 +235,27 @@ public final class NsdServiceInfo implements Parcelable {
     public void writeToParcel(Parcel dest, int flags) {
         dest.writeString(mServiceName);
         dest.writeString(mServiceType);
-        dest.writeParcelable(mTxtRecord, flags);
         if (mHost != null) {
-            dest.writeByte((byte)1);
+            dest.writeInt(1);
             dest.writeByteArray(mHost.getAddress());
         } else {
-            dest.writeByte((byte)0);
+            dest.writeInt(0);
         }
         dest.writeInt(mPort);
+
+        // TXT record key/value pairs.
+        dest.writeInt(mTxtRecord.size());
+        for (String key : mTxtRecord.keySet()) {
+            byte[] value = mTxtRecord.get(key);
+            if (value != null) {
+                dest.writeInt(1);
+                dest.writeInt(value.length);
+                dest.writeByteArray(value);
+            } else {
+                dest.writeInt(0);
+            }
+            dest.writeString(key);
+        }
     }
 
     /** Implement the Parcelable interface */
@@ -134,15 +265,26 @@ public final class NsdServiceInfo implements Parcelable {
                 NsdServiceInfo info = new NsdServiceInfo();
                 info.mServiceName = in.readString();
                 info.mServiceType = in.readString();
-                info.mTxtRecord = in.readParcelable(null);
 
-                if (in.readByte() == 1) {
+                if (in.readInt() == 1) {
                     try {
                         info.mHost = InetAddress.getByAddress(in.createByteArray());
                     } catch (java.net.UnknownHostException e) {}
                 }
 
                 info.mPort = in.readInt();
+
+                // TXT record key/value pairs.
+                int recordCount = in.readInt();
+                for (int i = 0; i < recordCount; ++i) {
+                    byte[] valueArray = null;
+                    if (in.readInt() == 1) {
+                        int valueLength = in.readInt();
+                        valueArray = new byte[valueLength];
+                        in.readByteArray(valueArray);
+                    }
+                    info.mTxtRecord.put(in.readString(), valueArray);
+                }
                 return info;
             }
 
diff --git a/services/core/java/com/android/server/NsdService.java b/services/core/java/com/android/server/NsdService.java
index 93799550dcd7..8df93f128354 100644
--- a/services/core/java/com/android/server/NsdService.java
+++ b/services/core/java/com/android/server/NsdService.java
@@ -38,10 +38,13 @@ import android.util.SparseArray;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
+import java.io.UnsupportedEncodingException;
 import java.net.InetAddress;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Locale;
+import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 
 import com.android.internal.app.IBatteryStats;
@@ -443,14 +446,14 @@ public class NsdService extends INsdManager.Stub {
                     case NativeResponseCode.SERVICE_FOUND:
                         /* NNN uniqueId serviceName regType domain */
                         if (DBG) Slog.d(TAG, "SERVICE_FOUND Raw: " + raw);
-                        servInfo = new NsdServiceInfo(cooked[2], cooked[3], null);
+                        servInfo = new NsdServiceInfo(cooked[2], cooked[3]);
                         clientInfo.mChannel.sendMessage(NsdManager.SERVICE_FOUND, 0,
                                 clientId, servInfo);
                         break;
                     case NativeResponseCode.SERVICE_LOST:
                         /* NNN uniqueId serviceName regType domain */
                         if (DBG) Slog.d(TAG, "SERVICE_LOST Raw: " + raw);
-                        servInfo = new NsdServiceInfo(cooked[2], cooked[3], null);
+                        servInfo = new NsdServiceInfo(cooked[2], cooked[3]);
                         clientInfo.mChannel.sendMessage(NsdManager.SERVICE_LOST, 0,
                                 clientId, servInfo);
                         break;
@@ -463,7 +466,7 @@ public class NsdService extends INsdManager.Stub {
                     case NativeResponseCode.SERVICE_REGISTERED:
                         /* NNN regId serviceName regType */
                         if (DBG) Slog.d(TAG, "SERVICE_REGISTERED Raw: " + raw);
-                        servInfo = new NsdServiceInfo(cooked[2], null, null);
+                        servInfo = new NsdServiceInfo(cooked[2], null);
                         clientInfo.mChannel.sendMessage(NsdManager.REGISTER_SERVICE_SUCCEEDED,
                                 id, clientId, servInfo);
                         break;
@@ -679,9 +682,22 @@ public class NsdService extends INsdManager.Stub {
     private boolean registerService(int regId, NsdServiceInfo service) {
         if (DBG) Slog.d(TAG, "registerService: " + regId + " " + service);
         try {
-            //Add txtlen and txtdata
-            mNativeConnector.execute("mdnssd", "register", regId, service.getServiceName(),
+            Command cmd = new Command("mdnssd", "register", regId, service.getServiceName(),
                     service.getServiceType(), service.getPort());
+
+            // Add TXT records as additional arguments.
+            Map<String, byte[]> txtRecords = service.getAttributes();
+            for (String key : txtRecords.keySet()) {
+                try {
+                    // TODO: Send encoded TXT record as bytes once NDC/netd supports binary data.
+                    cmd.appendArg(String.format(Locale.US, "%s=%s", key,
+                            new String(txtRecords.get(key), "UTF_8")));
+                } catch (UnsupportedEncodingException e) {
+                    Slog.e(TAG, "Failed to encode txtRecord " + e);
+                }
+            }
+
+            mNativeConnector.execute(cmd);
         } catch(NativeDaemonConnectorException e) {
             Slog.e(TAG, "Failed to execute registerService " + e);
             return false;
diff --git a/tests/CoreTests/android/core/NsdServiceInfoTest.java b/tests/CoreTests/android/core/NsdServiceInfoTest.java
new file mode 100644
index 000000000000..5bf0167f1cf0
--- /dev/null
+++ b/tests/CoreTests/android/core/NsdServiceInfoTest.java
@@ -0,0 +1,163 @@
+package android.core;
+
+import android.test.AndroidTestCase;
+
+import android.os.Bundle;
+import android.os.Parcel;
+import android.os.StrictMode;
+import android.net.nsd.NsdServiceInfo;
+import android.util.Log;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+
+
+public class NsdServiceInfoTest extends AndroidTestCase {
+
+    public final static InetAddress LOCALHOST;
+    static {
+        // Because test.
+        StrictMode.ThreadPolicy policy = new StrictMode.ThreadPolicy.Builder().permitAll().build();
+        StrictMode.setThreadPolicy(policy);
+
+        InetAddress _host = null;
+        try {
+            _host = InetAddress.getLocalHost();
+        } catch (UnknownHostException e) { }
+        LOCALHOST = _host;
+    }
+
+    public void testLimits() throws Exception {
+        NsdServiceInfo info = new NsdServiceInfo();
+
+        // Non-ASCII keys.
+        boolean exceptionThrown = false;
+        try {
+            info.setAttribute("猫", "meow");
+        } catch (IllegalArgumentException e) {
+            exceptionThrown = true;
+        }
+        assertTrue(exceptionThrown);
+        assertEmptyServiceInfo(info);
+
+        // ASCII keys with '=' character.
+        exceptionThrown = false;
+        try {
+            info.setAttribute("kitten=", "meow");
+        } catch (IllegalArgumentException e) {
+            exceptionThrown = true;
+        }
+        assertTrue(exceptionThrown);
+        assertEmptyServiceInfo(info);
+
+        // Single key + value length too long.
+        exceptionThrown = false;
+        try {
+            String longValue = "loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo" +
+                    "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo" +
+                    "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo" +
+                    "ooooooooooooooooooooooooooooong";  // 248 characters.
+            info.setAttribute("longcat", longValue);  // Key + value == 255 characters.
+        } catch (IllegalArgumentException e) {
+            exceptionThrown = true;
+        }
+        assertTrue(exceptionThrown);
+        assertEmptyServiceInfo(info);
+
+        // Total TXT record length too long.
+        exceptionThrown = false;
+        int recordsAdded = 0;
+        try {
+            for (int i = 100; i < 300; ++i) {
+                // 6 char key + 5 char value + 2 bytes overhead = 13 byte record length.
+                String key = String.format("key%d", i);
+                info.setAttribute(key, "12345");
+                recordsAdded++;
+            }
+        } catch (IllegalArgumentException e) {
+            exceptionThrown = true;
+        }
+        assertTrue(exceptionThrown);
+        assertTrue(100 == recordsAdded);
+        assertTrue(info.getTxtRecord().length == 1300);
+    }
+
+    public void testParcel() throws Exception {
+        NsdServiceInfo emptyInfo = new NsdServiceInfo();
+        checkParcelable(emptyInfo);
+
+        NsdServiceInfo fullInfo = new NsdServiceInfo();
+        fullInfo.setServiceName("kitten");
+        fullInfo.setServiceType("_kitten._tcp");
+        fullInfo.setPort(4242);
+        fullInfo.setHost(LOCALHOST);
+        checkParcelable(fullInfo);
+
+        NsdServiceInfo noHostInfo = new NsdServiceInfo();
+        noHostInfo.setServiceName("kitten");
+        noHostInfo.setServiceType("_kitten._tcp");
+        noHostInfo.setPort(4242);
+        checkParcelable(noHostInfo);
+
+        NsdServiceInfo attributedInfo = new NsdServiceInfo();
+        attributedInfo.setServiceName("kitten");
+        attributedInfo.setServiceType("_kitten._tcp");
+        attributedInfo.setPort(4242);
+        attributedInfo.setHost(LOCALHOST);
+        attributedInfo.setAttribute("color", "pink");
+        attributedInfo.setAttribute("sound", (new String("にゃあ")).getBytes("UTF-8"));
+        attributedInfo.setAttribute("adorable", (String) null);
+        attributedInfo.setAttribute("sticky", "yes");
+        attributedInfo.setAttribute("siblings", new byte[] {});
+        attributedInfo.setAttribute("edge cases", new byte[] {0, -1, 127, -128});
+        attributedInfo.removeAttribute("sticky");
+        checkParcelable(attributedInfo);
+
+        // Sanity check that we actually wrote attributes to attributedInfo.
+        assertTrue(attributedInfo.getAttributes().keySet().contains("adorable"));
+        String sound = new String(attributedInfo.getAttributes().get("sound"), "UTF-8");
+        assertTrue(sound.equals("にゃあ"));
+        byte[] edgeCases = attributedInfo.getAttributes().get("edge cases");
+        assertTrue(Arrays.equals(edgeCases, new byte[] {0, -1, 127, -128}));
+        assertFalse(attributedInfo.getAttributes().keySet().contains("sticky"));
+    }
+
+    public void checkParcelable(NsdServiceInfo original) {
+        // Write to parcel.
+        Parcel p = Parcel.obtain();
+        Bundle writer = new Bundle();
+        writer.putParcelable("test_info", original);
+        writer.writeToParcel(p, 0);
+
+        // Extract from parcel.
+        p.setDataPosition(0);
+        Bundle reader = p.readBundle();
+        reader.setClassLoader(NsdServiceInfo.class.getClassLoader());
+        NsdServiceInfo result = reader.getParcelable("test_info");
+
+        // Assert equality of base fields.
+        assertEquality(original.getServiceName(), result.getServiceName());
+        assertEquality(original.getServiceType(), result.getServiceType());
+        assertEquality(original.getHost(), result.getHost());
+        assertTrue(original.getPort() == result.getPort());
+
+        // Assert equality of attribute map.
+        Map<String, byte[]> originalMap = original.getAttributes();
+        Map<String, byte[]> resultMap = result.getAttributes();
+        assertEquality(originalMap.keySet(), resultMap.keySet());
+        for (String key : originalMap.keySet()) {
+            assertTrue(Arrays.equals(originalMap.get(key), resultMap.get(key)));
+        }
+    }
+
+    public void assertEquality(Object expected, Object result) {
+        assertTrue(expected == result || expected.equals(result));
+    }
+
+    public void assertEmptyServiceInfo(NsdServiceInfo shouldBeEmpty) {
+        assertTrue(null == shouldBeEmpty.getTxtRecord());
+    }
+}