Fixed server dh key gen

4 files changed, 23 insertions, 16 deletions

diff --git a/lib/net/ssh/transport/algorithms.rb b/lib/net/ssh/transport/algorithms.rb
index 9d270af..981e5e6 100644
--- a/lib/net/ssh/transport/algorithms.rb
+++ b/lib/net/ssh/transport/algorithms.rb
@@ -373,8 +373,7 @@ module Net; module SSH; module Transport
             :server_algorithm_packet => @server_packet,
             :client_algorithm_packet => @client_packet,
             :server_side => true,
-            :server_keys => options[:server_keys],
-            :server_dh => options[:server_dh]
+            :server_keys => options[:server_keys]
           }
         else
           debug { "cli.server_packet: #{@server_packet}"}
diff --git a/lib/net/ssh/transport/kex/diffie_hellman_group14_sha1.rb b/lib/net/ssh/transport/kex/diffie_hellman_group14_sha1.rb
index 28a8553..9a91eb8 100644
--- a/lib/net/ssh/transport/kex/diffie_hellman_group14_sha1.rb
+++ b/lib/net/ssh/transport/kex/diffie_hellman_group14_sha1.rb
@@ -31,6 +31,12 @@ module Net; module SSH; module Transport; module Kex
     # The group constant
     G = 2
 
+    def self.dh
+      ret = OpenSSL::PKey::DH.new
+      ret.p, ret.g = OpenSSL::BN.new(P_s, P_r), G
+      ret
+    end
+
     private
 
     def get_p
diff --git a/lib/net/ssh/transport/kex/diffie_hellman_group_exchange_sha1.rb b/lib/net/ssh/transport/kex/diffie_hellman_group_exchange_sha1.rb
index f763fc1..e7fc8cd 100644
--- a/lib/net/ssh/transport/kex/diffie_hellman_group_exchange_sha1.rb
+++ b/lib/net/ssh/transport/kex/diffie_hellman_group_exchange_sha1.rb
@@ -33,6 +33,18 @@ module Net::SSH::Transport::Kex
         data[:need_bytes] = need_bits / 8
       end
 
+      def choose_dh min_bits, need_bits, max_bits
+        DiffieHellmanGroup14SHA1.dh
+      end
+
+      def dh_gen_key dh, need_bits
+        pbits = dh.p.num_bits
+        length = [need_bits * 2, pbits - 1].min
+        dh.priv_key = OpenSSL::BN.rand(length)
+        dh.generate_key!
+        raise unless dh.valid?
+      end
+
       def read_and_handle_get_request
         buffer = connection.next_message
         if buffer.type == KEXDH_GEX_REQUEST
@@ -47,18 +59,13 @@ module Net::SSH::Transport::Kex
           need_bits = [max_bits,need_bits].min
           @data[:need_bits] = need_bits
 
-          dh = data[:server_dh][need_bits]
-          if dh.nil?
-            puts "Generating DH #{need_bits}"
-            debug {"Generating DH"}
-            dh = OpenSSL::PKey::DH.new(need_bits)
-            puts "Generated DH"
-            debug {"Generated DH"}
-            puts "Sending KEXDH_GEX_GROUP"
-          end
+          dh = choose_dh min_bits, need_bits, max_bits
 
           buffer = Net::SSH::Buffer.from(:byte,KEXDH_GEX_GROUP, :bignum, dh.p ,:bignum, dh.g)
           connection.send_message(buffer)
+
+          dh_gen_key dh, need_bits # TODO is need_bits good
+
           return dh
         else
           raise Net::SSH::Exception, "expected KEXDH_GEX_REQUEST, got #{buffer.type}"
diff --git a/test_server.rb b/test_server.rb
index af24980..012960a 100644
--- a/test_server.rb
+++ b/test_server.rb
@@ -18,10 +18,6 @@ puts "Setting up server keys..."
 server_keys = Net::SSH::Server::Keys.new(logger: logger, server_keys_directory: '.')
 server_keys.load_or_generate
 
-puts "Precomputing dh keys..."
-key_sizes = [1024]
-server_dhs = Hash[key_sizes.map {|i| [i,OpenSSL::PKey::DH.new(i)]}]
-
 puts "Listening on port #{PORT}..."
 Thread.start do
   server = TCPServer.new PORT
@@ -35,7 +31,6 @@ Thread.start do
       options[:host_key] = server_keys.types
       options[:kex] = ['diffie-hellman-group-exchange-sha256']
       options[:hmac] = ['hmac-md5']
-      options[:server_dh] = server_dhs
       session = Net::SSH::Transport::ServerSession.new(client,options)
       session.run_loop do |connection|
         connection.on_open_channel('session') do |session, channel, packet|