diff --git a/src/directory.zig b/src/directory.zig index 0f41d87..9ed1fad 100644 --- a/src/directory.zig +++ b/src/directory.zig @@ -27,6 +27,7 @@ pub const DirEntry = struct { fn init(alloc: std.mem.Allocator, hdr: DirHeader, rdr: io.AnyReader) !DirEntry { const raw = try rdr.readStruct(RawDirEntryStart); const name = try alloc.alloc(u8, raw.name_size + 1); + errdefer alloc.free(name); _ = try rdr.read(name); return .{ .block_start = hdr.inode_block_start, diff --git a/src/file.zig b/src/file.zig index 439e9b7..696ddda 100644 --- a/src/file.zig +++ b/src/file.zig @@ -113,6 +113,7 @@ pub const File = struct { } try self.readDirEntries(rdr); var files = try rdr.alloc.alloc(File, self.dirEntries.?.count()); + errdefer rdr.alloc.free(files); var dirEntryIter = self.dirEntries.?.valueIterator(); var i: u32 = 0; while (dirEntryIter.next()) |ent| : (i += 1) { @@ -208,7 +209,7 @@ pub const File = struct { pub fn extract(self: *File, rdr: *Reader, config: ExtractConfig, path: []const u8) (ExtractError || anyerror)!void { var pol: std.Thread.Pool = undefined; try pol.init(.{ - .allocator = std.heap.smp_allocator, + .allocator = rdr.alloc, .n_jobs = config.thread_count, }); defer pol.deinit(); @@ -233,22 +234,22 @@ pub const File = struct { } var iter = try self.iterator(rdr); defer iter.deinit(); - while (iter.next()) |*f| { + while (iter.next()) |f| { const extr_path = try std.mem.concat(rdr.alloc, u8, &[3][]const u8{ real_path, "/", f.name }); defer rdr.alloc.free(extr_path); - try @constCast(f).extractReal(rdr, config, pool, extr_path, false); + try f.extractReal(rdr, config, pool, extr_path, false); } }, .file, .ext_file => { if ((!first and exists) or (first and exists and stat.?.kind != .directory)) return ExtractError.FileExists; - const extr_path = if (first and exists and stat.?.kind == .directory) blk: { - break :blk try std.mem.concat(rdr.alloc, u8, &[3][]const u8{ real_path, "/", self.name }); - } else blk: { - const tmp = try rdr.alloc.alloc(u8, real_path.len); - @memcpy(tmp, real_path); - break :blk tmp; - }; + var extr_path: []u8 = undefined; + if (first and exists and stat.?.kind == .directory) { + extr_path = try std.mem.concat(rdr.alloc, u8, &[3][]const u8{ real_path, "/", self.name }); + } else { + extr_path = try rdr.alloc.alloc(u8, real_path.len); + @memcpy(extr_path, real_path); + } defer rdr.alloc.free(extr_path); var ext = try self.extractor(rdr); defer ext.deinit(); @@ -289,10 +290,10 @@ const FileIterator = struct { curIndex: u32 = 0, - pub fn next(self: *FileIterator) ?File { + pub fn next(self: *FileIterator) ?*File { if (self.curIndex >= self.files.len) return null; defer self.curIndex += 1; - return self.files[self.curIndex]; + return &self.files[self.curIndex]; } pub fn reset(self: *FileIterator) void { self.curIndex = 0; diff --git a/src/reader.zig b/src/reader.zig index 5291a40..98505ae 100644 --- a/src/reader.zig +++ b/src/reader.zig @@ -54,10 +54,17 @@ pub const Reader = struct { return out; } pub fn deinit(self: *Reader) void { + self.frag_table.deinit(self.alloc); + self.export_table.deinit(self.alloc); + self.id_table.deinit(self.alloc); self.root.deinit(self.alloc); self.holder.deinit(); } + pub fn open(self: *Reader, path: []const u8) !File { + return self.root.open(self, path); + } + fn fileFromRef(self: *Reader, ref: inode.InodeRef, name: []const u8) !File { var offset_rdr = self.holder.readerAt(ref.block_start + self.super.inode_table_start); var meta_rdr: MetadataReader = .init( @@ -92,8 +99,40 @@ test "root iter" { test "extract" { const test_sfs_path = "testing/LinuxPATest.sfs"; const extract_path = "testing/testExtract"; - try std.fs.cwd().deleteTree(extract_path); + std.fs.cwd().deleteTree(extract_path) catch |err| { + if (err != std.fs.Dir.DeleteFileError.FileNotFound) { + return err; + } + }; var rdr: Reader = try .init(std.testing.allocator, test_sfs_path, 0); defer rdr.deinit(); try rdr.root.extract(&rdr, try .init(), extract_path); } + +test "extract single file" { + const test_sfs_path = "testing/LinuxPATest.sfs"; + const sfs_file_path = "Start.exe"; + const extract_path = "testing/Start.exe"; + std.fs.cwd().deleteFile(extract_path) catch |err| { + if (err != std.fs.Dir.DeleteFileError.FileNotFound) { + return err; + } + }; + var rdr: Reader = try .init(std.testing.allocator, test_sfs_path, 0); + defer rdr.deinit(); + var fil = try rdr.open(sfs_file_path); + defer fil.deinit(std.testing.allocator); + try fil.extract(&rdr, try .init(), extract_path); +} + +test "extract single directory" { + const test_sfs_path = "testing/LinuxPATest.sfs"; + const sfs_file_path = "Documents"; + const extract_path = "testing/Documents"; + try std.fs.cwd().deleteTree(extract_path); + var rdr: Reader = try .init(std.testing.allocator, test_sfs_path, 0); + defer rdr.deinit(); + var fil = try rdr.open(sfs_file_path); + defer fil.deinit(std.testing.allocator); + try fil.extract(&rdr, try .init(), extract_path); +} diff --git a/src/readers/data_extractor.zig b/src/readers/data_extractor.zig index b204d21..ed92856 100644 --- a/src/readers/data_extractor.zig +++ b/src/readers/data_extractor.zig @@ -79,7 +79,8 @@ pub const DataExtractor = struct { if (self.frag_data != null) self.alloc.free(self.frag_data.?); } - fn processBlockToFile(self: *DataExtractor, errs: *std.ArrayList(anyerror), block_ind: usize, fil: *fs.File) void { + fn processBlockToFile(self: *DataExtractor, wg: *std.Thread.WaitGroup, errs: *std.ArrayList(anyerror), block_ind: usize, fil: *fs.File) void { + defer wg.finish(); const offset_rdr = self.holder.readerAt(self.block_offset[block_ind]); var fil_wrtr: FileOffsetWriter = .init(fil, block_ind * self.block_size); var limit = std.io.limitedReader(offset_rdr, self.sizes[block_ind].size); @@ -94,6 +95,15 @@ pub const DataExtractor = struct { }; } + fn fragmentToFile(self: *DataExtractor, wg: *std.Thread.WaitGroup, errs: *std.ArrayList(anyerror), fil: *fs.File) void { + defer wg.finish(); + fil.pwriteAll(self.frag_data.?, self.block_size * self.sizes.len) catch |err| { + errs.append(err) catch |ignored_err| { + std.debug.print("{}\n", .{ignored_err}); + }; + }; + } + /// Write the data completely to the given file. /// Ignores the file's current offset and writes from the beginning of the file. /// Returns the amount of bytes written. @@ -104,9 +114,15 @@ pub const DataExtractor = struct { var errs: std.ArrayList(anyerror) = .init(self.alloc); defer errs.deinit(); for (0..self.sizes.len) |i| { - pool.spawnWg(&wg, processBlockToFile, .{ self, &errs, i, fil }); + wg.start(); + try pool.spawn(processBlockToFile, .{ self, &wg, &errs, i, fil }); + } + if (self.frag_data != null) { + wg.start(); + try pool.spawn(fragmentToFile, .{ self, &wg, &errs, fil }); } wg.wait(); + //TODO: see if there's any errors } // fn processBlock(self: *DataExtractor, errs: std.ArrayList(anyerror), data_out: std.AutoHashMap([]u8), block_ind: u32) void { diff --git a/src/readers/metadata.zig b/src/readers/metadata.zig index 63978f7..33d87ab 100644 --- a/src/readers/metadata.zig +++ b/src/readers/metadata.zig @@ -64,7 +64,7 @@ pub const MetadataReader = struct { while (cur_read < bytes.len) { if (self.offset >= self.block.len) try self.readNextBlock(); to_read = @min(bytes.len - cur_read, self.block.len - self.offset); - @memcpy(bytes[cur_read..], self.block[self.offset .. @as(usize, self.offset) + to_read]); + @memcpy(bytes[cur_read .. cur_read + to_read], self.block[self.offset .. @as(usize, self.offset) + to_read]); self.offset += @truncate(to_read); cur_read += to_read; } diff --git a/src/table.zig b/src/table.zig index 9987cd3..e3acab4 100644 --- a/src/table.zig +++ b/src/table.zig @@ -26,22 +26,21 @@ pub fn Table( }; } pub fn deinit(self: *Self, alloc: std.mem.Allocator) void { - if (self.table.len == 0) alloc.free(self.table); + if (self.table.len != 0) alloc.free(self.table); } pub fn getValue(self: *Self, read: *Reader, i: u64) !T { if (i >= self.item_count) return TableError.InvalidIndex; if (self.table.len > i) return self.table[i]; - var meta_rdr: MetadataReader = undefined; var offset_rdr: FileOffsetReader = undefined; + var meta_rdr: MetadataReader = undefined; var meta_buf: [8]u8 = [1]u8{0} ** 8; - var meta_offset: u64 = 0; + const meta_offset = std.mem.bytesAsValue(u64, &meta_buf); var to_read: u32 = 0; while (self.table.len <= i) { _ = try read.holder.file.preadAll(&meta_buf, self.offset); self.offset += 8; - meta_offset = std.mem.bytesToValue(u8, &meta_buf); - offset_rdr = read.holder.readerAt(meta_offset); + offset_rdr = read.holder.readerAt(meta_offset.*); meta_rdr = .init(read.alloc, self.decomp, offset_rdr.any()); defer meta_rdr.deinit(); to_read = @min(self.item_count - self.table.len, comptime 8192 / @sizeOf(T)); diff --git a/src/zig_unsquashfs.zig b/src/zig_unsquashfs.zig index 966ab91..ae8fbc7 100644 --- a/src/zig_unsquashfs.zig +++ b/src/zig_unsquashfs.zig @@ -2,6 +2,7 @@ const std = @import("std"); const config = @import("config"); const Reader = @import("reader.zig").Reader; +const ExtractConfig = @import("file.zig").File.ExtractConfig; const stdout = std.io.getStdOut(); @@ -121,26 +122,37 @@ pub fn main() !void { _ = try stdout.writeAll("no extract location given\n"); return; } - var rdr: Reader = .init( + var rdr = Reader.init( alloc.allocator(), filename, offset, ) catch |err| { - try std.fmt.format(stdout.writer(), "Error opening {s} as squashfs: {any}", "\n", .{ filename, err }); + try std.fmt.format(stdout.writer(), "Error opening {s} as squashfs: {any}\n", .{ filename, err }); + return; }; if (list == .None) { + var conf = ExtractConfig.init() catch |err| { + try std.fmt.format(stdout.writer(), "Error getting system info: {any}\n", .{err}); + return; + }; + conf.deref_sym = deref; + conf.unbreak_sym = unbreak; + conf.verbose = verbose; if (extr_files.items.len == 0) { - rdr.root.extract(&rdr, extr_location) catch |err| { - try std.fmt.format(stdout.writer(), "Error extracting archive: {any}", "\n", .{err}); + rdr.root.extract(&rdr, conf, extr_location) catch |err| { + try std.fmt.format(stdout.writer(), "Error extracting archive: {any}\n", .{err}); + return; }; } else { for (extr_files.items) |path| { var fil = rdr.root.open(&rdr, path) catch |err| { - try std.fmt.format(stdout.writer(), "Error extracting {s}: {any}", "\n", .{ path, err }); + try std.fmt.format(stdout.writer(), "Error extracting {s}: {any}\n", .{ path, err }); + return; }; defer fil.deinit(alloc.allocator()); - fil.extract(&rdr, extr_location) catch |err| { - try std.fmt.format(stdout.writer(), "Error extracting {s}: {any}", "\n", .{ path, err }); + fil.extract(&rdr, conf, extr_location) catch |err| { + try std.fmt.format(stdout.writer(), "Error extracting {s}: {any}\n", .{ path, err }); + return; }; } }