Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added SubIO: IO-like object that could be used for copy-free substream access #9

Merged
merged 2 commits into from
Jul 26, 2022
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 108 additions & 2 deletions lib/kaitai/struct/struct.rb
Original file line number Diff line number Diff line change
@@ -97,10 +97,10 @@ def initialize(actual, expected)
def initialize(arg)
if arg.is_a?(String)
@_io = StringIO.new(arg)
elsif arg.is_a?(IO)
elsif arg.is_a?(IO) or arg.is_a?(SubIO)
@_io = arg
else
raise TypeError.new('can be initialized with IO or String only')
raise TypeError.new('can be initialized with IO, SubIO or String only')
end
align_to_byte
end
@@ -516,6 +516,20 @@ def self.process_rotate_left(data, amount, group_size)

# @!endgroup

##
# Reserves next n bytes from current stream as a
# Kaitai::Struct::Stream substream. Substream has its own pointer
# and addressing in the range of [0, n) bytes. This stream's pointer
# is advanced to the position right after this substream.
# @param n [Fixnum] number of bytes to reserve for a substream
# @return [Stream] substream covering n bytes from the current
# position
def substream(n)
sub = Stream.new(SubIO.new(@_io, @_io.pos, n))
@_io.seek(@_io.pos + n)
sub
end

##
# Resolves value using enum: if the value is not found in the map,
# we'll just use literal value per se.
@@ -565,6 +579,98 @@ def self.inspect_values(*args)
end
end

##
# Substream IO implementation: a IO object which wraps existing IO object
# and provides similar byte/bytes reading functionality, but only for a
# limited set of bytes starting from specified offset and spanning up to
# specified length.
class SubIO
attr_reader :parent_io
attr_reader :parent_offset
attr_reader :parent_len
attr_reader :pos

def initialize(parent_io, parent_start, parent_len)
@parent_io = parent_io
@parent_start = parent_start
@parent_len = parent_len
@parent_end = @parent_start + @parent_len
@pos = 0
@closed = false
end

def eof?
raise IOError.new("closed stream") if @closed

@pos >= @parent_len
end

def seek(amount, whence = IO::SEEK_SET)
raise IOError.new("closed stream") if @closed
raise ArgumentError.new("Anything but IO::SEEK_SET is not supported in SubIO::seek") if whence != IO::SEEK_SET
raise TypeError.new("Need an integer argument for amount in SubIO::seek") unless amount.respond_to?(:to_int)
raise Errno::EINVAL.new("Negative position requested") if amount < 0
@pos = amount.to_int
return 0
end

def getc
raise IOError.new("closed stream") if @closed

return nil if @pos >= @parent_len

# remember position in parent IO
old_pos = @parent_io.pos
@parent_io.seek(@parent_start + @pos)
res = @parent_io.getc
@pos += 1

# restore position in parent IO
@parent_io.seek(old_pos)

res
end

def read(len = nil)
raise IOError.new("closed stream") if @closed

# remember position in parent IO
old_pos = @parent_io.pos

# read until the end of substream
if len.nil?
len = @parent_len - @pos
return "" if len < 0
else
# special case to requesting exactly 0 bytes
return "" if len == 0

# cap intent to read if going beyond substream boundary
left = @parent_len - @pos

# if actually requested reading and we're beyond the boundary, return nil
return nil if left <= 0

# otherwise, still return something, but less than requested
len = left if len > left
end

@parent_io.seek(@parent_start + @pos)
res = @parent_io.read(len)
read_len = res.size
@pos += read_len

# restore position in parent IO
@parent_io.seek(old_pos)

res
end

def close
@closed = true
end
end

##
# Common ancestor for all error originating from Kaitai Struct usage.
# Stores KSY source path, pointing to an element supposedly guilty of
229 changes: 229 additions & 0 deletions spec/subio_spec.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,229 @@
require 'kaitai/struct/struct'
require 'stringio'

RSpec.describe Kaitai::Struct::SubIO do
context "in 12345 asking for 234" do
before(:each) do
parent_io = StringIO.new("12345")
@io = Kaitai::Struct::SubIO.new(parent_io, 1, 3)
@normal_io = StringIO.new("234")
end

describe "#seek" do
it "can seek to 0" do
expect(@normal_io.seek(0)).to eq(0)
expect(@io.seek(0)).to eq(0)

expect(@normal_io.pos).to eq(0)
expect(@io.pos).to eq(0)
end

it "can seek to 2" do
expect(@normal_io.seek(2)).to eq(0)
expect(@io.seek(2)).to eq(0)

expect(@normal_io.pos).to eq(2)
expect(@io.pos).to eq(2)
end

it "can seek to 10 (beyond EOF)" do
expect(@normal_io.seek(10)).to eq(0)
expect(@io.seek(10)).to eq(0)

expect(@normal_io.pos).to eq(10)
expect(@io.pos).to eq(10)
end

