Merge "[AWARE] Enhance TLV utils to support new TLV structs in NAN"
diff --git a/wifi/java/android/net/wifi/aware/TlvBufferUtils.java b/wifi/java/android/net/wifi/aware/TlvBufferUtils.java
index 29f10e9..b3b5b29 100644
--- a/wifi/java/android/net/wifi/aware/TlvBufferUtils.java
+++ b/wifi/java/android/net/wifi/aware/TlvBufferUtils.java
@@ -61,6 +61,7 @@
     public static class TlvConstructor {
         private int mTypeSize;
         private int mLengthSize;
+        private ByteOrder mByteOrder = ByteOrder.BIG_ENDIAN;
 
         private byte[] mArray;
         private int mArrayLength;
@@ -84,6 +85,20 @@
             }
             mTypeSize = typeSize;
             mLengthSize = lengthSize;
+            mPosition = 0;
+        }
+
+        /**
+         * Configure the TLV constructor to use a particular byte order. Should be
+         * {@link ByteOrder#BIG_ENDIAN} (the default at construction) or
+         * {@link ByteOrder#LITTLE_ENDIAN}.
+         *
+         * @return The constructor to facilitate chaining
+         *         {@code ctr.putXXX(..).putXXX(..)}.
+         */
+        public TlvConstructor setByteOrder(ByteOrder byteOrder) {
+            mByteOrder = byteOrder;
+            return this;
         }
 
         /**
@@ -96,6 +111,7 @@
         public TlvConstructor wrap(@Nullable byte[] array) {
             mArray = array;
             mArrayLength = (array == null) ? 0 : array.length;
+            mPosition = 0;
             return this;
         }
 
@@ -109,6 +125,7 @@
         public TlvConstructor allocate(int capacity) {
             mArray = new byte[capacity];
             mArrayLength = capacity;
+            mPosition = 0;
             return this;
         }
 
@@ -155,6 +172,18 @@
         }
 
         /**
+         * Copies a raw byte into the TLV buffer - without a type or a length.
+         *
+         * @param b The byte to be inserted into the structure.
+         * @return The constructor to facilitate chaining {@code cts.putXXX(..).putXXX(..)}.
+         */
+        public TlvConstructor putRawByte(byte b) {
+            checkRawLength(1);
+            mArray[mPosition++] = b;
+            return this;
+        }
+
+        /**
          * Copies a byte array into the TLV with the indicated type. For an LV
          * formatted structure (i.e. typeLength=0 in {@link TlvConstructor
          * TlvConstructor(int, int)} ) the type field is ignored.
@@ -193,6 +222,22 @@
         }
 
         /**
+         * Copies a byte array into the TLV - without a type or a length.
+         *
+         * @param array The array to be copied (in full) into the TLV structure.
+         * @return The constructor to facilitate chaining
+         *         {@code ctr.putXXX(..).putXXX(..)}.
+         */
+        public TlvConstructor putRawByteArray(@Nullable byte[] array) {
+            if (array == null) return this;
+
+            checkRawLength(array.length);
+            System.arraycopy(array, 0, mArray, mPosition, array.length);
+            mPosition += array.length;
+            return this;
+        }
+
+        /**
          * Places a zero length element (i.e. Length field = 0) into the TLV.
          * For an LV formatted structure (i.e. typeLength=0 in
          * {@link TlvConstructor TlvConstructor(int, int)} ) the type field is
@@ -221,7 +266,7 @@
         public TlvConstructor putShort(int type, short data) {
             checkLength(2);
             addHeader(type, 2);
-            Memory.pokeShort(mArray, mPosition, data, ByteOrder.BIG_ENDIAN);
+            Memory.pokeShort(mArray, mPosition, data, mByteOrder);
             mPosition += 2;
             return this;
         }
@@ -239,7 +284,7 @@
         public TlvConstructor putInt(int type, int data) {
             checkLength(4);
             addHeader(type, 4);
-            Memory.pokeInt(mArray, mPosition, data, ByteOrder.BIG_ENDIAN);
+            Memory.pokeInt(mArray, mPosition, data, mByteOrder);
             mPosition += 4;
             return this;
         }
@@ -294,18 +339,24 @@
             }
         }
 
+        private void checkRawLength(int dataLength) {
+            if (mPosition + dataLength > mArrayLength) {
+                throw new BufferOverflowException();
+            }
+        }
+
         private void addHeader(int type, int length) {
             if (mTypeSize == 1) {
                 mArray[mPosition] = (byte) type;
             } else if (mTypeSize == 2) {
-                Memory.pokeShort(mArray, mPosition, (short) type, ByteOrder.BIG_ENDIAN);
+                Memory.pokeShort(mArray, mPosition, (short) type, mByteOrder);
             }
             mPosition += mTypeSize;
 
             if (mLengthSize == 1) {
                 mArray[mPosition] = (byte) length;
             } else if (mLengthSize == 2) {
-                Memory.pokeShort(mArray, mPosition, (short) length, ByteOrder.BIG_ENDIAN);
+                Memory.pokeShort(mArray, mPosition, (short) length, mByteOrder);
             }
             mPosition += mLengthSize;
         }
@@ -330,13 +381,19 @@
         public int length;
 
         /**
+         * Control of the endianess of the TLV element - true for big-endian, false for little-
+         * endian.
+         */
+        public ByteOrder byteOrder = ByteOrder.BIG_ENDIAN;
+
+        /**
          * The Value (V) field - a raw byte array representing the current TLV
          * element where the entry starts at {@link TlvElement#offset}.
          */
