1 files changed, 15 insertions, 2 deletions

diff --git a/embassy-stm32-wpan/build.rs b/embassy-stm32-wpan/build.rs
index 94aac070d..7ab458bf2 100644
--- a/embassy-stm32-wpan/build.rs
+++ b/embassy-stm32-wpan/build.rs
@@ -18,9 +18,22 @@ fn main() {
     // stm32wb tl_mbox link sections
 
     let out_file = out_dir.join("tl_mbox.x").to_string_lossy().to_string();
-    fs::write(out_file, fs::read_to_string("tl_mbox.x.in").unwrap()).unwrap();
+    let in_file;
+    if env::var_os("CARGO_FEATURE_EXTENDED").is_some() {
+        if env::vars()
+            .map(|(a, _)| a)
+            .any(|x| x.starts_with("CARGO_FEATURE_STM32WB1"))
+        {
+            in_file = "tl_mbox_extended_wb1.x.in";
+        } else {
+            in_file = "tl_mbox_extended_wbx5.x.in";
+        }
+    } else {
+        in_file = "tl_mbox.x.in";
+    }
+    fs::write(out_file, fs::read_to_string(in_file).unwrap()).unwrap();
     println!("cargo:rustc-link-search={}", out_dir.display());
-    println!("cargo:rerun-if-changed=tl_mbox.x.in");
+    println!("cargo:rerun-if-changed={}", in_file);
 }
 
 enum GetOneError {