it "cannot seek to -1" do
expect { @normal_io.seek(-1) }.to raise_error(Errno::EINVAL)
expect { @io.seek(-1) }.to raise_error(Errno::EINVAL)
end

it "cannot seek to \"foo\"" do
expect { @normal_io.seek("foo") }.to raise_error(TypeError)
expect { @io.seek("foo") }.to raise_error(TypeError)
end

it "can seek to 2.3" do
expect(@normal_io.seek(2.3)).to eq(0)
expect(@io.seek(2.3)).to eq(0)

expect(@normal_io.pos).to eq(2)
expect(@io.pos).to eq(2)
end
end

describe "#pos" do
it "returns 0 by default" do
expect(@normal_io.pos).to eq(0)
expect(@io.pos).to eq(0)
end

it "returns 2 after reading 2 bytes" do
@normal_io.read(2)
@io.read(2)

expect(@normal_io.pos).to eq(2)
expect(@io.pos).to eq(2)
end

it "returns 3 after reading 4 bytes" do
@normal_io.read(4)
@io.read(4)

expect(@normal_io.pos).to eq(3)
expect(@io.pos).to eq(3)
end
end

describe "#eof?" do
it "returns false by default" do
expect(@normal_io.eof?).to eq(false)
expect(@io.eof?).to eq(false)
end

it "returns false after reading 2 bytes" do
@normal_io.read(2)
@io.read(2)

expect(@normal_io.eof?).to eq(false)
expect(@io.eof?).to eq(false)
end

it "returns true after reading 3 bytes" do
@normal_io.read(3)
@io.read(3)

expect(@normal_io.eof?).to eq(true)
expect(@io.eof?).to eq(true)
end

it "returns true after reading 4 bytes" do
@normal_io.read(4)
@io.read(4)

expect(@normal_io.eof?).to eq(true)
expect(@io.eof?).to eq(true)
end

it "returns true after seeking at 3 bytes" do
@normal_io.seek(3)
@io.seek(3)

expect(@normal_io.eof?).to eq(true)
expect(@io.eof?).to eq(true)
end

it "returns true after seeking at 10 bytes" do
@normal_io.seek(10)
@io.seek(10)

expect(@normal_io.eof?).to eq(true)
expect(@io.eof?).to eq(true)
end
end

describe "#read" do
it "reads 234 with no arguments" do
expect(@normal_io.read).to eq("234")
expect(@io.read).to eq("234")
end

it "reads 23 when asked to read 2" do
expect(@normal_io.read(2)).to eq("23")
expect(@io.read(2)).to eq("23")
end

it "reads 234 when asked to read 3" do
expect(@normal_io.read(3)).to eq("234")
expect(@io.read(3)).to eq("234")
end

it "reads 234 when asked to read 4" do
expect(@normal_io.read(4)).to eq("234")
expect(@io.read(4)).to eq("234")
end

it "reads 234 when asked to read 10" do
expect(@normal_io.read(10)).to eq("234")
expect(@io.read(10)).to eq("234")
end

it "reads 234 + empty when asked to read + read" do
expect(@normal_io.read).to eq("234")
expect(@io.read).to eq("234")

expect(@normal_io.read).to eq("")
expect(@io.read).to eq("")
end

it "reads 2 + 34 when asked to read(1) + read" do
expect(@normal_io.read(1)).to eq("2")
expect(@io.read(1)).to eq("2")

expect(@normal_io.read).to eq("34")
expect(@io.read).to eq("34")
end

it "reads 2 + 34 when asked to read(1) + read(2)" do
expect(@normal_io.read(1)).to eq("2")
expect(@io.read(1)).to eq("2")

expect(@normal_io.read(2)).to eq("34")
expect(@io.read(2)).to eq("34")
end

it "reads 2 + 34 when asked to read(1) + read(10)" do
expect(@normal_io.read(1)).to eq("2")
expect(@io.read(1)).to eq("2")

expect(@normal_io.read(10)).to eq("34")
expect(@io.read(10)).to eq("34")
end

context("after seek to EOF") do
before(:each) do
@normal_io.seek(3)
@io.seek(3)
end

it "reads nil when asked to read(1)" do
expect(@normal_io.read(1)).to eq(nil)
expect(@io.read(1)).to eq(nil)
end

it "reads empty when asked to read()" do
expect(@normal_io.read).to eq("")
expect(@io.read).to eq("")
end

it "reads empty when asked to read(0)" do
expect(@normal_io.read(0)).to eq("")
expect(@io.read(0)).to eq("")
end
end

context("after seek beyond EOF") do
before(:each) do
@normal_io.seek(10)
@io.seek(10)
end

it "reads nil when asked to read(1)" do
expect(@normal_io.read(1)).to eq(nil)
expect(@io.read(1)).to eq(nil)
end

it "reads empty when asked to read()" do
expect(@normal_io.read).to eq("")
expect(@io.read).to eq("")
end

it "reads empty when asked to read(0)" do
expect(@normal_io.read(0)).to eq("")
expect(@io.read(0)).to eq("")
end
end
end
end
end