Skip to content

Commit a477798

Browse files
icewind1991icewind1991
authored
Merge pull request #48 from esp-rs/merge-segments
merge adjecent segments and better fit ram segments between rom segments
2 parents 868fbfe + 6a24aea commit a477798

File tree

9 files changed

+211
-328
lines changed

9 files changed

+211
-328
lines changed

‎espflash/src/chip/esp32.rs

Lines changed: 99 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
1-
use bytemuck::bytes_of;
2-
use sha2::{Digest, Sha256};
3-
1+
use crate::elf::merge_segments;
42
use crate::{
53
chip::{
64
encode_flash_size, get_segment_padding, save_flash_segment, save_segment, Chip, ChipType,
@@ -10,7 +8,8 @@ use crate::{
108
elf::{FirmwareImage, RomSegment, ESP_CHECKSUM_MAGIC},
119
Error, PartitionTable,
1210
};
13-
11+
use bytemuck::bytes_of;
12+
use sha2::{Digest, Sha256};
1413
use std::{borrow::Cow, io::Write, iter::once};
1514

1615
pub struct Esp32;
@@ -75,95 +74,6 @@ impl ChipType for Esp32 {
7574
};
7675
let partition_table = partition_table.to_bytes();
7776

78-
fn get_data<'a>(image: &'a FirmwareImage) -> Result<RomSegment<'a>, Error> {
79-
let mut data = Vec::new();
80-
81-
let header = EspCommonHeader {
82-
magic: ESP_MAGIC,
83-
segment_count: 0,
84-
flash_mode: image.flash_mode as u8,
85-
flash_config: encode_flash_size(image.flash_size)? + image.flash_frequency as u8,
86-
entry: image.entry,
87-
};
88-
data.write_all(bytes_of(&header))?;
89-
90-
let extended_header = ExtendedHeader {
91-
wp_pin: WP_PIN_DISABLED,
92-
clk_q_drv: 0,
93-
d_cs_drv: 0,
94-
gd_wp_drv: 0,
95-
chip_id: 0,
96-
min_rev: 0,
97-
padding: [0; 8],
98-
append_digest: 1,
99-
};
100-
data.write_all(bytes_of(&extended_header))?;
101-
102-
let mut checksum = ESP_CHECKSUM_MAGIC;
103-
104-
let _ = image.segments().collect::<Vec<_>>();
105-
106-
let mut flash_segments: Vec<_> = image.rom_segments(Chip::Esp32).collect();
107-
flash_segments.sort();
108-
let mut ram_segments: Vec<_> = image.ram_segments(Chip::Esp32).collect();
109-
ram_segments.sort();
110-
let mut ram_segments = ram_segments.into_iter();
111-
112-
let mut segment_count = 0;
113-
114-
for segment in flash_segments {
115-
loop {
116-
let pad_len = get_segment_padding(data.len(), &segment);
117-
if pad_len > 0 {
118-
if pad_len > SEG_HEADER_LEN {
119-
if let Some(ram_segment) = ram_segments.next() {
120-
checksum = save_segment(&mut data, &ram_segment, checksum)?;
121-
segment_count += 1;
122-
continue;
123-
}
124-
}
125-
let pad_header = SegmentHeader {
126-
addr: 0,
127-
length: pad_len as u32,
128-
};
129-
data.write_all(bytes_of(&pad_header))?;
130-
for _ in 0..pad_len {
131-
data.write_all(&[0])?;
132-
}
133-
segment_count += 1;
134-
} else {
135-
break;
136-
}
137-
}
138-
checksum = save_flash_segment(&mut data, &segment, checksum)?;
139-
segment_count += 1;
140-
}
141-
142-
for segment in ram_segments {
143-
checksum = save_segment(&mut data, &segment, checksum)?;
144-
segment_count += 1;
145-
}
146-
147-
let padding = 15 - (data.len() % 16);
148-
let padding = &[0u8; 16][0..padding as usize];
149-
data.write_all(padding)?;
150-
151-
data.write_all(&[checksum])?;
152-
153-
// since we added some dummy segments, we need to patch the segment count
154-
data[1] = segment_count as u8;
155-
156-
let mut hasher = Sha256::new();
157-
hasher.update(&data);
158-
let hash = hasher.finalize();
159-
data.write_all(&hash)?;
160-
161-
Ok(RomSegment {
162-
addr: APP_ADDR,
163-
data: Cow::Owned(data),
164-
})
165-
}
166-
16777
Box::new(
16878
once(Ok(RomSegment {
16979
addr: BOOT_ADDR,
@@ -173,7 +83,7 @@ impl ChipType for Esp32 {
17383
addr: PARTION_ADDR,
17484
data: Cow::Owned(partition_table),
17585
})))
176-
.chain(once(get_data(image))),
86+
.chain(once(get_data(image, 0))),
17787
)
17888
}
17989
}
@@ -196,3 +106,98 @@ fn test_esp32_rom() {
196106
assert_eq!(expected_bin.len(), buff.len());
197107
assert_eq!(&expected_bin.as_slice(), &buff);
198108
}
109+
110+
// shared between all esp32 family chips
111+
pub(crate) fn get_data<'a>(
112+
image: &'a FirmwareImage,
113+
chip_id: u16,
114+
) -> Result<RomSegment<'a>, Error> {
115+
let mut data = Vec::new();
116+
117+
let header = EspCommonHeader {
118+
magic: ESP_MAGIC,
119+
segment_count: 0,
120+
flash_mode: image.flash_mode as u8,
121+
flash_config: encode_flash_size(image.flash_size)? + image.flash_frequency as u8,
122+
entry: image.entry,
123+
};
124+
data.write_all(bytes_of(&header))?;
125+
126+
let extended_header = ExtendedHeader {
127+
wp_pin: WP_PIN_DISABLED,
128+
clk_q_drv: 0,
129+
d_cs_drv: 0,
130+
gd_wp_drv: 0,
131+
chip_id,
132+
min_rev: 0,
133+
padding: [0; 8],
134+
append_digest: 1,
135+
};
136+
data.write_all(bytes_of(&extended_header))?;
137+
138+
let mut checksum = ESP_CHECKSUM_MAGIC;
139+
140+
let _ = image.segments().collect::<Vec<_>>();
141+
142+
let flash_segments: Vec<_> = merge_segments(image.rom_segments(Chip::Esp32s2).collect());
143+
let mut ram_segments: Vec<_> = merge_segments(image.ram_segments(Chip::Esp32s2).collect());
144+
145+
let mut segment_count = 0;
146+
147+
for segment in flash_segments {
148+
loop {
149+
let pad_len = get_segment_padding(data.len(), &segment);
150+
if pad_len > 0 {
151+
if pad_len > SEG_HEADER_LEN {
152+
if let Some(ram_segment) = ram_segments.first_mut() {
153+
// save up to `pad_len` from the ram segment, any remaining bits in the ram segments will be saved later
154+
let pad_segment = ram_segment.split_off(pad_len as usize);
155+
checksum = save_segment(&mut data, &pad_segment, checksum)?;
156+
if ram_segment.data().is_empty() {
157+
ram_segments.remove(0);
158+
}
159+
segment_count += 1;
160+
continue;
161+
}
162+
}
163+
let pad_header = SegmentHeader {
164+
addr: 0,
165+
length: pad_len as u32,
166+
};
167+
data.write_all(bytes_of(&pad_header))?;
168+
for _ in 0..pad_len {
169+
data.write_all(&[0])?;
170+
}
171+
segment_count += 1;
172+
} else {
173+
break;
174+
}
175+
}
176+
checksum = save_flash_segment(&mut data, &segment, checksum)?;
177+
segment_count += 1;
178+
}
179+
180+
for segment in ram_segments {
181+
checksum = save_segment(&mut data, &segment, checksum)?;
182+
segment_count += 1;
183+
}
184+
185+
let padding = 15 - (data.len() % 16);
186+
let padding = &[0u8; 16][0..padding as usize];
187+
data.write_all(padding)?;
188+
189+
data.write_all(&[checksum])?;
190+
191+
// since we added some dummy segments, we need to patch the segment count
192+
data[1] = segment_count as u8;
193+
194+
let mut hasher = Sha256::new();
195+
hasher.update(&data);
196+
let hash = hasher.finalize();
197+
data.write_all(&hash)?;
198+
199+
Ok(RomSegment {
200+
addr: APP_ADDR,
201+
data: Cow::Owned(data),
202+
})
203+
}

