Merge pull request #145 from adarna-crypto/variable-push-fix

fixes issue #144
quilt · Apr 20, 2024 · 9123fe1 · 9123fe1
2 parents 5b0c962 + 8e675cb
commit 9123fe1
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 4 deletions.
diff --git a/etk-asm/src/asm.rs b/etk-asm/src/asm.rs
@@ -218,7 +218,7 @@ pub struct Assembler {
     undeclared_labels: HashSet<String>,
 
     /// Pushes that are variable-sized and need to be backpatched.
-    variable_sized_push: Vec<AbstractOp>,
+    variable_sized_push: Vec<PushDef>,
 }
 
 /// A label definition.
@@ -243,6 +243,30 @@ impl LabelDef {
     }
 }
 
+/// A push definition.
+#[derive(Clone, Debug, PartialEq)]
+pub struct PushDef {
+    position: usize,
+    op: AbstractOp,
+}
+
+impl PushDef {
+    /// Create a new `PushDef`.
+    pub fn new(op: AbstractOp, position: usize) -> Self {
+        Self { op, position }
+    }
+
+    /// Get the position of the push.
+    pub fn position(&self) -> usize {
+        self.position
+    }
+
+    /// Get the op from the push.
+    pub fn op(&self) -> &AbstractOp {
+        &self.op
+    }
+}
+
 impl Assembler {
     /// Create a new `Assembler`.
     pub fn new() -> Self {
@@ -358,7 +382,10 @@ impl Assembler {
                             // Here, we set the size of the push to 2 bytes (min possible value),
                             //  as we don't know the final value of the label yet.
                             self.concrete_len += 2;
-                            self.variable_sized_push.push(op.clone());
+                            self.variable_sized_push.push(PushDef {
+                                position: self.concrete_len,
+                                op: op.clone(),
+                            });
                         } else {
                             self.concrete_len += op.size().unwrap();
                         }
@@ -390,8 +417,8 @@ impl Assembler {
     }
 
     fn backpatch_labels(&mut self) -> Result<(), Error> {
-        for op in self.variable_sized_push.iter() {
-            if let AbstractOp::Push(imm) = op {
+        for pushdef in self.variable_sized_push.iter() {
+            if let AbstractOp::Push(imm) = &pushdef.op {
                 let exp = imm
                     .tree
                     .eval_with_context((&self.declared_labels, &self.declared_macros).into());
@@ -403,7 +430,12 @@ impl Assembler {
                     if imm_size > 1 {
                         for label_value in self.declared_labels.values_mut() {
                             let labeldef = label_value.as_ref().unwrap();
+                            if labeldef.position < pushdef.position {
+                                // don't move labels that are declared earlier than this push
+                                continue;
+                            };
                             self.concrete_len += imm_size as usize - 1;
+
                             *label_value = Some(LabelDef {
                                 position: labeldef.position + imm_size as usize - 1,
                                 updated: true,

diff --git a/etk-asm/tests/asm.rs b/etk-asm/tests/asm.rs
@@ -321,5 +321,11 @@ fn test_variable_sized_push2() -> Result<(), Error> {
     ingester.ingest_file(source(&["variable-push2", "main3.etk"]))?;
     assert_eq!(output, hex!("610107015801"));
 
+    let mut output = Vec::new();
+    let mut ingester = Ingest::new(&mut output);
+    ingester.ingest_file(source(&["variable-push2", "main4.etk"]))?;
+    let first_push = output.get(0..2).unwrap();
+    assert_eq!(first_push, hex!("6004"));
+
     Ok(())
 }
diff --git a/etk-asm/tests/asm/variable-push2/main4.etk b/etk-asm/tests/asm/variable-push2/main4.etk
@@ -0,0 +1,32 @@
+    %macro insert_instructions()
+        push32 0x01
+        pop
+        push32 0x02
+        pop
+        push32 0x03
+        pop
+        push32 0x04
+        pop
+        push32 0x05
+        pop
+        push32 0x06
+        pop
+        push32 0x07
+        pop
+        push32 0x08
+        pop
+    %end
+
+    %push(first)
+    jump
+    stop
+first:
+    jumpdest
+    pc
+    pc
+    push1 0x01
+    %push(second)
+    jumpi
+    %insert_instructions()
+second:
+    jumpdest