-        public byte[] refArray;
+        private byte[] mRefArray;
 
         /**
-         * The offset to be used into {@link TlvElement#refArray} to access the
+         * The offset to be used into {@link TlvElement#mRefArray} to access the
          * raw data representing the current TLV element.
          */
         public int offset;
@@ -344,7 +401,7 @@
         private TlvElement(int type, int length, @Nullable byte[] refArray, int offset) {
             this.type = type;
             this.length = length;
-            this.refArray = refArray;
+            mRefArray = refArray;
             this.offset = offset;
 
             if (offset + length > refArray.length) {
@@ -353,6 +410,15 @@
         }
 
         /**
+         * Return the raw byte array of the Value (V) field.
+         *
+         * @return The Value (V) field as a byte array.
+         */
+        public byte[] getRawData() {
+            return Arrays.copyOfRange(mRefArray, offset, offset + length);
+        }
+
+        /**
          * Utility function to return a byte representation of a TLV element of
          * length 1. Note: an attempt to call this function on a TLV item whose
          * {@link TlvElement#length} is != 1 will result in an exception.
@@ -364,7 +430,7 @@
                 throw new IllegalArgumentException(
                         "Accesing a byte from a TLV element of length " + length);
             }
-            return refArray[offset];
+            return mRefArray[offset];
         }
 
         /**
@@ -379,7 +445,7 @@
                 throw new IllegalArgumentException(
                         "Accesing a short from a TLV element of length " + length);
             }
-            return Memory.peekShort(refArray, offset, ByteOrder.BIG_ENDIAN);
+            return Memory.peekShort(mRefArray, offset, byteOrder);
         }
 
         /**
@@ -394,7 +460,7 @@
                 throw new IllegalArgumentException(
                         "Accesing an int from a TLV element of length " + length);
             }
-            return Memory.peekInt(refArray, offset, ByteOrder.BIG_ENDIAN);
+            return Memory.peekInt(mRefArray, offset, byteOrder);
         }
 
         /**
@@ -403,7 +469,7 @@
          * @return String repersentation of the current TLV element.
          */
         public String getString() {
-            return new String(refArray, offset, length);
+            return new String(mRefArray, offset, length);
         }
     }
 
