handshake_state.rb in noise-ruby-0.8.4

- old
+ new

@@ -16,27 +16,24 @@
     #
     # message_patterns: A sequence of message patterns.
     #     Each message pattern is a sequence of tokens from the set ("e", "s", "ee", "es", "se", "ss").
     class HandshakeState
       attr_reader :message_patterns, :symmetric_state
-      attr_reader :s, :rs
+      attr_reader :s, :rs, :e, :re
 
-      def initialize(connection, protocol, initiator, prologue, keypairs)
+      def initialize(connection, protocol, initiator, prologue, local_keypairs, remote_keys)
         @connection = connection
         @protocol = protocol
         @symmetric_state = SymmetricState.new
         @symmetric_state.initialize_symmetric(@protocol, connection)
         @symmetric_state.mix_hash(prologue)
         @initiator = initiator
-        @s = keypairs[:s]
-        @e = keypairs[:e]
-        @rs = keypairs[:rs]
-        @re = keypairs[:re]
+        @s = local_keypairs[:s]
+        @e = local_keypairs[:e]
+        @rs = remote_keys[:rs]
+        @re = remote_keys[:re]
 
-        # TODO : Calls MixHash() once for each public key listed in the pre-messages from  handshake_pattern, with the
-        # specified public key as input (see Section 7 for an explanation of pre-messages). If both initiator and
-        # responder have pre-messages, the initiator's public keys are hashed first.
         get_local_keypair = ->(token) { instance_variable_get('@' + token).public_key }
         get_remote_keypair = ->(token) { instance_variable_get('@r' + token) }
 
         if initiator
           initiator_keypair_getter = get_local_keypair
@@ -44,20 +41,48 @@
         else
           initiator_keypair_getter = get_remote_keypair
           responder_keypair_getter = get_local_keypair
         end
 
+        # Sets message_patterns to the message patterns from handshake_pattern
+        @message_patterns = @protocol.pattern.tokens.dup
+
         @protocol.pattern.initiator_pre_messages&.map do |message|
           keypair = initiator_keypair_getter.call(message)
           @symmetric_state.mix_hash(keypair)
         end
 
+        if @protocol.pattern.fallback
+          message = @message_patterns.delete_at(0).first
+          public_key = initiator_keypair_getter.call(message)
+          @symmetric_state.mix_hash(public_key)
+        end
+
         @protocol.pattern.responder_pre_messages&.map do |message|
           keypair = responder_keypair_getter.call(message)
           @symmetric_state.mix_hash(keypair)
         end
-        # Sets message_patterns to the message patterns from handshake_pattern
-        @message_patterns = @protocol.pattern.tokens.dup
+      end
+
+      def expected_message_length(payload_size)
+        has_key = @symmetric_state.cipher_state.key?
+        pattern = @message_patterns.first
+        len = pattern.inject(0) do |len, token|
+          case token
+          when 'e'
+            len += @protocol.dh_fn.dhlen
+            has_key = true if @protocol.psk_handshake?
+          when 's'
+            len += @protocol.dh_fn.dhlen
+            len += 16 if has_key
+          when 'ee', 'es', 'se', 'ss', 'psk'
+            has_key = true
+          end
+          len
+        end
+        len += payload_size
+        len += 16 if has_key
+        len
       end
 
       # Takes a payload byte sequence which may be zero-length, and a message_buffer to write the output into
       def write_message(payload, message_buffer)
         pattern = @message_patterns.shift