Merge pull request #7 from LNP-WG/fix/size-encoding

Fix encodings of TLV data

Cargo.toml

@@ -18,7 +18,7 @@ path = "src/lib.rs"
 crate-type = ["rlib", "staticlib"]
 
 [dependencies]
-amplify = "3.13.0"
+amplify = "3.14.2"
 strict_encoding = "0.9.0"
 lightning_encoding_derive = { version = "0.9.1", path = "derive", optional = true }
 lnpbp_chain = "0.9.0"

src/bitcoin.rs

@@ -90,15 +90,17 @@ impl Strategy for hlc::HashPreimage {
 
 impl LightningEncode for Script {
     #[inline]
-    fn lightning_encode<E: Write>(&self, e: E) -> Result<usize, Error> {
-        self.as_bytes().lightning_encode(e)
+    fn lightning_encode<E: Write>(&self, mut e: E) -> Result<usize, Error> {
+        e.write_all(self.as_bytes())?;
+        Ok(self.len())
     }
 }
 
 impl LightningDecode for Script {
-    fn lightning_decode<D: Read>(d: D) -> Result<Self, Error> {
-        let value = Vec::<u8>::lightning_decode(d)?;
-        let bytes = consensus::serialize(&value);
+    fn lightning_decode<D: Read>(mut d: D) -> Result<Self, Error> {
+        let mut buf = vec![];
+        d.read_to_end(&mut buf)?;
+        let bytes = consensus::serialize(&buf);
         consensus::deserialize(&bytes)
             .map_err(|err| Error::DataIntegrityError(err.to_string()))
     }
@@ -124,10 +126,11 @@ mod test {
     fn real_clightning_scriptpubkey() {
         // Real scriptpubkey sent by clightning
         let msg_recv = [
-            0, 22, 0, 20, 42, 238, 172, 27, 222, 161, 61, 181, 251, 208, 97,
+            0u8, 22, 0, 20, 42, 238, 172, 27, 222, 161, 61, 181, 251, 208, 97,
             79, 71, 255, 98, 8, 213, 205, 114, 94,
         ];
 
-        PubkeyScript::lightning_deserialize(&msg_recv).unwrap();
+        let script = PubkeyScript::lightning_deserialize(&msg_recv).unwrap();
+        assert_eq!(script.lightning_serialize().unwrap(), msg_recv);
     }
 }

src/primitives.rs

@@ -11,15 +11,14 @@
 // along with this software.
 // If not, see <https://opensource.org/licenses/MIT>.
 
-use std::io;
 use std::io::{Read, Write};
 
 use amplify::flags::FlagVec;
 use amplify::num::u24;
 use amplify::{Slice32, Wrapper};
 
 use super::{strategies, Strategy};
-use crate::{Error, LightningDecode, LightningEncode};
+use crate::{BigSize, Error, LightningDecode, LightningEncode};
 
 impl LightningEncode for u8 {
     fn lightning_encode<E: Write>(&self, mut e: E) -> Result<usize, Error> {
@@ -101,55 +100,37 @@ impl LightningDecode for u64 {
 }
 
 impl LightningEncode for usize {
-    fn lightning_encode<E: Write>(&self, mut e: E) -> Result<usize, Error> {
-        let count = match *self {
-            count if count > u16::MAX as usize => {
-                return Err(Error::TooLargeData(count))
-            }
-            count => count as u16,
-        };
-        let bytes = count.to_be_bytes();
-        e.write_all(&bytes)?;
-        Ok(bytes.len())
+    fn lightning_encode<E: Write>(&self, e: E) -> Result<usize, Error> {
+        let size = BigSize::from(*self);
+        size.lightning_encode(e)
     }
 }
 
 impl LightningDecode for usize {
-    fn lightning_decode<D: Read>(mut d: D) -> Result<Self, Error> {
-        let mut buf = [0u8; 2];
-        d.read_exact(&mut buf)?;
-        Ok(u16::from_be_bytes(buf) as usize)
+    fn lightning_decode<D: Read>(d: D) -> Result<Self, Error> {
+        BigSize::lightning_decode(d).map(|size| size.into_inner() as usize)
     }
 }
 
 impl LightningEncode for FlagVec {
-    fn lightning_encode<E: Write>(&self, e: E) -> Result<usize, Error> {
+    fn lightning_encode<E: Write>(&self, mut e: E) -> Result<usize, Error> {
         let flags = self.shrunk();
-        flags.as_inner().lightning_encode(e)
+        let mut vec = flags.as_inner().to_vec();
+        vec.reverse();
+        let len = vec.len() as u16;
+        len.lightning_encode(&mut e)?;
+        e.write_all(&vec)?;
+        Ok(vec.len() + 2)
     }
 }
 
 impl LightningDecode for FlagVec {
-    fn lightning_decode<D: Read>(d: D) -> Result<Self, Error> {
-        let flags = Vec::<u8>::lightning_decode(d)?;
-        Ok(FlagVec::from_inner(flags))
-    }
-    fn lightning_deserialize(data: impl AsRef<[u8]>) -> Result<Self, Error> {
-        let bytes = data.as_ref();
-        if bytes.is_empty() {
-            Ok(FlagVec::default())
-        } else {
-            let mut decoder = io::Cursor::new(data.as_ref());
-            let rv = Self::lightning_decode(&mut decoder)?;
-            let consumed = decoder.position() as usize;
-
-            // Fail if data are not consumed entirely.
-            if consumed == data.as_ref().len() {
-                Ok(rv)
-            } else {
-                Err(Error::DataNotEntirelyConsumed)
-            }
-        }
+    fn lightning_decode<D: Read>(mut d: D) -> Result<Self, Error> {
+        let len = u16::lightning_decode(&mut d)?;
+        let mut buf = vec![0u8; len as usize];
+        d.read_exact(&mut buf)?;
+        buf.reverse();
+        Ok(FlagVec::from_inner(buf))
     }
 }