From dfdfef48e3b2f4c10726dd36d0a6bae61f53346c Mon Sep 17 00:00:00 2001
From: Joseph Livesey <joseph@blockchaintp.com>
Date: Fri, 17 Jun 2022 16:30:47 -0400
Subject: [PATCH] feature/no-fork-cli-flag

Signed-off-by: Joseph Livesey <joseph@blockchaintp.com>
---
 src/engine.rs | 60 +++++++++++++++++++++++++++++++--------------------
 src/main.rs   | 17 +++++++++------
 2 files changed, 47 insertions(+), 30 deletions(-)

diff --git a/src/engine.rs b/src/engine.rs
index 54044c40..ff8716c2 100644
--- a/src/engine.rs
+++ b/src/engine.rs
@@ -212,11 +212,13 @@ impl DevmodeService {
     }
 }
 
-pub struct DevmodeEngine {}
+pub struct DevmodeEngine {
+    no_forks: bool,
+}
 
 impl DevmodeEngine {
-    pub fn new() -> Self {
-        DevmodeEngine {}
+    pub fn new(no_forks: bool) -> Self {
+        DevmodeEngine { no_forks }
     }
 }
 
@@ -282,31 +284,43 @@ impl Engine for DevmodeEngine {
                                 DisplayBlock(&block)
                             );
 
-                            // Advance the chain if possible.
-                            if block.block_num > chain_head.block_num
-                                || (block.block_num == chain_head.block_num
-                                    && block.block_id > chain_head.block_id)
-                            {
-                                info!("Committing {}", DisplayBlock(&block));
-                                service.commit_block(block_id);
-                            } else if block.block_num < chain_head.block_num {
-                                let mut chain_block = chain_head;
-                                loop {
-                                    chain_block = service.get_block(&chain_block.previous_id);
-                                    if chain_block.block_num == block.block_num {
-                                        break;
-                                    }
-                                }
-                                if block.block_id > chain_block.block_id {
-                                    info!("Switching to new fork {}", DisplayBlock(&block));
+                            if self.no_forks {
+                                if block.block_num > chain_head.block_num
+                                    && block.block_num == chain_head.block_num + 1
+                                {
+                                    info!("Committing {}", DisplayBlock(&block));
                                     service.commit_block(block_id);
                                 } else {
-                                    info!("Ignoring fork {}", DisplayBlock(&block));
+                                    info!("Ignoring {}", DisplayBlock(&block));
                                     service.ignore_block(block_id);
                                 }
                             } else {
-                                info!("Ignoring {}", DisplayBlock(&block));
-                                service.ignore_block(block_id);
+                                // Advance the chain if possible.
+                                if block.block_num > chain_head.block_num
+                                    || (block.block_num == chain_head.block_num
+                                        && block.block_id > chain_head.block_id)
+                                {
+                                    info!("Committing {}", DisplayBlock(&block));
+                                    service.commit_block(block_id);
+                                } else if block.block_num < chain_head.block_num {
+                                    let mut chain_block = chain_head;
+                                    loop {
+                                        chain_block = service.get_block(&chain_block.previous_id);
+                                        if chain_block.block_num == block.block_num {
+                                            break;
+                                        }
+                                    }
+                                    if block.block_id > chain_block.block_id {
+                                        info!("Switching to new fork {}", DisplayBlock(&block));
+                                        service.commit_block(block_id);
+                                    } else {
+                                        info!("Ignoring fork {}", DisplayBlock(&block));
+                                        service.ignore_block(block_id);
+                                    }
+                                } else {
+                                    info!("Ignoring {}", DisplayBlock(&block));
+                                    service.ignore_block(block_id);
+                                }
                             }
                         }
 
diff --git a/src/main.rs b/src/main.rs
index 492d675c..f656f56b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -42,13 +42,17 @@ fn main() {
         (@arg connect: -C --connect +takes_value
          "connection endpoint for validator")
         (@arg verbose: -v --verbose +multiple
-         "increase output verbosity"))
+         "increase output verbosity")
+        (@arg no_forks: -n --no-forks
+         "prevent chain forking"))
     .get_matches();
 
     let endpoint = matches
         .value_of("connect")
         .unwrap_or("tcp://localhost:5050");
 
+    let no_forks: bool = matches.is_present("no_forks");
+
     let console_log_level;
     match matches.occurrences_of("verbose") {
         0 => console_log_level = LogLevelFilter::Warn,
@@ -77,10 +81,9 @@ fn main() {
     });
 
     let (driver, _stop) = ZmqDriver::new();
-    driver
-        .start(endpoint, DevmodeEngine::new())
-        .unwrap_or_else(|err| {
-            error!("{}", err);
-            process::exit(1);
-        });
+    let engine = DevmodeEngine::new(no_forks);
+    driver.start(endpoint, engine).unwrap_or_else(|err| {
+        error!("{}", err);
+        process::exit(1);
+    });
 }