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

# -*- 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