tools/embed: refine null segment patcher function

Use lone's ELF functions and types to implement the segment patcher. Gets rid of duplicated code for each ELF class, properly handles ELF data encodings and improves error handling.
lone-lang · Aug 22, 2024 · 0e4c99c · 0e4c99c
1 parent 3a8aa2b
commit 0e4c99c
Showing 1 changed file with 50 additions and 88 deletions.
diff --git a/source/tools/lone-embed.c b/source/tools/lone-embed.c
@@ -431,110 +431,72 @@ static void move_and_expand_pht_if_needed(struct elf *elf)
 	analyze(elf);
 }
 
-static void set_lone_segments(struct elf *elf)
+static void
+set_segment(struct elf *elf, struct lone_elf_header *header,
+		struct lone_elf_segment *segment, enum lone_elf_segment_type type)
 {
-	size_t entry_count = elf->segments.table.segment.count;
-	void *table = elf->segments.table.segments;
-	bool set_load_segment = false, set_lone_segment = false;
-	Elf32_Phdr *phdr32;
-	Elf64_Phdr *phdr64;
-	size_t i;
-
-	if (!has_required_null_segments(elf)) {
-		/* Not enough null segments to patch this ELF */ linux_exit(9);
-	}
-
-	switch (elf->class) {
-	case ELFCLASS64:
-		for (i = 0; i < entry_count; ++i) {
-			phdr64 = ((Elf64_Phdr *) table) + i;
+	lone_elf_umax address, size;
+	bool set_true_size;
 
-			switch (phdr64->p_type) {
-			case PT_NULL: // linker allocated spare segment
+	set_true_size = type == LONE_ELF_SEGMENT_TYPE_LONE;
 
-				if (!set_load_segment) {
-					phdr64->p_type = PT_LOAD;
+	lone_elf_segment_write_type(header, segment, type);
+	lone_elf_segment_write_file_offset(header, segment, elf->data.offset);
 
-					phdr64->p_offset = elf->data.offset;
-					phdr64->p_vaddr = phdr64->p_paddr = align_to_page(elf, elf->limits.end.virtual);
-					phdr64->p_filesz = phdr64->p_memsz = align_to_page(elf, elf->data.size);
-					phdr64->p_align = elf->page_size;
-					phdr64->p_flags = PF_R;
+	address = align_to_page(elf, elf->limits.end.virtual);
+	lone_elf_segment_write_virtual_address(header, segment, address);
+	lone_elf_segment_write_physical_address(header, segment, address);
 
-					set_load_segment = true;
+	size = elf->data.size;
+	size = set_true_size? size : align_to_page(elf, size);
+	lone_elf_segment_write_size_in_file(header, segment, size);
+	lone_elf_segment_write_size_in_memory(header, segment, size);
 
-				} else if (!set_lone_segment) {
-					phdr64->p_type = PT_LONE;
+	lone_elf_segment_write_alignment(header, segment, set_true_size? 1 : elf->page_size);
+	lone_elf_segment_write_flags(header, segment, LONE_ELF_SEGMENT_FLAGS_READABLE);
+}
 
-					phdr64->p_offset = elf->data.offset;
-					phdr64->p_vaddr = phdr64->p_paddr = align_to_page(elf, elf->limits.end.virtual);
-					phdr64->p_filesz = phdr64->p_memsz = elf->data.size;
-					phdr64->p_align = 1;
-					phdr64->p_flags = PF_R;
+static void set_lone_segments(struct elf *elf)
+{
+	struct lone_elf_header *header;
+	struct lone_elf_segments segments;
+	struct lone_elf_segment *segment;
+	struct lone_optional_u32 type;
+	bool set_load_segment, set_lone_segment;
+	lone_u16 i;
 
-					set_lone_segment = true;
+	if (!has_required_null_segments(elf)) { not_enough_nulls(); }
 
-				} else {
-					break;
-				}
+	header = hdr(elf);
+	segments = elf->segments.table;
+	set_load_segment = false;
+	set_lone_segment = false;
 
-			__attribute__((fallthrough));
-			default:
-				continue;
-			}
+	for (i = 0; i < segments.segment.count; ++i) {
+		segment = lone_elf_segment_at(segments, i);
+		type = lone_elf_segment_read_type(header, segment);
+		if (!type.present) { invalid_elf(); }
 
-			if (set_lone_segment && set_load_segment) {
+		switch ((enum lone_elf_segment_type) type.value) {
+		case LONE_ELF_SEGMENT_TYPE_NULL: // linker allocated spare segment
+			if (!set_load_segment) {
+				set_segment(elf, header, segment, LONE_ELF_SEGMENT_TYPE_LOAD);
+				set_load_segment = true;
+			} else if (!set_lone_segment) {
+				set_segment(elf, header, segment, LONE_ELF_SEGMENT_TYPE_LONE);
+				set_lone_segment = true;
+			} else {
 				break;
 			}
-		}
 
-		break;
-	case ELFCLASS32:
-		for (i = 0; i < entry_count; ++i) {
-			phdr32 = ((Elf32_Phdr *) table) + i;
-
-			switch (phdr32->p_type) {
-			case PT_NULL: // linker allocated spare segment
-
-				if (!set_load_segment) {
-					phdr32->p_type = PT_LOAD;
-
-					phdr32->p_offset = elf->data.offset;
-					phdr32->p_vaddr = phdr32->p_paddr = align_to_page(elf, elf->limits.end.virtual);
-					phdr32->p_filesz = phdr32->p_memsz = align_to_page(elf, elf->data.size);
-					phdr32->p_align = elf->page_size;
-					phdr32->p_flags = PF_R;
-
-					set_load_segment = true;
-
-				} else if (!set_lone_segment) {
-					phdr32->p_type = PT_LONE;
-
-					phdr32->p_offset = elf->data.offset;
-					phdr32->p_vaddr = phdr32->p_paddr = align_to_page(elf, elf->limits.end.virtual);
-					phdr32->p_filesz = phdr32->p_memsz = elf->data.size;
-					phdr32->p_align = 1;
-					phdr32->p_flags = PF_R;
-
-					set_lone_segment = true;
-
-				} else {
-					break;
-				}
-
-			__attribute__((fallthrough));
-			default:
-				continue;
-			}
-
-			if (set_lone_segment && set_load_segment) {
-				break;
-			}
+		__attribute__((fallthrough));
+		default:
+			continue;
 		}
 
-		break;
-	default:
-		/* Invalid ELF class but somehow made it here? */ linux_exit(8);
+		if (set_lone_segment && set_load_segment) {
+			break;
+		}
 	}
 }