diff --git a/lib/net/ssh/transport/packet_stream.rb b/lib/net/ssh/transport/packet_stream.rb index 407a02ac7..750ef2521 100644 --- a/lib/net/ssh/transport/packet_stream.rb +++ b/lib/net/ssh/transport/packet_stream.rb @@ -169,8 +169,8 @@ def if_needs_rekey? # the states and generally prepares the object for use as a packet stream. def initialize_ssh @hints = {} - @server = State.new(self) - @client = State.new(self) + @server = State.new(self, :server) + @client = State.new(self, :client) @packet = nil initialize_buffered_io end diff --git a/lib/net/ssh/transport/state.rb b/lib/net/ssh/transport/state.rb index 319dfefe2..eec0e6f1e 100644 --- a/lib/net/ssh/transport/state.rb +++ b/lib/net/ssh/transport/state.rb @@ -31,6 +31,12 @@ class State # The number of data blocks processed since the last call to #reset! attr_reader :blocks + # The cipher algorithm in use for this socket endpoint. + attr_reader :cipher + + # The role that this state plays (either :client or :server) + attr_reader :role + # The maximum number of packets that this endpoint wants to process before # needing a rekey. attr_accessor :max_packets @@ -45,15 +51,15 @@ class State # Creates a new state object, belonging to the given socket. Initializes # the algorithms to "none". - def initialize(socket) + def initialize(socket, role) @socket = socket + @role = role @sequence_number = @packets = @blocks = 0 @cipher = CipherFactory.get("none") @hmac = HMAC.get("none") @compression = nil @compressor = @decompressor = nil - @next_iv = nil - @cipher_needs_reset = false + @next_iv = "" end # A convenience method for quickly setting multiple values in a single @@ -65,25 +71,16 @@ def set(values) reset! end - # The cipher algorithm in use for this socket endpoint. - def cipher - if @cipher_needs_reset - @cipher.reset - @cipher.iv = @next_iv - @cipher_needs_reset = false - end - - @cipher - end - def update_cipher(data) - @next_iv = data[-cipher.iv_len..-1] - cipher.update(data) + result = cipher.update(data) + update_next_iv(role == :client ? result : data) + return result end def final_cipher - @cipher_needs_reset = true - cipher.final + result = cipher.final + update_next_iv(role == :client ? result : "", true) + return result end # Increments the counters. The sequence number is incremented (and remapped @@ -185,6 +182,20 @@ def needs_rekey? max_packets && packets > max_packets || max_blocks && blocks > max_blocks end + + private + + def update_next_iv(data, reset=false) + @next_iv << data + @next_iv = @next_iv[-cipher.iv_len..-1] + + if reset + cipher.reset + cipher.iv = @next_iv + end + + return data + end end end; end; end diff --git a/test/transport/test_state.rb b/test/transport/test_state.rb index 17038d0ac..87fad5cfe 100644 --- a/test/transport/test_state.rb +++ b/test/transport/test_state.rb @@ -166,7 +166,7 @@ def socket end def state - @state ||= Net::SSH::Transport::State.new(socket) + @state ||= Net::SSH::Transport::State.new(socket, :test) end end