From 8ef37a60cb4f724e453a5b4806a99409184d9484 Mon Sep 17 00:00:00 2001
From: Song GUO <songguo@google.com>
Date: Wed, 20 Oct 2021 14:07:00 +0800
Subject: [PATCH] [python] Add unsigned int support to CHIP TLV Reader / Writer
 in Python (#10639)

* Add tlv.sint and tlv.uint to distinguish signed and unsigned int when encoding

* Use Pythonic NewType

* Revert and add a note
---
 .github/workflows/build.yaml                  |   5 +
 src/controller/python/chip/tlv/__init__.py    |  18 ++-
 .../python/test/unit_tests/test_tlv.py        | 153 ++++++++++++++++++
 3 files changed, 172 insertions(+), 4 deletions(-)
 create mode 100644 src/controller/python/test/unit_tests/test_tlv.py

diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index bf0c6697bcb19c..d3146cae769ec5 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -231,6 +231,11 @@ jobs:
                       scripts/build/gn_build.sh
                       scripts/tests/gn_tests.sh
                   done
+            - name: Run Python library specific unit tests
+              timeout-minutes: 5
+              run: |
+                  scripts/run_in_build_env.sh 'pip3 install ./out/controller/python/chip-0.0-cp37-abi3-linux_x86_64.whl'
+                  scripts/run_in_build_env.sh '(cd src/controller/python/test/unit_tests/ && python3 -m unittest -v)'
             # TODO Log Upload https://github.com/project-chip/connectedhomeip/issues/2227
             # TODO https://github.com/project-chip/connectedhomeip/issues/1512
             # - name: Run Code Coverage
diff --git a/src/controller/python/chip/tlv/__init__.py b/src/controller/python/chip/tlv/__init__.py
index e1d4d0a3a23a0f..4f5c771cba930a 100644
--- a/src/controller/python/chip/tlv/__init__.py
+++ b/src/controller/python/chip/tlv/__init__.py
@@ -112,6 +112,13 @@
 }
 
 
+class uint(int):
+    '''
+    NewType will not return a class until Python 3.10, as Python 3.10 is not widely used, we still need to construct a class so it can work as a type.
+    '''
+    pass
+
+
 class TLVWriter(object):
     def __init__(self, encoding=None, implicitProfile=None):
         self._encoding = encoding if encoding is not None else bytearray()
@@ -177,11 +184,10 @@ def put(self, tag, val):
             self.putNull(tag)
         elif isinstance(val, bool):
             self.putBool(tag, val)
+        elif isinstance(val, uint):
+            self.putUnsignedInt(tag, val)
         elif isinstance(val, int):
-            if val < 0:
-                self.putSignedInt(tag, val)
-            else:
-                self.putUnsignedInt(tag, val)
+            self.putSignedInt(tag, val)
         elif isinstance(val, float):
             self.putFloat(tag, val)
         elif isinstance(val, str):
@@ -553,6 +559,7 @@ def _decodeVal(self, tlv, decoding):
             (decoding["value"],) = struct.unpack(
                 "<B", tlv[self._bytesRead: self._bytesRead + 1]
             )
