1use byteorder::{ByteOrder, LittleEndian};
54
55use crate::block_storage::BlockCompression;
56use crate::{Result, SochDBError};
57
58#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
60#[repr(u8)]
61pub enum FormatVersion {
62 V1 = 1,
64 V2 = 2,
66}
67
68impl FormatVersion {
69 pub const CURRENT: FormatVersion = FormatVersion::V2;
71
72 pub fn from_byte(b: u8) -> Option<Self> {
74 match b {
75 1 => Some(FormatVersion::V1),
76 2 => Some(FormatVersion::V2),
77 _ => None,
78 }
79 }
80
81 pub fn header_size(&self) -> usize {
83 match self {
84 FormatVersion::V1 => V1_HEADER_SIZE,
85 FormatVersion::V2 => V2_HEADER_SIZE,
86 }
87 }
88}
89
90const V1_MAGIC: [u8; 4] = *b"TBLK";
92const V2_MAGIC: [u8; 4] = *b"TBL2";
93
94const V1_HEADER_SIZE: usize = 17;
96const V2_HEADER_SIZE: usize = 21;
97
98#[derive(Debug, Clone, Copy, Default)]
100pub struct BlockFlags {
101 pub encrypted: bool,
103 pub extended_checksum: bool,
105 pub spanning: bool,
107 pub has_metadata: bool,
109}
110
111impl BlockFlags {
112 pub fn to_byte(&self) -> u8 {
114 let mut b = 0u8;
115 if self.encrypted {
116 b |= 0x01;
117 }
118 if self.extended_checksum {
119 b |= 0x02;
120 }
121 if self.spanning {
122 b |= 0x04;
123 }
124 if self.has_metadata {
125 b |= 0x08;
126 }
127 b
128 }
129
130 pub fn from_byte(b: u8) -> Self {
132 Self {
133 encrypted: (b & 0x01) != 0,
134 extended_checksum: (b & 0x02) != 0,
135 spanning: (b & 0x04) != 0,
136 has_metadata: (b & 0x08) != 0,
137 }
138 }
139}
140
141#[derive(Debug, Clone)]
150pub struct V1Header {
151 pub compression: BlockCompression,
152 pub original_size: u32,
153 pub compressed_size: u32,
154 pub checksum: u32,
155}
156
157impl V1Header {
158 pub fn from_bytes(buf: &[u8]) -> Result<Self> {
160 if buf.len() < V1_HEADER_SIZE {
161 return Err(SochDBError::InvalidData(format!(
162 "V1 header too short: {} < {}",
163 buf.len(),
164 V1_HEADER_SIZE
165 )));
166 }
167
168 if &buf[0..4] != V1_MAGIC.as_slice() {
169 return Err(SochDBError::InvalidData(format!(
170 "Invalid V1 magic: {:?}",
171 &buf[0..4]
172 )));
173 }
174
175 Ok(Self {
176 compression: BlockCompression::from_byte(buf[4]),
177 original_size: LittleEndian::read_u32(&buf[5..9]),
178 compressed_size: LittleEndian::read_u32(&buf[9..13]),
179 checksum: LittleEndian::read_u32(&buf[13..17]),
180 })
181 }
182
183 pub fn to_bytes(&self) -> [u8; V1_HEADER_SIZE] {
185 let mut buf = [0u8; V1_HEADER_SIZE];
186 buf[0..4].copy_from_slice(&V1_MAGIC);
187 buf[4] = self.compression.to_byte();
188 LittleEndian::write_u32(&mut buf[5..9], self.original_size);
189 LittleEndian::write_u32(&mut buf[9..13], self.compressed_size);
190 LittleEndian::write_u32(&mut buf[13..17], self.checksum);
191 buf
192 }
193
194 pub fn upgrade_to_v2(&self) -> V2Header {
196 V2Header {
197 format_version: FormatVersion::V2,
198 compression: self.compression,
199 flags: BlockFlags::default(),
200 original_size: self.original_size,
201 compressed_size: self.compressed_size,
202 checksum: self.checksum,
203 }
204 }
205}
206
207#[derive(Debug, Clone)]
219pub struct V2Header {
220 pub format_version: FormatVersion,
221 pub compression: BlockCompression,
222 pub flags: BlockFlags,
223 pub original_size: u32,
224 pub compressed_size: u32,
225 pub checksum: u32,
226}
227
228impl V2Header {
229 pub fn from_bytes(buf: &[u8]) -> Result<Self> {
231 if buf.len() < V2_HEADER_SIZE {
232 return Err(SochDBError::InvalidData(format!(
233 "V2 header too short: {} < {}",
234 buf.len(),
235 V2_HEADER_SIZE
236 )));
237 }
238
239 if &buf[0..4] != V2_MAGIC.as_slice() {
240 return Err(SochDBError::InvalidData(format!(
241 "Invalid V2 magic: {:?}",
242 &buf[0..4]
243 )));
244 }
245
246 let format_version = FormatVersion::from_byte(buf[4]).ok_or_else(|| {
247 SochDBError::InvalidData(format!("Unknown format version: {}", buf[4]))
248 })?;
249
250 Ok(Self {
251 format_version,
252 compression: BlockCompression::from_byte(buf[5]),
253 flags: BlockFlags::from_byte(buf[6]),
254 original_size: LittleEndian::read_u32(&buf[7..11]),
255 compressed_size: LittleEndian::read_u32(&buf[11..15]),
256 checksum: LittleEndian::read_u32(&buf[15..19]),
257 })
258 }
259
260 pub fn to_bytes(&self) -> [u8; V2_HEADER_SIZE] {
262 let mut buf = [0u8; V2_HEADER_SIZE];
263 buf[0..4].copy_from_slice(&V2_MAGIC);
264 buf[4] = self.format_version as u8;
265 buf[5] = self.compression.to_byte();
266 buf[6] = self.flags.to_byte();
267 LittleEndian::write_u32(&mut buf[7..11], self.original_size);
268 LittleEndian::write_u32(&mut buf[11..15], self.compressed_size);
269 LittleEndian::write_u32(&mut buf[15..19], self.checksum);
270 buf
272 }
273
274 pub fn downgrade_to_v1(&self) -> V1Header {
276 V1Header {
277 compression: self.compression,
278 original_size: self.original_size,
279 compressed_size: self.compressed_size,
280 checksum: self.checksum,
281 }
282 }
283}
284
285#[derive(Debug, Clone)]
287pub enum BlockHeader {
288 V1(V1Header),
289 V2(V2Header),
290}
291
292impl BlockHeader {
293 pub fn from_bytes(buf: &[u8]) -> Result<Self> {
295 if buf.len() < 4 {
296 return Err(SochDBError::InvalidData(
297 "Buffer too short for magic detection".to_string(),
298 ));
299 }
300
301 let magic = &buf[0..4];
302
303 if magic == V1_MAGIC.as_slice() {
304 Ok(BlockHeader::V1(V1Header::from_bytes(buf)?))
305 } else if magic == V2_MAGIC.as_slice() {
306 Ok(BlockHeader::V2(V2Header::from_bytes(buf)?))
307 } else {
308 Err(SochDBError::InvalidData(format!(
309 "Unknown block magic: {:?}",
310 magic
311 )))
312 }
313 }
314
315 pub fn version(&self) -> FormatVersion {
317 match self {
318 BlockHeader::V1(_) => FormatVersion::V1,
319 BlockHeader::V2(_) => FormatVersion::V2,
320 }
321 }
322
323 pub fn header_size(&self) -> usize {
325 self.version().header_size()
326 }
327
328 pub fn compression(&self) -> BlockCompression {
330 match self {
331 BlockHeader::V1(h) => h.compression,
332 BlockHeader::V2(h) => h.compression,
333 }
334 }
335
336 pub fn original_size(&self) -> u32 {
338 match self {
339 BlockHeader::V1(h) => h.original_size,
340 BlockHeader::V2(h) => h.original_size,
341 }
342 }
343
344 pub fn compressed_size(&self) -> u32 {
346 match self {
347 BlockHeader::V1(h) => h.compressed_size,
348 BlockHeader::V2(h) => h.compressed_size,
349 }
350 }
351
352 pub fn checksum(&self) -> u32 {
354 match self {
355 BlockHeader::V1(h) => h.checksum,
356 BlockHeader::V2(h) => h.checksum,
357 }
358 }
359
360 pub fn flags(&self) -> BlockFlags {
362 match self {
363 BlockHeader::V1(_) => BlockFlags::default(),
364 BlockHeader::V2(h) => h.flags,
365 }
366 }
367
368 pub fn upgrade(&self) -> Self {
370 match self {
371 BlockHeader::V1(h) => BlockHeader::V2(h.upgrade_to_v2()),
372 BlockHeader::V2(_) => self.clone(),
373 }
374 }
375
376 pub fn to_bytes(&self) -> Vec<u8> {
378 match self {
379 BlockHeader::V1(h) => h.to_bytes().to_vec(),
380 BlockHeader::V2(h) => h.to_bytes().to_vec(),
381 }
382 }
383
384 pub fn is_current(&self) -> bool {
386 self.version() == FormatVersion::CURRENT
387 }
388
389 pub fn needs_migration(&self) -> bool {
391 self.version() < FormatVersion::CURRENT
392 }
393}
394
395#[derive(Debug, Clone)]
397pub struct MigratableBlock {
398 pub header: BlockHeader,
399 pub data: Vec<u8>,
400}
401
402impl MigratableBlock {
403 pub fn new(
405 data: Vec<u8>,
406 compression: BlockCompression,
407 original_size: u32,
408 compressed_size: u32,
409 checksum: u32,
410 ) -> Self {
411 Self {
412 header: BlockHeader::V2(V2Header {
413 format_version: FormatVersion::CURRENT,
414 compression,
415 flags: BlockFlags::default(),
416 original_size,
417 compressed_size,
418 checksum,
419 }),
420 data,
421 }
422 }
423
424 pub fn with_flags(mut self, flags: BlockFlags) -> Self {
426 if let BlockHeader::V2(ref mut h) = self.header {
427 h.flags = flags;
428 }
429 self
430 }
431
432 pub fn from_bytes(buf: &[u8]) -> Result<Self> {
434 let header = BlockHeader::from_bytes(buf)?;
435 let header_size = header.header_size();
436 let data_size = header.compressed_size() as usize;
437
438 if buf.len() < header_size + data_size {
439 return Err(SochDBError::InvalidData(format!(
440 "Block buffer too short: {} < {}",
441 buf.len(),
442 header_size + data_size
443 )));
444 }
445
446 Ok(Self {
447 header,
448 data: buf[header_size..header_size + data_size].to_vec(),
449 })
450 }
451
452 pub fn to_bytes(&self) -> Vec<u8> {
454 let header_bytes = self.header.to_bytes();
455 let mut result = Vec::with_capacity(header_bytes.len() + self.data.len());
456 result.extend_from_slice(&header_bytes);
457 result.extend_from_slice(&self.data);
458 result
459 }
460
461 pub fn migrate(&mut self) {
463 self.header = self.header.upgrade();
464 }
465
466 pub fn needs_migration(&self) -> bool {
468 self.header.needs_migration()
469 }
470
471 pub fn verify_checksum(&self) -> Result<()> {
473 let computed = crc32fast::hash(&self.data);
474 let stored = self.header.checksum();
475
476 if computed != stored {
477 return Err(SochDBError::DataCorruption {
478 details: format!(
479 "Checksum mismatch: computed {} != stored {}",
480 computed, stored
481 ),
482 location: "block data".to_string(),
483 hint: "Block may be corrupted, try restoring from backup".to_string(),
484 });
485 }
486
487 Ok(())
488 }
489}
490
491#[derive(Debug, Default)]
493pub struct MigrationStats {
494 pub blocks_read: u64,
495 pub blocks_migrated: u64,
496 pub v1_blocks_found: u64,
497 pub v2_blocks_found: u64,
498 pub checksum_failures: u64,
499}
500
501impl MigrationStats {
502 pub fn record_read(&mut self, version: FormatVersion) {
503 self.blocks_read += 1;
504 match version {
505 FormatVersion::V1 => self.v1_blocks_found += 1,
506 FormatVersion::V2 => self.v2_blocks_found += 1,
507 }
508 }
509
510 pub fn record_migration(&mut self) {
511 self.blocks_migrated += 1;
512 }
513
514 pub fn record_checksum_failure(&mut self) {
515 self.checksum_failures += 1;
516 }
517
518 pub fn migration_progress(&self) -> f64 {
520 if self.v1_blocks_found == 0 {
521 100.0
522 } else {
523 (self.blocks_migrated as f64 / self.v1_blocks_found as f64) * 100.0
524 }
525 }
526}
527
528pub struct FormatMigrator {
530 stats: MigrationStats,
531 verify_checksums: bool,
532}
533
534impl FormatMigrator {
535 pub fn new() -> Self {
536 Self {
537 stats: MigrationStats::default(),
538 verify_checksums: true,
539 }
540 }
541
542 pub fn with_checksum_verification(mut self, verify: bool) -> Self {
544 self.verify_checksums = verify;
545 self
546 }
547
548 pub fn migrate_block(&mut self, block: &mut MigratableBlock) -> Result<bool> {
550 self.stats.record_read(block.header.version());
551
552 if self.verify_checksums
553 && let Err(e) = block.verify_checksum()
554 {
555 self.stats.record_checksum_failure();
556 return Err(e);
557 }
558
559 if block.needs_migration() {
560 block.migrate();
561 self.stats.record_migration();
562 Ok(true)
563 } else {
564 Ok(false)
565 }
566 }
567
568 pub fn migrate_blocks(&mut self, blocks: &mut [MigratableBlock]) -> Result<usize> {
570 let mut migrated = 0;
571 for block in blocks {
572 if self.migrate_block(block)? {
573 migrated += 1;
574 }
575 }
576 Ok(migrated)
577 }
578
579 pub fn stats(&self) -> &MigrationStats {
581 &self.stats
582 }
583}
584
585impl Default for FormatMigrator {
586 fn default() -> Self {
587 Self::new()
588 }
589}
590
591#[cfg(test)]
592mod tests {
593 use super::*;
594
595 #[test]
596 fn test_v1_header_roundtrip() {
597 let header = V1Header {
598 compression: BlockCompression::None,
599 original_size: 1024,
600 compressed_size: 1024,
601 checksum: 0xDEADBEEF,
602 };
603
604 let bytes = header.to_bytes();
605 assert_eq!(bytes.len(), V1_HEADER_SIZE);
606
607 let parsed = V1Header::from_bytes(&bytes).unwrap();
608 assert_eq!(parsed.compression, BlockCompression::None);
609 assert_eq!(parsed.original_size, 1024);
610 assert_eq!(parsed.compressed_size, 1024);
611 assert_eq!(parsed.checksum, 0xDEADBEEF);
612 }
613
614 #[test]
615 fn test_v2_header_roundtrip() {
616 let header = V2Header {
617 format_version: FormatVersion::V2,
618 compression: BlockCompression::Lz4,
619 flags: BlockFlags {
620 encrypted: true,
621 extended_checksum: false,
622 spanning: true,
623 has_metadata: false,
624 },
625 original_size: 2048,
626 compressed_size: 1500,
627 checksum: 0xCAFEBABE,
628 };
629
630 let bytes = header.to_bytes();
631 assert_eq!(bytes.len(), V2_HEADER_SIZE);
632
633 let parsed = V2Header::from_bytes(&bytes).unwrap();
634 assert_eq!(parsed.format_version, FormatVersion::V2);
635 assert_eq!(parsed.compression, BlockCompression::Lz4);
636 assert!(parsed.flags.encrypted);
637 assert!(!parsed.flags.extended_checksum);
638 assert!(parsed.flags.spanning);
639 assert!(!parsed.flags.has_metadata);
640 assert_eq!(parsed.original_size, 2048);
641 assert_eq!(parsed.compressed_size, 1500);
642 assert_eq!(parsed.checksum, 0xCAFEBABE);
643 }
644
645 #[test]
646 fn test_version_detection() {
647 let v1_header = V1Header {
649 compression: BlockCompression::None,
650 original_size: 100,
651 compressed_size: 100,
652 checksum: 0x12345678,
653 };
654 let v1_bytes = v1_header.to_bytes();
655
656 let detected = BlockHeader::from_bytes(&v1_bytes).unwrap();
657 assert_eq!(detected.version(), FormatVersion::V1);
658
659 let v2_header = V2Header {
661 format_version: FormatVersion::V2,
662 compression: BlockCompression::Zstd,
663 flags: BlockFlags::default(),
664 original_size: 200,
665 compressed_size: 150,
666 checksum: 0x87654321,
667 };
668 let v2_bytes = v2_header.to_bytes();
669
670 let detected = BlockHeader::from_bytes(&v2_bytes).unwrap();
671 assert_eq!(detected.version(), FormatVersion::V2);
672 }
673
674 #[test]
675 fn test_v1_to_v2_upgrade() {
676 let v1_header = V1Header {
677 compression: BlockCompression::Lz4,
678 original_size: 500,
679 compressed_size: 300,
680 checksum: 0xABCDEF00,
681 };
682
683 let v2_header = v1_header.upgrade_to_v2();
684
685 assert_eq!(v2_header.format_version, FormatVersion::V2);
686 assert_eq!(v2_header.compression, BlockCompression::Lz4);
687 assert_eq!(v2_header.original_size, 500);
688 assert_eq!(v2_header.compressed_size, 300);
689 assert_eq!(v2_header.checksum, 0xABCDEF00);
690 assert!(!v2_header.flags.encrypted);
692 }
693
694 #[test]
695 fn test_block_migration() {
696 let data = b"Hello, SochDB!";
698 let checksum = crc32fast::hash(data);
699
700 let v1_header = V1Header {
701 compression: BlockCompression::None,
702 original_size: data.len() as u32,
703 compressed_size: data.len() as u32,
704 checksum,
705 };
706
707 let mut buf = v1_header.to_bytes().to_vec();
708 buf.extend_from_slice(data);
709
710 let mut block = MigratableBlock::from_bytes(&buf).unwrap();
712 assert!(block.needs_migration());
713 assert_eq!(block.header.version(), FormatVersion::V1);
714
715 block.migrate();
716 assert!(!block.needs_migration());
717 assert_eq!(block.header.version(), FormatVersion::V2);
718
719 assert_eq!(block.data, data);
721 block.verify_checksum().unwrap();
722 }
723
724 #[test]
725 fn test_format_migrator() {
726 let data1 = b"Block one data";
727 let data2 = b"Block two data";
728
729 let mut blocks: Vec<MigratableBlock> = vec![
731 MigratableBlock {
732 header: BlockHeader::V1(V1Header {
733 compression: BlockCompression::None,
734 original_size: data1.len() as u32,
735 compressed_size: data1.len() as u32,
736 checksum: crc32fast::hash(data1),
737 }),
738 data: data1.to_vec(),
739 },
740 MigratableBlock {
741 header: BlockHeader::V1(V1Header {
742 compression: BlockCompression::None,
743 original_size: data2.len() as u32,
744 compressed_size: data2.len() as u32,
745 checksum: crc32fast::hash(data2),
746 }),
747 data: data2.to_vec(),
748 },
749 ];
750
751 let mut migrator = FormatMigrator::new();
752 let migrated = migrator.migrate_blocks(&mut blocks).unwrap();
753
754 assert_eq!(migrated, 2);
755 assert_eq!(migrator.stats().blocks_read, 2);
756 assert_eq!(migrator.stats().blocks_migrated, 2);
757 assert_eq!(migrator.stats().v1_blocks_found, 2);
758
759 for block in &blocks {
760 assert_eq!(block.header.version(), FormatVersion::V2);
761 }
762 }
763
764 #[test]
765 fn test_checksum_verification_failure() {
766 let data = b"Test data";
767 let block = MigratableBlock {
768 header: BlockHeader::V1(V1Header {
769 compression: BlockCompression::None,
770 original_size: data.len() as u32,
771 compressed_size: data.len() as u32,
772 checksum: 0xBADBAD, }),
774 data: data.to_vec(),
775 };
776
777 let result = block.verify_checksum();
778 assert!(result.is_err());
779 }
780
781 #[test]
782 fn test_block_flags() {
783 let flags = BlockFlags {
784 encrypted: true,
785 extended_checksum: true,
786 spanning: false,
787 has_metadata: true,
788 };
789
790 let byte = flags.to_byte();
791 let parsed = BlockFlags::from_byte(byte);
792
793 assert!(parsed.encrypted);
794 assert!(parsed.extended_checksum);
795 assert!(!parsed.spanning);
796 assert!(parsed.has_metadata);
797 }
798
799 #[test]
800 fn test_migration_progress() {
801 let mut stats = MigrationStats::default();
802
803 assert_eq!(stats.migration_progress(), 100.0);
805
806 stats.v1_blocks_found = 10;
808 stats.blocks_migrated = 5;
809 assert_eq!(stats.migration_progress(), 50.0);
810
811 stats.blocks_migrated = 10;
812 assert_eq!(stats.migration_progress(), 100.0);
813 }
814
815 #[test]
816 fn test_block_complete_roundtrip() {
817 let data = b"Complete block test with some data";
818 let checksum = crc32fast::hash(data);
819
820 let block = MigratableBlock::new(
821 data.to_vec(),
822 BlockCompression::None,
823 data.len() as u32,
824 data.len() as u32,
825 checksum,
826 );
827
828 let bytes = block.to_bytes();
830
831 let parsed = MigratableBlock::from_bytes(&bytes).unwrap();
833
834 assert_eq!(parsed.header.version(), FormatVersion::V2);
835 assert_eq!(parsed.data, data);
836 parsed.verify_checksum().unwrap();
837 }
838
839 #[test]
840 fn test_block_with_flags() {
841 let data = b"Encrypted data";
842 let checksum = crc32fast::hash(data);
843
844 let block = MigratableBlock::new(
845 data.to_vec(),
846 BlockCompression::Zstd,
847 data.len() as u32,
848 data.len() as u32,
849 checksum,
850 )
851 .with_flags(BlockFlags {
852 encrypted: true,
853 extended_checksum: false,
854 spanning: false,
855 has_metadata: true,
856 });
857
858 let bytes = block.to_bytes();
859 let parsed = MigratableBlock::from_bytes(&bytes).unwrap();
860
861 assert!(parsed.header.flags().encrypted);
862 assert!(parsed.header.flags().has_metadata);
863 assert!(!parsed.header.flags().spanning);
864 }
865
866 #[test]
867 fn test_unknown_magic_error() {
868 let bad_magic = b"XXXX\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00";
869 let result = BlockHeader::from_bytes(bad_magic);
870 assert!(result.is_err());
871 }
872
873 #[test]
874 fn test_buffer_too_short_error() {
875 let short_buf = b"TBL";
876 let result = BlockHeader::from_bytes(short_buf);
877 assert!(result.is_err());
878 }
879}