@@ -413,6 +479,7 @@
     public static class TlvIterable implements Iterable<TlvElement> {
         private int mTypeSize;
         private int mLengthSize;
+        private ByteOrder mByteOrder = ByteOrder.BIG_ENDIAN;
         private byte[] mArray;
         private int mArrayLength;
 
@@ -440,6 +507,13 @@
         }
 
         /**
+         * Configure the TLV iterator to use little-endian byte ordering.
+         */
+        public void setByteOrder(ByteOrder byteOrder) {
+            mByteOrder = byteOrder;
+        }
+
+        /**
          * Prints out a parsed representation of the TLV-formatted byte array.
          * Whenever possible bytes, shorts, and integer are printed out (for
          * fields whose length is 1, 2, or 4 respectively).
@@ -486,7 +560,7 @@
         public List<byte[]> toList() {
             List<byte[]> list = new ArrayList<>();
             for (TlvElement tlv : this) {
-                list.add(Arrays.copyOfRange(tlv.refArray, tlv.offset, tlv.offset + tlv.length));
+                list.add(Arrays.copyOfRange(tlv.mRefArray, tlv.offset, tlv.offset + tlv.length));
             }
 
             return list;
@@ -516,7 +590,7 @@
                     if (mTypeSize == 1) {
                         type = mArray[mOffset];
                     } else if (mTypeSize == 2) {
-                        type = Memory.peekShort(mArray, mOffset, ByteOrder.BIG_ENDIAN);
+                        type = Memory.peekShort(mArray, mOffset, mByteOrder);
                     }
                     mOffset += mTypeSize;
 
@@ -524,11 +598,12 @@
                     if (mLengthSize == 1) {
                         length = mArray[mOffset];
                     } else if (mLengthSize == 2) {
-                        length = Memory.peekShort(mArray, mOffset, ByteOrder.BIG_ENDIAN);
+                        length = Memory.peekShort(mArray, mOffset, mByteOrder);
                     }
                     mOffset += mLengthSize;
 
                     TlvElement tlv = new TlvElement(type, length, mArray, mOffset);
+                    tlv.byteOrder = mByteOrder;
                     mOffset += length;
                     return tlv;
                 }
@@ -543,7 +618,8 @@
 
     /**
      * Validates that a (T)LV array is constructed correctly. I.e. that its specified Length
-     * fields correctly fill the specified length (and do not overshoot).
+     * fields correctly fill the specified length (and do not overshoot). Uses big-endian
+     * byte ordering.
      *
      * @param array The (T)LV array to verify.
      * @param typeSize The size (in bytes) of the type field. Valid values are 0, 1, or 2.
@@ -551,6 +627,22 @@
      * @return A boolean indicating whether the array is valid (true) or invalid (false).
      */
     public static boolean isValid(@Nullable byte[] array, int typeSize, int lengthSize) {
+        return isValidEndian(array, typeSize, lengthSize, ByteOrder.BIG_ENDIAN);
+    }
+
+    /**
+     * Validates that a (T)LV array is constructed correctly. I.e. that its specified Length
+     * fields correctly fill the specified length (and do not overshoot).
+     *
+     * @param array The (T)LV array to verify.
+     * @param typeSize The size (in bytes) of the type field. Valid values are 0, 1, or 2.
+     * @param lengthSize The size (in bytes) of the length field. Valid values are 1 or 2.
+     * @param byteOrder The endianness of the byte array: {@link ByteOrder#BIG_ENDIAN} or
+     *                  {@link ByteOrder#LITTLE_ENDIAN}.
+     * @return A boolean indicating whether the array is valid (true) or invalid (false).
+     */
+    public static boolean isValidEndian(@Nullable byte[] array, int typeSize, int lengthSize,
+            ByteOrder byteOrder) {
         if (typeSize < 0 || typeSize > 2) {
             throw new IllegalArgumentException(
                     "Invalid arguments - typeSize must be 0, 1, or 2: typeSize=" + typeSize);
@@ -569,8 +661,7 @@
             if (lengthSize == 1) {
                 nextTlvIndex += lengthSize + array[nextTlvIndex];
             } else {
-                nextTlvIndex += lengthSize + Memory.peekShort(array, nextTlvIndex,
-                        ByteOrder.BIG_ENDIAN);
+                nextTlvIndex += lengthSize + Memory.peekShort(array, nextTlvIndex, byteOrder);
             }
         }
 
diff --git a/wifi/tests/src/android/net/wifi/aware/TlvBufferUtilsTest.java b/wifi/tests/src/android/net/wifi/aware/TlvBufferUtilsTest.java
index 83affed..971aa8e 100644
--- a/wifi/tests/src/android/net/wifi/aware/TlvBufferUtilsTest.java
+++ b/wifi/tests/src/android/net/wifi/aware/TlvBufferUtilsTest.java
@@ -28,6 +28,7 @@
 import java.util.ArrayList;
 import java.util.List;
 
+
 /**
  * Unit test harness for TlvBufferUtils class.
  */
@@ -69,6 +70,24 @@
     }
 
     /**
+     * Validate that re-using a TLV by any of the reallocation method resets it completely.
+     */
+    @Test
+    public void testTlvReuse() {
+        TlvBufferUtils.TlvConstructor tlv = new TlvBufferUtils.TlvConstructor(1, 1);
+
+        tlv.allocate(10);
+        tlv.putByte(0, (byte) 2);
+        tlv.putByte(1, (byte) 104);
+
+        collector.checkThat("initial", tlv.getArray(), equalTo(new byte[]{0, 1, 2, 1, 1, 104}));
+
+        tlv.allocate(8);
+        tlv.putByte(5, (byte) 7);
+        collector.checkThat("re-alloc", tlv.getArray(), equalTo(new byte[]{5, 1, 7}));
+    }
+
+    /**
      * Verify that can build a valid TLV from a List of byte[].
      */
     @Test