+            decoding["value"] = uint(decoding["value"])
             self._bytesRead += 1
         elif decoding["type"] == "Signed Integer 1-byte value":
             (decoding["value"],) = struct.unpack(
@@ -563,6 +570,7 @@ def _decodeVal(self, tlv, decoding):
             (decoding["value"],) = struct.unpack(
                 "<H", tlv[self._bytesRead: self._bytesRead + 2]
             )
+            decoding["value"] = uint(decoding["value"])
             self._bytesRead += 2
         elif decoding["type"] == "Signed Integer 2-byte value":
             (decoding["value"],) = struct.unpack(
@@ -573,6 +581,7 @@ def _decodeVal(self, tlv, decoding):
             (decoding["value"],) = struct.unpack(
                 "<L", tlv[self._bytesRead: self._bytesRead + 4]
             )
+            decoding["value"] = uint(decoding["value"])
             self._bytesRead += 4
         elif decoding["type"] == "Signed Integer 4-byte value":
             (decoding["value"],) = struct.unpack(
@@ -583,6 +592,7 @@ def _decodeVal(self, tlv, decoding):
             (decoding["value"],) = struct.unpack(
                 "<Q", tlv[self._bytesRead: self._bytesRead + 8]
             )
+            decoding["value"] = uint(decoding["value"])
             self._bytesRead += 8
         elif decoding["type"] == "Signed Integer 8-byte value":
             (decoding["value"],) = struct.unpack(
diff --git a/src/controller/python/test/unit_tests/test_tlv.py b/src/controller/python/test/unit_tests/test_tlv.py
new file mode 100644
index 00000000000000..7e83f9c26863f2
--- /dev/null
+++ b/src/controller/python/test/unit_tests/test_tlv.py
@@ -0,0 +1,153 @@
+#
+#    Copyright (c) 2021 Project CHIP Authors
+#    All rights reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+
+
+from chip.tlv import TLVWriter, TLVReader
+from chip.tlv import uint as tlvUint
+
+import unittest
+
+
+class TestTLVWriter(unittest.TestCase):
+    def _getEncoded(self, val, tag=None):
+        writer = TLVWriter()
+        writer.put(tag, val)
+        return writer.encoding
+
+    def test_int(self):
+        encodedVal = self._getEncoded(0x00deadbeefca00fe)
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000011,
+                                    0xfe, 0x00, 0xca, 0xef, 0xbe, 0xad, 0xde, 0x00]))
+        encodedVal = self._getEncoded(0x7cadbeef)
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000010, 0xef, 0xbe, 0xad, 0x7c]))
+        encodedVal = self._getEncoded(0x7cad)
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000001, 0xad, 0x7c]))
+        encodedVal = self._getEncoded(0x7c)
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000000, 0x7c]))
+        # Negative numbers
+        encodedVal = self._getEncoded(-(0x5555555555555555))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000011,
+                                    0xab, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa]))
+        encodedVal = self._getEncoded(-(0x55555555))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000010, 0xab, 0xaa, 0xaa, 0xaa]))
+        encodedVal = self._getEncoded(-(0x5555))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000001, 0xab, 0xaa]))
+        encodedVal = self._getEncoded(-(0x55))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000000, 0xab]))
+        # The following numbers are positive values but exceeds the upper bounds of the type they seems to be.
+        encodedVal = self._getEncoded(0xdeadbeef)
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000011, 0xef, 0xbe, 0xad, 0xde, 0x00, 0x00, 0x00, 0x00]))
+        encodedVal = self._getEncoded(0xdead)
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000010, 0xad, 0xde, 0x00, 0x00]))
+        encodedVal = self._getEncoded(0xad)
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000001, 0xad, 0x00]))
+        # Similar, these negative numbers also exceedes the width of the type they seems to be.
+        encodedVal = self._getEncoded(-(0xaaaaaaaa))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000011, 0x56, 0x55, 0x55, 0x55, 0xff, 0xff, 0xff, 0xff]))
+        encodedVal = self._getEncoded(-(0xaaaa))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000010, 0x56, 0x55, 0xff, 0xff]))
+        encodedVal = self._getEncoded(-(0xaa))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000001, 0x56, 0xff]))
+        try:
+            encodedVal = self._getEncoded(0xf000f000f000f000)
+            self.fail("Signed number exceeds INT64_MAX but no exception received")
+        except Exception:
+            pass
+
+        try:
+            encodedVal = self._getEncoded(-(0xf000f000f000f000))
+            self.fail("Signed number exceeds INT64_MIN but no exception received")
+        except Exception:
+            pass
+
+    def test_uint(self):
+        encodedVal = self._getEncoded(tlvUint(0xdeadbeefca0000fe))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000111,
+                                    0xfe, 0x00, 0x00, 0xca, 0xef, 0xbe, 0xad, 0xde]))
+        encodedVal = self._getEncoded(tlvUint(0xdeadbeef))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000110, 0xef, 0xbe, 0xad, 0xde]))
+        encodedVal = self._getEncoded(tlvUint(0xdead))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000101, 0xad, 0xde]))
+        encodedVal = self._getEncoded(tlvUint(0xde))
+        self.assertEqual(encodedVal,
+                         bytearray([0b00000100, 0xde]))
+        try:
+            encodedVal = self._getEncoded(tlvUint(-1))
+            self.fail("Negative unsigned int but no exception raised.")
+        except Exception:
+            pass
+        try:
+            encodedVal = self._getEncoded(tlvUint(0x10000000000000000))
+            self.fail("Overflowed uint but no exception raised.")
+        except Exception:
+            pass
+
+
+class TestTLVReader(unittest.TestCase):
+    def _read_case(self, input, answer):
+        decoded = TLVReader(bytearray(input)).get()["Any"]
+        self.assertEqual(type(decoded), type(answer))
+        self.assertEqual(decoded, answer)
+
+    def test_int(self):
+        self._read_case([0b00000011,
+                         0xfe, 0x00, 0xca, 0xef, 0xbe, 0xad, 0xde, 0x00], 0x00deadbeefca00fe)
+        self._read_case([0b00000010, 0xef, 0xbe, 0xad, 0x7c], 0x7cadbeef)
+        self._read_case([0b00000001, 0xad, 0x7c], 0x7cad)
+        self._read_case([0b00000000, 0x7c], 0x7c)
+
+        self._read_case([0b00000011,
+                         0xab, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa], -(0x5555555555555555))
+        self._read_case([0b00000010, 0xab, 0xaa, 0xaa, 0xaa], -(0x55555555))
+        self._read_case([0b00000001, 0xab, 0xaa], -(0x5555))
+        self._read_case([0b00000000, 0xab], -(0x55))
+
+    def test_uint(self):
+        self._read_case([0b00000111,
+                         0xfe, 0x00, 0xca, 0xef, 0xbe, 0xad, 0xde, 0x00], tlvUint(0x00deadbeefca00fe))
+        self._read_case([0b00000110, 0xef, 0xbe, 0xad, 0x7c],
+                        tlvUint(0x7cadbeef))
+        self._read_case([0b00000101, 0xad, 0x7c], tlvUint(0x7cad))
+        self._read_case([0b00000100, 0x7c], tlvUint(0x7c))
+
+        self._read_case([0b00000111,
+                         0xab, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa], tlvUint(0xaaaaaaaaaaaaaaab))
+        self._read_case([0b00000110, 0xab, 0xaa, 0xaa, 0xaa],
+                        tlvUint(0xaaaaaaab))
+        self._read_case([0b00000101, 0xab, 0xaa], tlvUint(0xaaab))
+        self._read_case([0b00000100, 0xab], tlvUint(0xab))
+
+
+if __name__ == '__main__':
+    unittest.main()