Save

src/natpmpcrystal.cr

@@ -11,30 +11,62 @@ module NatPMP
     UNSUPPORTED_OPCODE        = 5
   end
 
-  struct MappingPacket
+  enum OP : UInt8
-    enum OP : UInt8
+    UDP = 1_u8
-      UDP = 1_u8
+    TCP = 2_u8
-      TCP = 2_u8
-    end
-
-    property vers : UInt8 = 0_u8
-    property op : UInt8
-    property reserved : UInt16 = 0_u16
-    property internal_port : UInt16
-    property external_port : UInt16
-    property lifetime : UInt32
-
-    def initialize(@op = UDP, @internal_port = nil, @external_port = nil, @lifetime = 0)
-    end
   end
 
-  # private record MappingPacket,
+  struct MappingPacket
-  #   vers : UInt8 = 0,
+    @vers : UInt8 = 0_u8
-  #   op : UInt8 = 0,
+    @op : UInt8 = OP::UDP.value
-  #   reserved : UInt16 = 0,
+    @reserved : UInt16 = 0_u16
-  #   internal_port : UInt16 = 0,
+    @internal_port : UInt16
-  #   external_port : UInt16 = 0,
+    @external_port : UInt16
-  #   lifetime : UInt32 = 0,
+    @lifetime : UInt32 = 0_u32
+
+    def initialize(@internal_port, @external_port, @lifetime = 0, @op = UDP)
+      unless [1, 2].include?(@op)
+        raise ArgumentError, "Operation should be either '1_u8' for UDP or '2_u8' for TCP (default: UDP)"
+      end
+    end
+
+    def initialize(@internal_port, @external_port)
+    end
+
+    def to_io
+      io = IO::Memory.new(12)
+      io.write_bytes(@vers, IO::ByteFormat::BigEndian)
+      io.write_bytes(@op, IO::ByteFormat::BigEndian)
+      io.write_bytes(@reserved, IO::ByteFormat::BigEndian)
+      io.write_bytes(@internal_port, IO::ByteFormat::BigEndian)
+      io.write_bytes(@external_port, IO::ByteFormat::BigEndian)
+      io.write_bytes(@lifetime, IO::ByteFormat::BigEndian)
+      return io
+    end
+
+    def to_slice
+      slice = uninitialized UInt8[12]
+      IO::ByteFormat::BigEndian.encode(@vers, v = Bytes.new(1))
+      IO::ByteFormat::BigEndian.encode(@op, o = Bytes.new(1))
+      # IO::ByteFormat::BigEndian.encode(@reserved, r = Bytes.new(2))
+      IO::ByteFormat::BigEndian.encode(@internal_port, i = Bytes.new(2))
+      IO::ByteFormat::BigEndian.encode(@external_port, e = Bytes.new(2))
+      IO::ByteFormat::BigEndian.encode(@lifetime, l = Bytes.new(4))
+      slice[0] = v[0]
+      slice[1] = o[0]
+      slice[2] = 0 # RESERVED
+      slice[3] = 0 # RESERVED
+      slice[4] = i[0]
+      slice[5] = i[1]
+      slice[6] = e[0]
+      slice[7] = e[1]
+      slice[8] = l[0]
+      slice[9] = l[1]
+      slice[10] = l[2]
+      slice[11] = l[3]
+      return slice
+    end
+  end
 
   class Client
     @client : UDPSocket
@@ -45,8 +77,7 @@ module NatPMP
       initialize(gateway_ip.path)
     end
 
-    def initialize(gateway_ip : String)
+    def initialize(@gateway_ip : String)
-      @gateway_ip = gateway_ip
       @client = UDPSocket.new
       # A given host may have more than one independent
       # NAT-PMP client running at the same time, and address announcements
@@ -77,9 +108,27 @@ module NatPMP
     def send_public_address_request
       @client.send("\x00\x00")
       msg = Bytes.new(12)