@@ -121,6 +140,23 @@
         List<byte[]> data = new TlvBufferUtils.TlvIterable(0, 1, invalidTlv01).toList();
     }
 
+    /**
+     * Validate the API which places raw bytes into the TLV (without a TL structure).
+     */
+    @Test
+    public void testRawPuts() {
+        TlvBufferUtils.TlvConstructor tlv = new TlvBufferUtils.TlvConstructor(1, 1);
+
+        tlv.allocate(10);
+        tlv.putByte(0, (byte) 2);
+        tlv.putRawByte((byte) 55);
+        tlv.putByte(1, (byte) 104);
+        tlv.putRawByteArray(new byte[]{66, 77});
+
+        collector.checkThat("data", tlv.getArray(),
+                equalTo(new byte[]{0, 1, 2, 55, 1, 1, 104, 66, 77}));
+    }
+
     @Test
     public void testTlvIterate() {
         final String ascii = "ABC";
@@ -163,6 +199,7 @@
         tlv02.putByte(0, (byte) 2);
         tlv02.putString(0, ascii);
         tlv02.putString(0, nonAscii);
+        tlv02.putByteArray(0, new byte[]{5, 4, 3, 2, 1});
 
         TlvBufferUtils.TlvIterable tlv02It = new TlvBufferUtils.TlvIterable(0, 2, tlv02.getArray());
         count = 0;
@@ -181,6 +218,11 @@
                         equalTo(nonAscii.getBytes().length));
                 collector.checkThat("tlv02-correct-iteration-DATA",
                         tlv.getString().equals(nonAscii), equalTo(true));
+            } else if (count == 3) {
+                collector.checkThat("tlv02-correct-iteration-mLength", tlv.length,
+                        equalTo(5));
+                collector.checkThat("tlv02-correct-iteration-DATA", tlv.getRawData(),
+                        equalTo(new byte[]{5, 4, 3, 2, 1}));
             } else {
                 collector.checkThat("Invalid number of iterations in loop - tlv02", true,
                         equalTo(false));
@@ -188,7 +230,7 @@
             ++count;
         }
         collector.checkThat("Invalid number of iterations outside loop - tlv02", count,
-                equalTo(3));
+                equalTo(4));
 
         collector.checkThat("tlv22-valid",
                 TlvBufferUtils.isValid(tlv22.getArray(), 2, 2),