forked from lolo859/vystem
184 lines
6.8 KiB
C++
184 lines
6.8 KiB
C++
// SPDX-License-Identifier: MPL-2.0
|
|
#include "vybuild.hpp"
|
|
#include <zlib.h>
|
|
using namespace std;
|
|
#pragma pack(1)
|
|
struct mbr_entry {
|
|
uint8_t boot_indicator=0x00;
|
|
uint8_t start_chs[3]={0,0x02,0};
|
|
uint8_t partition_type=0xEE;
|
|
uint8_t end_chs[3]={0xFF,0xFF,0xFF};
|
|
uint32_t start_lba=1;
|
|
uint32_t last_lba=0xFFFFFFFF;
|
|
};
|
|
#pragma pack()
|
|
#pragma pack(1)
|
|
struct mbr_protect {
|
|
uint8_t bootstrap_code[440];
|
|
uint32_t disk_signature;
|
|
uint16_t reserved=0;
|
|
mbr_entry first_entry;
|
|
uint8_t other_entries[48]={0};
|
|
uint16_t signature=0xAA55;
|
|
};
|
|
#pragma pack()
|
|
#pragma pack(1)
|
|
struct gpt_header {
|
|
char signature[8]={'E','F','I',' ','P','A','R','T'};
|
|
uint32_t revision=0x00010000;
|
|
uint32_t header_size=92;
|
|
uint32_t header_crc32;
|
|
uint32_t reserved=0;
|
|
uint64_t current_lba=1;
|
|
uint64_t backup_lba;
|
|
uint64_t first_usable_lba;
|
|
uint64_t last_usable_lba;
|
|
uint8_t disk_guid[16];
|
|
uint64_t partition_entries_lba=2;
|
|
uint32_t num_partitions_entries=128;
|
|
uint32_t size_of_partitions_entries=128;
|
|
uint32_t partitions_entries_crc32;
|
|
};
|
|
#pragma pack()
|
|
#pragma pack(1)
|
|
struct partition_entry {
|
|
uint8_t type_guid[16];
|
|
uint8_t unique_guid[16];
|
|
uint64_t start_lba;
|
|
uint64_t last_lba;
|
|
uint64_t attributes;
|
|
uint16_t name[36];
|
|
};
|
|
#pragma pack()
|
|
struct disk_container {
|
|
mbr_protect mbr;
|
|
gpt_header header;
|
|
array<partition_entry,128> parts;
|
|
vector<uint8_t> data;
|
|
size_t part_count=0;
|
|
size_t lba_bumb=0;
|
|
};
|
|
array<uint8_t,16> vybuild::disk::uuid4_to_bytes(const string& uuid) {
|
|
array<uint8_t,16> raw{};
|
|
int byte_index=0;
|
|
for (size_t i=0;i<uuid.size();i++) {
|
|
char c=uuid[i];
|
|
if (c=='-') continue;
|
|
auto hex_val=[](char x) -> uint8_t {
|
|
if (x>='0' && x<='9') return x-'0';
|
|
if (x>='a' && x<='f') return x-'a'+10;
|
|
if (x>='A' && x<='F') return x-'A'+10;
|
|
throw std::runtime_error("Invalid hex char in UUID.");
|
|
};
|
|
uint8_t high=hex_val(c);
|
|
if (++i>=uuid.size()) throw std::runtime_error("Malformed UUID.");
|
|
uint8_t low=hex_val(uuid[i]);
|
|
if (byte_index>=16) throw std::runtime_error("UUID too long.");
|
|
raw[byte_index++]=(high<<4)|low;
|
|
}
|
|
if (byte_index!=16) throw std::runtime_error("Invalid UUID size.");
|
|
array<uint8_t,16> out{};
|
|
out[0]=raw[3];
|
|
out[1]=raw[2];
|
|
out[2]=raw[1];
|
|
out[3]=raw[0];
|
|
out[4]=raw[5];
|
|
out[5]=raw[4];
|
|
out[6]=raw[7];
|
|
out[7]=raw[6];
|
|
for (size_t i=8;i<16;i++) out[i]=raw[i];
|
|
return out;
|
|
}
|
|
map<string,disk_container> disk_map;
|
|
void vybuild::disk::create_disk(string key,size_t size,string disk_guid) {
|
|
if (size%512!=0) throw std::runtime_error("Disk image size isn't a multiple of 512 bytes.");
|
|
disk_container cont;
|
|
size_t lba_count=size/512;
|
|
size_t last_lba=lba_count-1;
|
|
cont.header.backup_lba=last_lba;
|
|
size_t first_usable_lba=2+((cont.header.size_of_partitions_entries*cont.header.num_partitions_entries)/512);
|
|
cont.header.first_usable_lba=first_usable_lba;
|
|
cont.header.last_usable_lba=last_lba-1;
|
|
auto disk_id=uuid4_to_bytes(disk_guid);
|
|
memcpy(cont.header.disk_guid,disk_id.data(),16);
|
|
memset(cont.parts.data(),0,sizeof(cont.parts));
|
|
cont.header.partitions_entries_crc32=crc32(0,(uint8_t*)cont.parts.data(),sizeof(cont.parts));
|
|
cont.header.header_crc32=0;
|
|
cont.header.header_crc32=crc32(0,(uint8_t*)&cont.header,sizeof(cont.header));
|
|
cont.data.resize((cont.header.last_usable_lba-cont.header.first_usable_lba)*512);
|
|
disk_map[key]=cont;
|
|
return;
|
|
}
|
|
array<uint16_t,36> gpt_name_to_utf16le(const string& str) {
|
|
std::array<uint16_t,36> out{};
|
|
size_t len=str.size();
|
|
if (len>35) len=35;
|
|
for (size_t i=0;i<len;i++) out[i]=static_cast<uint16_t>(static_cast<unsigned char>(str[i]));
|
|
out[len]=0;
|
|
return out;
|
|
}
|
|
void add_partition(disk_container& disk,size_t partition_size_bytes,string type_guid,string unique_guid,uint64_t attributes,string name,uint8_t* in_buf,size_t in_size) {
|
|
if (partition_size_bytes==0 || partition_size_bytes%512!=0) throw std::runtime_error("Invalid partition size.");
|
|
if ((disk.lba_bumb*512)+partition_size_bytes>=disk.data.size()) throw std::runtime_error("Provided partition size would overflow the disk.");
|
|
if (in_size>partition_size_bytes) throw std::runtime_error("Provided buffer is too big for provided partition size.");
|
|
if (disk.part_count==128) throw std::runtime_error("Too much partitions.");
|
|
partition_entry entry{};
|
|
entry.attributes=attributes;
|
|
size_t lba_count=partition_size_bytes/512;
|
|
entry.start_lba=disk.header.first_usable_lba+disk.lba_bumb;
|
|
entry.last_lba=entry.start_lba+lba_count-1;
|
|
auto gpt_name=gpt_name_to_utf16le(name);
|
|
memcpy(entry.name,gpt_name.data(),sizeof(gpt_name));
|
|
auto type=vybuild::disk::uuid4_to_bytes(type_guid);
|
|
auto unique=vybuild::disk::uuid4_to_bytes(unique_guid);
|
|
memcpy(entry.type_guid,type.data(),16);
|
|
memcpy(entry.unique_guid,unique.data(),16);
|
|
disk.parts[disk.part_count]=entry;
|
|
memcpy(disk.data.data()+disk.lba_bumb*512,in_buf,in_size);
|
|
disk.part_count++;
|
|
disk.lba_bumb+=lba_count;
|
|
return;
|
|
}
|
|
void vybuild::disk::add_partition(string key,size_t partition_size_bytes,string type_guid,string unique_guid,uint64_t attributes,string name,uint8_t *in_buf,size_t in_size) {
|
|
if (disk_map.find(key)==disk_map.end()) throw std::runtime_error(key+" isn't a valid disk key.");
|
|
disk_container& disk=disk_map[key];
|
|
add_partition(disk,partition_size_bytes,type_guid,unique_guid,attributes,name,in_buf,in_size);
|
|
return;
|
|
}
|
|
static void write_sector(vector<uint8_t>& out,size_t lba,const void* data,size_t size) {
|
|
size_t offset=lba*512;
|
|
if (out.size()<offset+size) out.resize(offset+size,0);
|
|
memcpy(out.data()+offset,data,size);
|
|
return;
|
|
}
|
|
vector<uint8_t> export_disk(disk_container& cont) {
|
|
vector<uint8_t> disk;
|
|
size_t total_sectors=(cont.header.backup_lba+1);
|
|
disk.resize(total_sectors*512,0);
|
|
write_sector(disk,0,&cont.mbr,sizeof(cont.mbr));
|
|
gpt_header header=cont.header;
|
|
auto parts_crc=crc32(0,(uint8_t*)cont.parts.data(),sizeof(cont.parts));
|
|
header.partitions_entries_crc32=parts_crc;
|
|
header.header_crc32=0;
|
|
header.header_crc32=crc32(0,(uint8_t*)&header,sizeof(header));
|
|
write_sector(disk,1,&header,sizeof(header));
|
|
write_sector(disk,cont.header.partition_entries_lba,cont.parts.data(),sizeof(cont.parts));
|
|
if (!cont.data.empty()) {
|
|
size_t data_lba=cont.header.first_usable_lba;
|
|
write_sector(disk,data_lba,cont.data.data(),cont.data.size());
|
|
}
|
|
write_sector(disk,cont.header.backup_lba-(sizeof(cont.parts)/512),cont.parts.data(),sizeof(cont.parts));
|
|
gpt_header backup=cont.header;
|
|
backup.current_lba=cont.header.backup_lba;
|
|
backup.backup_lba=1;
|
|
backup.header_crc32=0;
|
|
backup.header_crc32=crc32(0,(uint8_t*)&backup,sizeof(backup));
|
|
write_sector(disk,cont.header.backup_lba,&backup,sizeof(backup));
|
|
return disk;
|
|
}
|
|
vector<uint8_t> vybuild::disk::export_disk(string key) {
|
|
if (disk_map.find(key)==disk_map.end()) throw std::runtime_error(key+" isn't a valid disk key.");
|
|
disk_container& fat=disk_map[key];
|
|
return export_disk(fat);
|
|
}
|