-      @client.receive(msg)
+      @client.read_timeout = 250.milliseconds
-      vers = (msg[0])
+      8.times do |i|
-      op = (msg[1])
+        begin
+          @client.receive(msg)
+          break
+        rescue IO::TimeoutError
+          # If no
+          # NAT-PMP response is received from the gateway after 250 ms, the
+          # client retransmits its request and waits 500 ms.  The client SHOULD
+          # repeat this process with the interval between attempts doubling each
+          # time.
+          @client.read_timeout = @client.read_timeout.not_nil!*2
+          next
+        rescue
+          raise "The gateway '#{@gateway_ip}' does not support NAT-PMP"
+          break
+        end
+      end
+
+      vers = msg[0]
+      op = msg[1]
       result_code = get_result_code(msg[2..3])
       epoch = get_epoch(msg[4..7])
 
@@ -94,17 +143,18 @@ module NatPMP
       return vers, op, result_code, epoch, ip_address
     end
 
-    # def request_mapping(op : Int32, internal_port : Int32, external_port : Int32, lifetime : Int32)
+    def request_mapping(internal_port : UInt16, external_port : UInt16, lifetime : Uint32, operation : UInt8)
-    #   begin
+      request = MappingPacket.new(internal_port, external_port, lifetime, operation)
-    #   request_mapping(op.to_u8, internal_port.to_u16, external_port.to_u16, lifetime.to_u32)
+      msg = Bytes.new(16)
-    #   rescue ex
+      @client.receive(msg)
-    #     puts ex.message
+      vers = msg[0]
-    #     exit(1)
+      op = msg[1]
-    #   end
+      result_code = get_result_code(msg[2..3])
-    # end
+      epoch = get_epoch(msg[4..7])
-
+      internal_port = get_port(msg[8..9])
-    def request_mapping(op : UInt8, internal_port : UInt16, external_port : UInt16, lifetime : UInt32)
+      external_port = get_port(msg[10..11])
-      request = MappingPacket.new(op, internal_port, external_port, lifetime)
+      lifetime = get_lifetime(msg[12..15])
+      return vers, op, result_code, epoch, internal_port, external_port, lifetime
     end
 
     private def get_result_code(msg)
@@ -115,30 +165,50 @@ module NatPMP
 
     # Seconds Since Start of Epoch
     private def get_epoch(msg)
-      # epoch : Int32 = 0
-      # msg.each_with_index do |byte, index|
-      #   epoch |= (byte.to_i << (8 * (msg.size - 1 - index)))
-      # end
-
       # Responses also contain a 32-bit unsigned integer
       # corresponding to the number of seconds since the NAT gateway was
       # rebooted or since its port mapping state was otherwise reset.
       return IO::ByteFormat::BigEndian.decode(UInt32, msg)
     end
 
+    private def get_port(msg)
+      return IO::ByteFormat::BigEndian.decode(UInt16, msg)
+    end
+
     private def get_ip_address(msg)
       "#{msg[0]}.#{msg[1]}.#{msg[2]}.#{msg[3]}"
     end
+
   end
 end
 
-client = NatPMP::Client.new("192.168.1.1")
+# client = NatPMP::Client.new("192.168.0.1")
-pp client.send_public_address_request
+# pp client.send_public_address_request
-xd = client.send_public_address_request_raw
+pp mapping_packet = NatPMP::MappingPacket.new(25555,25555)
-pp xd
 
-request = client.request_mapping(NatPMP::MappingPacket::OP::UDP.value, 25580_u16, 25580_u16, 0_u32)
+Benchmark.ips do |x|
+  x.report("bytes") do
+    mapping_packet.to_io
+  end
+
+  x.report("staticarray") do
+    # pp mapping_packet.to_io.to_slice
+    mapping_packet.to_slice
+  end
+
+  x.report("staticarray to io") do
+    # pp mapping_packet.to_io.to_slice
+    mapping_packet.to_slice
+  end
+end
+
+pp typeof(mapping_packet)
+
+# xd = client.send_public_address_request_raw
+# pp xd
+
+# request = client.request_mapping(25580, 25580)
 
 # request2 = client.request_mapping(1, 255802, 25580, 0)
 
-pp request
+# pp request