‎espflash/src/chip/esp32c3.rs

Lines changed: 5 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,10 @@
1-
use bytemuck::bytes_of;
2-
use sha2::{Digest, Sha256};
3-
1+
use crate::chip::esp32::get_data;
42
use crate::{
5-
chip::{
6-
encode_flash_size, get_segment_padding, save_flash_segment, save_segment, Chip, ChipType,
7-
EspCommonHeader, ExtendedHeader, SegmentHeader, SpiRegisters, ESP_MAGIC, SEG_HEADER_LEN,
8-
WP_PIN_DISABLED,
9-
},
10-
elf::{FirmwareImage, RomSegment, ESP_CHECKSUM_MAGIC},
3+
chip::{ChipType, SpiRegisters},
4+
elf::{FirmwareImage, RomSegment},
115
Error, PartitionTable,
126
};
13-
14-
use std::{borrow::Cow, io::Write, iter::once};
7+
use std::{borrow::Cow, iter::once};
158

169
pub struct Esp32c3;
1710

@@ -76,95 +69,6 @@ impl ChipType for Esp32c3 {
7669
};
7770
let partition_table = partition_table.to_bytes();
7871

79-
fn get_data<'a>(image: &'a FirmwareImage) -> Result<RomSegment<'a>, Error> {
80-
let mut data = Vec::new();
81-
82-
let header = EspCommonHeader {
83-
magic: ESP_MAGIC,
84-
segment_count: 0,
85-
flash_mode: image.flash_mode as u8,
86-
flash_config: encode_flash_size(image.flash_size)? + image.flash_frequency as u8,
87-
entry: image.entry,
88-
};
89-
data.write_all(bytes_of(&header))?;
90-
91-
let extended_header = ExtendedHeader {
92-
wp_pin: WP_PIN_DISABLED,
93-
clk_q_drv: 0,
94-
d_cs_drv: 0,
95-
gd_wp_drv: 0,
96-
chip_id: 5,
97-
min_rev: 0,
98-
padding: [0; 8],
99-
append_digest: 1,
100-
};
101-
data.write_all(bytes_of(&extended_header))?;
102-
103-
let mut checksum = ESP_CHECKSUM_MAGIC;
104-
105-
let _ = image.segments().collect::<Vec<_>>();
106-
107-
let mut flash_segments: Vec<_> = image.rom_segments(Chip::Esp32c3).collect();
108-
flash_segments.sort();
109-
let mut ram_segments: Vec<_> = image.ram_segments(Chip::Esp32c3).collect();
110-
ram_segments.sort();
111-
let mut ram_segments = ram_segments.into_iter();
112-
113-
let mut segment_count = 0;
114-
115-
for segment in flash_segments {
116-
loop {
117-
let pad_len = get_segment_padding(data.len(), &segment);
118-
if pad_len > 0 {
119-
if pad_len > SEG_HEADER_LEN {
120-
if let Some(ram_segment) = ram_segments.next() {
121-
checksum = save_segment(&mut data, &ram_segment, checksum)?;
122-
segment_count += 1;
123-
continue;
124-
}
125-
}
126-
let pad_header = SegmentHeader {
127-
addr: 0,
128-
length: pad_len as u32,
129-
};
130-
data.write_all(bytes_of(&pad_header))?;
131-
for _ in 0..pad_len {
132-
data.write_all(&[0])?;
133-
}
134-
segment_count += 1;
135-
} else {
136-
break;
137-
}
138-
}
139-
checksum = save_flash_segment(&mut data, &segment, checksum)?;
140-
segment_count += 1;
141-
}
142-
143-
for segment in ram_segments {
144-
checksum = save_segment(&mut data, &segment, checksum)?;
145-
segment_count += 1;
146-
}
147-
148-
let padding = 15 - (data.len() % 16);
149-
let padding = &[0u8; 16][0..padding as usize];
150-
data.write_all(padding)?;
151-
152-
data.write_all(&[checksum])?;
153-
154-
// since we added some dummy segments, we need to patch the segment count
155-
data[1] = segment_count as u8;
156-
157-
let mut hasher = Sha256::new();
158-
hasher.update(&data);
159-
let hash = hasher.finalize();
160-
data.write_all(&hash)?;
161-
162-
Ok(RomSegment {
163-
addr: APP_ADDR,
164-
data: Cow::Owned(data),
165-
})
166-
}
167-
16872
Box::new(
16973
once(Ok(RomSegment {
17074
addr: BOOT_ADDR,
@@ -174,7 +78,7 @@ impl ChipType for Esp32c3 {
17478
addr: PARTITION_ADDR,
17579
data: Cow::Owned(partition_table),
17680
})))
177-
.chain(once(get_data(image))),
81+
.chain(once(get_data(image, 5))),
17882
)
17983
}
18084
}

0 commit comments

Comments
 (0)