metasploit-framework/lib/rex/io/ring_buffer.rb

370 lines
7.9 KiB
Ruby

# -*- coding: binary -*-
#
# This class implements a ring buffer with "cursors" in the form of sequence numbers.
# To use this class, pass in a file descriptor and a ring size, the class will read
# data from the file descriptor and store it in the ring. If the ring becomes full,
# the oldest item will be overwritten. To emulate a stream interface, call read_data
# to grab the last sequence number and any buffered data, call read_data again,
# passing in the sequence number and all data newer than that sequence will be
# returned, along with a new sequence to read from.
#
require 'rex/socket'
module Rex
module IO
class RingBuffer
attr_accessor :queue # The data queue, essentially an array of two-element arrays, containing a sequence and data buffer
attr_accessor :seq # The next available sequence number
attr_accessor :fd # The associated socket or IO object for this ring buffer
attr_accessor :size # The number of available slots in the queue
attr_accessor :mutex # The mutex locking access to the queue
attr_accessor :beg # The index of the earliest data fragment in the ring
attr_accessor :cur # The sequence number of the earliest data fragment in the ring
attr_accessor :monitor # The thread handle of the built-in monitor when used
attr_accessor :monitor_thread_error # :nodoc: #
#
# Create a new ring buffer
#
def initialize(socket, opts={})
self.size = opts[:size] || (1024 * 4)
self.fd = socket
self.seq = 0
self.beg = 0
self.cur = 0
self.queue = Array.new( self.size )
self.mutex = Mutex.new
end
def inspect
"#<Rex::IO::RingBuffer @size=#{size} @fd=#{fd} @seq=#{seq} @beg=#{beg} @cur=#{cur}>"
end
#
# Start the built-in monitor, not called when used in a larger framework
#
def start_monitor
self.monitor = monitor_thread if not self.monitor
end
#
# Stop the built-in monitor
#
def stop_monitor
self.monitor.kill if self.monitor
self.monitor = nil
end
#
# The built-in monitor thread (normally unused with Metasploit)
#
def monitor_thread
Thread.new do
begin
while self.fd
buff = self.fd.get_once(-1, 1.0)
next if not buff
store_data(buff)
end
rescue ::Exception => e
self.monitor_thread_error = e
end
end
end
#
# Push data back into the associated stream socket. Logging must occur
# elsewhere, this function is simply a passthrough.
#
def put(data, opts={})
self.fd.put(data, opts={})
end
#
# The clear_data method wipes the ring buffer
#
def clear_data
self.mutex.synchronize do
self.seq = 0
self.beg = 0
self.cur = 0
self.queue = Array.new( self.size )
end
end
#
# The store_data method is used to insert data into the ring buffer.
#
def store_data(data)
self.mutex.synchronize do
# self.cur points to the array index of queue containing the last item
# adding data will result in cur + 1 being used to store said data
# if cur is larger than size - 1, it will wrap back around. If cur
# is *smaller* beg, beg is increemnted to cur + 1 (and wrapped if
# necessary
loc = 0
if self.seq > 0
loc = ( self.cur + 1 ) % self.size
if loc <= self.beg
self.beg = (self.beg + 1) % self.size
end
end
self.queue[loc] = [self.seq += 1, data]
self.cur = loc
end
end
#
# The read_data method returns a two element array with the new reader cursor (a sequence number)
# and the returned data buffer (if any). A result of nil/nil indicates that no data is available
#
def read_data(ptr=nil)
self.mutex.synchronize do
# Verify that there is data in the queue
return [nil,nil] if not self.queue[self.beg]
# Configure the beginning read pointer (sequence number, not index)
ptr ||= self.queue[self.beg][0]
return [nil,nil] if not ptr
# If the pointer is below our baseline, we lost some data, so jump forward
if ptr < self.queue[self.beg][0]
ptr = self.queue[self.beg][0]
end
# Calculate how many blocks exist between the current sequence number
# and the requested pointer, this becomes the number of blocks we will
# need to read to satisfy the result. Due to the mutex block, we do
# not need to scan to find the sequence of the starting block or
# check the sequence of the ending block.
dis = self.seq - ptr
# If the requested sequnce number is less than our base pointer, it means
# that no new data is available and we should return empty.
return [nil,nil] if dis < 0
# Calculate the beginning block index and number of blocks to read
off = ptr - self.queue[self.beg][0]
set = (self.beg + off) % self.size
# Build the buffer by reading forward by the number of blocks needed
# and return the last read sequence number, plus one, as the new read
# pointer.
buff = ""
cnt = 0
lst = ptr
ptr.upto(self.seq) do |i|
block = self.queue[ (set + cnt) % self.size ]
lst,data = block[0],block[1]
buff += data
cnt += 1
end
return [lst + 1, buff]
end
end
#
# The base_sequence method returns the earliest sequence number in the queue. This is zero until
# all slots are filled and the ring rotates.
#
def base_sequence
self.mutex.synchronize do
return 0 if not self.queue[self.beg]
return self.queue[self.beg][0]
end
end
#
# The last_sequence method returns the "next" sequence number where new data will be
# available.
#
def last_sequence
self.seq
end
#
# The create_steam method assigns a IO::Socket compatible object to the ringer buffer
#
def create_stream
Stream.new(self)
end
#
# The select method returns when there is a chance of new data
# XXX: This is mostly useless and requires a rewrite to use a
# real select or notify mechanism
#
def select
::IO.select([ self.fd ], nil, [ self.fd ], 0.10)
end
#
# The wait method blocks until new data is available
#
def wait(seq)
nseq = nil
while not nseq
nseq,data = read_data(seq)
select
end
end
#
# The wait_for method blocks until new data is available or the timeout is reached
#
def wait_for(seq,timeout=1)
begin
::Timeout.timeout(timeout) do
wait(seq)
end
rescue ::Timeout::Error
end
end
#
# This class provides a backwards compatible "stream" socket that uses
# the parents ring buffer.
#
class Stream
attr_accessor :ring
attr_accessor :seq
attr_accessor :buff
def initialize(ring)
self.ring = ring
self.seq = ring.base_sequence
self.buff = ''
end
def read(len=nil)
if len and self.buff.length >= len
data = self.buff.slice!(0,len)
return data
end
while true
lseq, data = self.ring.read_data( self.seq )
return if not lseq
self.seq = lseq
self.buff << data
if len
if self.buff.length >= len
return self.buff.slice!(0,len)
else
IO.select(nil, nil, nil, 0.25)
next
end
end
data = self.buff
self.buff = ''
return data
# Not reached
break
end
end
def write(data)
self.ring.write(data)
end
end
end
end
end
=begin
server = Rex::Socket.create_tcp_server('LocalPort' => 0)
lport = server.getsockname[2]
client = Rex::Socket.create_tcp('PeerHost' => '127.0.0.1', 'PeerPort' => lport)
conn = server.accept
r = Rex::IO::RingBuffer.new(conn, {:size => 1024*1024})
client.put("1")
client.put("2")
client.put("3")
s,d = r.read_data
client.put("4")
client.put("5")
client.put("6")
s,d = r.read_data(s)
client.put("7")
client.put("8")
client.put("9")
s,d = r.read_data(s)
client.put("0")
s,d = r.read_data(s)
test_counter = 11
1.upto(100) do
client.put( "X" )
test_counter += 1
end
sleep(1)
s,d = r.read_data
p s
p d
fdata = ''
File.open("/bin/ls", "rb") do |fd|
fdata = fd.read(fd.stat.size)
fdata = fdata * 10
client.put(fdata)
end
sleep(1)
s,vdata = r.read_data(s)
if vdata != fdata
puts "DATA FAILED"
else
puts "DATA VERIFIED"
end
r.clear_data
a = r.create_stream
b = r.create_stream
client.put("ABC123")
sleep(1)
p a.read
p b.read
client.put("$$$$$$")
sleep(1)
p a.read
p b.read
c = r.create_stream
p c.read
=end