fix: resolve UB in _add, __sub, _mul for u8, u16, u32

sway-lib-core/src/ops.sw

@@ -173,17 +173,38 @@ impl Subtract for u64 {
     }
 }
 
-// unlike addition, underflowing subtraction does not need special treatment
-// because VM handles underflow
 impl Subtract for u32 {
     fn subtract(self, other: Self) -> Self {
-        __sub(self, other)
+        // any non-64-bit value is compiled to a u64 value under-the-hood
+        // constants (like Self::max() below) are also automatically promoted to u64
+        let res = __sub(self, other);
+        // integer underflow
+        if __gt(res, Self::max()) {
+            if panic_on_overflow_is_enabled() {
+                __revert(0)
+            } else {
+                // overflow enabled
+                __mod(res, __add(Self::max(), 1))
+            }
+        } else {
+            // no overflow
+            res
+        }
     }
 }
 
 impl Subtract for u16 {
     fn subtract(self, other: Self) -> Self {
-        __sub(self, other)
+        let res = __sub(self, other);
+        if __gt(res, Self::max()) {
+            if panic_on_overflow_is_enabled() {
+                __revert(0)
+            } else {
+                __mod(res, __add(Self::max(), 1))
+            }
+        } else {
+            res
+        }
     }
 }
 

test/src/in_language_tests/test_programs/math_inline_tests/src/main.sw

@@ -900,6 +900,22 @@ fn math_u256_overflow_mul_revert() {
     log(b);
 }
 
+#[test(should_revert)]
+fn math_u16_underflow_sub_revert() {
+    let a = 0u16;
+    let b = 1u16;
+    let c = a - b;
+    log(c);
+}
+
+#[test(should_revert)]
+fn math_u32_underflow_sub_revert() {
+    let a = 0u32;
+    let b = 1u32;
+    let c = a - b;
+    log(c);
+}
+
 #[test]
 fn math_u8_overflow_add() {
     let _ = disable_panic_on_overflow();
@@ -944,6 +960,26 @@ fn math_u16_overflow_add() {
     require(e == u16::max() - 2, e);
 }
 
+#[test]
+fn math_u16_underflow_sub() {
+    assert((u16::max() - u16::max()) == 0u16);
+    assert((u16::min() - u16::min()) == 0u16);
+    assert((10u16 - 5u16) == 5u16);
+
+    let _ = disable_panic_on_overflow();
+
+    let a = 0u16;
+    let b = 1u16;
+
+    let c = a - b;
+    assert(c == u16::max());
+
+    let d = u16::max();
+
+    let e = a - d;
+    assert(e == b);
+}
+
 #[test]
 fn math_u32_overflow_add() {
     let _ = disable_panic_on_overflow();
@@ -966,6 +1002,26 @@ fn math_u32_overflow_add() {
     require(e == u32::max() - 2, e);
 }
 
+#[test]
+fn math_u32_underflow_sub() {
+    assert((u32::max() - u32::max()) == 0u32);
+    assert((u32::min() - u32::min()) == 0u32);
+    assert((10u32 - 5u32) == 5u32);
+
+    let _ = disable_panic_on_overflow();
+
+    let a = 0u32;
+    let b = 1u32;
+
+    let c = a - b;
+    assert(c == u32::max());
+
+    let d = u32::max();
+
+    let e = a - d;
+    assert(e == b);
+}
+
 #[test]
 fn math_u64_overflow_add() {
     let _ = disable_panic_on_overflow();