From 11898d0f413629a9de3adb6d52f7744f6bef07c7 Mon Sep 17 00:00:00 2001
From: ChillerDragon <chillerdragon@gmail.com>
Date: Sat, 12 Nov 2022 12:51:09 +0100
Subject: [PATCH] Ensure tokens being in correct format

---
 lib/net_base.rb         | 10 +++++++++-
 lib/teeworlds_client.rb |  4 ++--
 lib/teeworlds_server.rb | 30 +++++++++++++-----------------
 lib/token.rb            | 20 ++++++++++++++++++++
 4 files changed, 44 insertions(+), 20 deletions(-)
 create mode 100644 lib/token.rb

diff --git a/lib/net_base.rb b/lib/net_base.rb
index f1ad6d5..3e98cb5 100644
--- a/lib/net_base.rb
+++ b/lib/net_base.rb
@@ -1,12 +1,15 @@
 # frozen_string_literal: true
 
+require_relative 'token'
+
 ##
 # NetBase
 #
 # Lowest network layer logic. Sends packets via udp.
 # Also adding the teeworlds protocol packet header.
 class NetBase
-  attr_accessor :peer_token, :ack
+  attr_accessor :ack
+  attr_reader :peer_token
 
   def initialize(opts = {})
     @verbose = opts[:verbose] || false
@@ -28,6 +31,11 @@ class NetBase
     @ack = 0
   end
 
+  def set_peer_token(token)
+    SecurityToken.validate(token)
+    @peer_token = token
+  end
+
   ##
   # Sends a packing setting the proper header for you
   #
diff --git a/lib/teeworlds_client.rb b/lib/teeworlds_client.rb
index 21dcb28..32209f1 100644
--- a/lib/teeworlds_client.rb
+++ b/lib/teeworlds_client.rb
@@ -281,7 +281,7 @@ class TeeworldsClient
     data += Packer.pack_int(inp[:wanted_weapon])
     data += Packer.pack_int(inp[:next_weapon])
     data += Packer.pack_int(inp[:prev_weapon])
-    msg = NetChunk.createl_header(vital: false, size: data.size + 1) +
+    msg = NetChunk.create_header(vital: false, size: data.size + 1) +
           [pack_msg_id(NETMSG_INPUT, system: true)] +
           data
     @netbase.send_packet(msg, 1)
@@ -289,7 +289,7 @@ class TeeworldsClient
 
   def on_msg_token(data)
     @token = bytes_to_str(data)
-    @netbase.peer_token = @token
+    @netbase.set_peer_token(@token)
     puts "Got token #{@token}"
     send_msg_connect
   end
diff --git a/lib/teeworlds_server.rb b/lib/teeworlds_server.rb
index c62f838..704a512 100644
--- a/lib/teeworlds_server.rb
+++ b/lib/teeworlds_server.rb
@@ -13,6 +13,7 @@ require_relative 'net_addr'
 require_relative 'packer'
 require_relative 'game_server'
 require_relative 'message'
+require_relative 'token'
 
 class Client
   attr_accessor :id, :addr, :vital_sent, :last_recv_time, :token
@@ -21,14 +22,9 @@ class Client
     @id = attr[:id]
     @addr = attr[:addr]
     @vital_sent = 0
-    @token = attr[:token]
-    unless @token.size == 4
-      raise "Invalid client token size\n" \
-            "got=#{@token.size} expected=4\n" \
-            "#{str_hex(@token)}"
-    end
-
     @last_recv_time = Time.now
+    @token = attr[:token]
+    SecurityToken.validate(@token)
   end
 
   # TODO: use or remove
@@ -145,16 +141,16 @@ class TeeworldsServer
   def send_ctrl_close(client, reason)
     msg = [NET_CTRLMSG_CLOSE]
     msg += Packer.pack_str(reason) unless reason.nil?
-    @netbase.peer_token = client.token
+    @netbase.set_peer_token(client.token)
     @netbase.send_packet(msg, 0, control: true, addr: client.addr)
-    # @netbase.peer_token = @server_token
+    # @netbase.set_peer_token(@server_token)
   end
 
   def send_ctrl_with_token(addr, token)
     msg = [NET_CTRLMSG_TOKEN] + str_bytes(@server_token)
-    @netbase.peer_token = token
+    @netbase.set_peer_token(token)
     @netbase.send_packet(msg, 0, control: true, addr:)
-    # @netbase.peer_token = @server_token
+    # @netbase.set_peer_token(@server_token)
   end
 
   def send_map(client)
@@ -200,9 +196,9 @@ class TeeworldsServer
   def on_ctrl_token(packet)
     u = Unpacker.new(packet.payload[1..])
     token = u.get_raw(4)
-    token = token.map { |b| b.to_s(16).rjust(2, '0') }.join
-    puts "got token #{token}"
-    send_ctrl_with_token(packet.addr, token)
+    token_str = token.map { |b| b.to_s(16).rjust(2, '0') }.join
+    puts "got token #{token_str}"
+    send_ctrl_with_token(packet.addr, token_str)
   end
 
   def on_ctrl_keep_alive(packet)
@@ -214,14 +210,14 @@ class TeeworldsServer
   end
 
   def on_ctrl_connect(packet)
-    puts 'got connection, sending accept'
-
     id = get_next_client_id
     if id == -1
       puts 'server full drop packet. TODO: tell the client'
       return
     end
-    client = Client.new(id:, addr: packet.addr, token: packet.payload[...4])
+    token = bytes_to_str(packet.payload[...4])
+    puts "got connection, sending accept (client token: #{token})"
+    client = Client.new(id:, addr: packet.addr, token:)
     @clients[id] = client
     @netbase.send_packet([NET_CTRLMSG_ACCEPT], 0, control: true, addr: packet.addr)
   end
diff --git a/lib/token.rb b/lib/token.rb
new file mode 100644
index 0000000..91f8f17
--- /dev/null
+++ b/lib/token.rb
@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+require_relative 'bytes'
+
+class SecurityToken
+  def self.err_msg(msg, token)
+    hex = ''
+    hex = "hex: #{str_hex(token)}\n" if token.instance_of?(String)
+    "Invalid token! Token should be a human readable hex string!\n" \
+      "  Good sample token: aabbccdd\n" \
+      "  #{msg}\n" \
+      "  token: #{token}:#{token.class}\n" \
+      "  #{hex}"
+  end
+
+  def self.validate(token)
+    raise err_msg("Expected type: String got: #{token.class}", token) unless token.instance_of?(String)
+    raise err_msg("Expected size: 8 got: #{token.size}", token) unless token.size == 8
+  end
+end