Add support for reading graph updates from a file, and example binary reading test.

arik-so · arik-so · commit 1110d4d69fd0 · 2021-11-09T19:17:28.000-08:00
diff --git a/lightning-graph-sync/src/lib.rs b/lightning-graph-sync/src/lib.rs
@@ -1,4 +1,5 @@
 use core::ops::Deref;
+use std::fs::File;
 
 use lightning::chain;
 use lightning::routing::network_graph;
@@ -9,19 +10,22 @@ use crate::error::GraphSyncError;
 pub mod error;
 pub mod processing;
 
-pub async fn sync_network_graph<CS: Deref>(
+pub async fn sync_network_graph_with_url<CS: Deref>(
     network_graph: &network_graph::NetworkGraph,
     sync_host: &str,
     sync_port: u16,
     sync_path: &str,
     chain_source: &Option<CS>,
 ) -> Result<(), GraphSyncError>
-    where
-        CS::Target: chain::Access,
+where
+    CS::Target: chain::Access,
 {
     // make sure there is precisely one leading slash
     let canonical_path = format!("/{}", sync_path.trim_start_matches("/"));
-    println!("Attempting sync from: http://{}:{}{}", sync_host, sync_port, canonical_path);
+    println!(
+        "Attempting sync from: http://{}:{}{}",
+        sync_host, sync_port, canonical_path
+    );
     let http_endpoint = HttpEndpoint::for_host(sync_host.into()).with_port(sync_port);
     let http_client_result = HttpClient::connect(&http_endpoint);
     let mut http_client = http_client_result?;
@@ -36,20 +40,29 @@ pub async fn sync_network_graph<CS: Deref>(
     processing::update_network_graph(&network_graph, &response_bytes[..], &chain_source)
 }
 
+pub fn sync_network_graph_with_file_path<CS: Deref>(
+    network_graph: &network_graph::NetworkGraph,
+    sync_path: &str,
+    chain_source: &Option<CS>,
+) -> Result<(), GraphSyncError>
+where
+    CS::Target: chain::Access,
+{
+    let mut file = File::open(sync_path)?;
+    processing::read_network_graph(&network_graph, &mut file, &chain_source)
+}
+
 #[cfg(test)]
 mod tests {
     use bitcoin::blockdata::constants::genesis_block;
     use bitcoin::Network;
 
     use lightning::routing::network_graph::NetworkGraph;
+
     use crate::error::GraphSyncError;
+    use crate::sync_network_graph_with_url;
 
-    use crate::sync_network_graph;
 
-    #[test]
-    fn it_works() {
-        assert_eq!(2 + 2, 4);
-    }
 
     #[tokio::test]
     async fn test_http_request() {
@@ -58,28 +71,21 @@ mod tests {
 
         let chain_source: Option<Box<dyn lightning::chain::Access>> = None;
 
-        let sync_result = sync_network_graph(
+        let sync_result = sync_network_graph_with_url(
             &network_graph,
             "localhost",
             3000,
             "delta/500",
             &chain_source,
-        ).await;
+        )
+        .await;
 
         if !sync_result.is_ok() {
             let error_string = match sync_result.as_ref().unwrap_err() {
-                GraphSyncError::IOError(error) => {
-                    error.to_string()
-                }
-                GraphSyncError::DecodeError(error) => {
-                    error.to_string()
-                }
-                GraphSyncError::LightningError(error) => {
-                    (&error.err).to_string()
-                }
-                GraphSyncError::ProcessingError(error) => {
-                    error.to_string()
-                }
+                GraphSyncError::IOError(error) => error.to_string(),
+                GraphSyncError::DecodeError(error) => error.to_string(),
+                GraphSyncError::LightningError(error) => (&error.err).to_string(),
+                GraphSyncError::ProcessingError(error) => error.to_string(),
             };
             println!("Error: {}", error_string);
         }
diff --git a/lightning-graph-sync/src/processing.rs b/lightning-graph-sync/src/processing.rs
@@ -1,36 +1,66 @@
 use core::ops::Deref;
 use std::io;
 use std::io::Read;
-use lightning::chain;
 
+use lightning::chain;
 use lightning::ln::msgs;
 use lightning::ln::wire::encode::Encode;
 use lightning::routing::network_graph;
-use lightning::util::ser::{BigSize, FixedLengthReader, Readable};
+use lightning::util::ser::{BigSize, Readable};
 
 use crate::error::GraphSyncError;
 
-pub fn update_network_graph<CS: Deref>(network_graph: &network_graph::NetworkGraph, update_data: &[u8], chain_source: &Option<CS>) -> Result<(), GraphSyncError> where CS::Target: chain::Access {
-    let total_input_length = update_data.len() as u64;
+pub fn update_network_graph<CS: Deref>(
+    network_graph: &network_graph::NetworkGraph,
+    update_data: &[u8],
+    chain_source: &Option<CS>,
+) -> Result<(), GraphSyncError>
+where
+    CS::Target: chain::Access,
+{
     let mut read_cursor = io::Cursor::new(update_data);
+    read_network_graph(&network_graph, &mut read_cursor, &chain_source)
+}
 
+pub(crate) fn read_network_graph<R: Read, CS: Deref>(
+    network_graph: &network_graph::NetworkGraph,
+    read_cursor: &mut R,
+    chain_source: &Option<CS>,
+) -> Result<(), GraphSyncError>
+where
+    CS::Target: chain::Access,
+{
     let mut prefix = Vec::with_capacity(4);
     prefix.resize(4, 0);
     read_cursor.read_exact(&mut prefix)?;
 
     let _deserialize_unsigned = match &prefix[..] {
         &[71, 83, 80, 1] => {
-            return Err(GraphSyncError::ProcessingError("Signed deserialization not supported".to_string()));
+            return Err(GraphSyncError::ProcessingError(
+                "Signed deserialization not supported".to_string(),
+            ));
         }
         &[76, 68, 75, 2] => true,
         _ => {
-            return Err(GraphSyncError::ProcessingError("Prefix must equal 'LDK' and the byte 2".to_string()));
+            return Err(GraphSyncError::ProcessingError(
+                "Prefix must equal 'LDK' and the byte 2".to_string(),
+            ));
         }
     };
 
+    // shadow variable for ownership release after length-restricted reads
+    let mut read_cursor = read_cursor;
     loop {
-        let message_length: BigSize = Readable::read(&mut read_cursor)?;
-        let mut restricted_cursor = FixedLengthReader::new(&mut read_cursor, message_length.0);
+        let message_length_result = Readable::read(read_cursor);
+        let message_length: BigSize = if let Ok(message_length) = message_length_result {
+            message_length
+        } else if let Err(msgs::DecodeError::ShortRead) = message_length_result {
+            break;
+        } else {
+            return Err(message_length_result.err().unwrap().into());
+        };
+        // let mut restricted_cursor = FixedLengthReader::new(&read_cursor, message_length.0);
+        let mut restricted_cursor = read_cursor.take(message_length.0);
 
         let message_type: u16 = Readable::read(&mut restricted_cursor)?;
 
@@ -60,10 +90,62 @@ pub fn update_network_graph<CS: Deref>(network_graph: &network_graph::NetworkGra
         // propagate the error
         let _update = update_result?;
 
-        if read_cursor.position() == total_input_length {
-            break;
-        }
+        read_cursor = restricted_cursor.into_inner();
     }
 
     Ok(())
-}
+}
+
+#[cfg(test)]
+mod tests {
+    use bitcoin::blockdata::constants::genesis_block;
+    use bitcoin::Network;
+
+    use lightning::routing::network_graph::NetworkGraph;
+
+    use crate::processing::update_network_graph;
+
+    #[test]
+    fn network_graph_updates_from_example_binary_input() {
+        let block_hash = genesis_block(Network::Bitcoin).block_hash();
+        let network_graph = NetworkGraph::new(block_hash);
+
+        let before = network_graph.to_string();
+        assert_eq!(before.len(), 31);
+
+        let chain_source: Option<Box<dyn lightning::chain::Access>> = None;
+
+        let example_input = vec![
+            76, 68, 75, 2, 176, 1, 0, 0, 0, 111, 226, 140, 10, 182, 241, 179, 114, 193, 166, 162,
+            70, 174, 99, 247, 79, 147, 30, 131, 101, 225, 90, 8, 156, 104, 214, 25, 0, 0, 0, 0, 0,
+            10, 209, 73, 0, 2, 144, 0, 0, 2, 96, 250, 182, 51, 6, 110, 215, 177, 217, 185, 184,
+            160, 250, 200, 126, 21, 121, 209, 112, 158, 135, 77, 40, 160, 209, 113, 161, 245, 196,
+            59, 184, 119, 2, 180, 101, 113, 124, 204, 150, 157, 103, 192, 42, 110, 145, 224, 176,
+            14, 26, 197, 136, 6, 143, 13, 46, 3, 200, 220, 34, 228, 191, 104, 238, 134, 80, 3, 167,
+            152, 93, 150, 15, 179, 70, 61, 160, 42, 198, 253, 53, 249, 37, 29, 21, 176, 28, 136,
+            240, 237, 138, 165, 218, 254, 14, 245, 123, 228, 145, 150, 2, 25, 120, 160, 148, 15,
+            56, 179, 42, 198, 99, 96, 81, 85, 53, 227, 51, 62, 134, 4, 175, 0, 151, 217, 104, 180,
+            201, 249, 117, 196, 87, 94, 150, 74, 1, 2, 111, 226, 140, 10, 182, 241, 179, 114, 193,
+            166, 162, 70, 174, 99, 247, 79, 147, 30, 131, 101, 225, 90, 8, 156, 104, 214, 25, 0, 0,
+            0, 0, 0, 10, 209, 73, 0, 2, 144, 0, 0, 97, 139, 55, 91, 1, 2, 0, 40, 0, 0, 0, 0, 0, 0,
+            3, 232, 0, 0, 0, 0, 0, 0, 0, 100, 0, 0, 0, 0, 123, 235, 5, 192, 74, 1, 2, 111, 226,
+            140, 10, 182, 241, 179, 114, 193, 166, 162, 70, 174, 99, 247, 79, 147, 30, 131, 101,
+            225, 90, 8, 156, 104, 214, 25, 0, 0, 0, 0, 0, 10, 209, 73, 0, 2, 144, 0, 0, 97, 138,
+            144, 189, 1, 1, 0, 40, 0, 0, 0, 0, 0, 0, 3, 232, 0, 0, 0, 0, 0, 0, 5, 220, 0, 0, 0, 0,
+            123, 235, 5, 192,
+        ];
+        let update_result = update_network_graph(&network_graph, &example_input[..], &chain_source);
+        assert!(update_result.is_ok());
+
+        let after = network_graph.to_string();
+        assert_eq!(after.len(), 1132);
+        assert!(after.contains("779484474903625728: features: 0000"));
+        assert!(after.contains(
+            "node_one: 0260fab633066ed7b1d9b9b8a0fac87e1579d1709e874d28a0d171a1f5c43bb877"
+        ));
+        assert!(after.contains(
+            "node_two: 02b465717ccc969d67c02a6e91e0b00e1ac588068f0d2e03c8dc22e4bf68ee8650"
+        ));
+        assert!(after.contains("channels: [779484474903625728]"));
+    }
+}
