From 03472de9f8f38abab735cd5ca5e462a59038b790 Mon Sep 17 00:00:00 2001
From: jonnosan <jonnosan@93682198-c243-4bdb-bd91-e943c89aac3b>
Date: Thu, 26 Mar 2009 10:09:17 +0000
Subject: [PATCH] add "write request" (file upload) support to tftp server

git-svn-id: http://svn.code.sf.net/p/netboot65/code@60 93682198-c243-4bdb-bd91-e943c89aac3b
---
 server/lib/tftp_server.rb | 76 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 72 insertions(+), 4 deletions(-)

diff --git a/server/lib/tftp_server.rb b/server/lib/tftp_server.rb
index 62319c8..86355d1 100644
--- a/server/lib/tftp_server.rb
+++ b/server/lib/tftp_server.rb
@@ -1,10 +1,10 @@
 #
 # minimal TFTP server implementation for use with netboot65
 #
-# supports RRQ and a custom DIR request (opcode 0x65 - returns null terminated list of filenames that match specified filemask)
+# supports RRQ, WRQ and a custom DIR request (opcode 0x65 - returns null terminated list of filenames that match specified filemask)
 # Jonno Downes (jonno@jamtronix.com) - January, 2009
 # 
-#
+# TFTP spec : http://www.ietf.org/rfc/rfc1350.txt
 
 require 'socket'
 class Netboot65TFTPServer
@@ -81,6 +81,59 @@ class Netboot65TFTPServer
     end
   end
 
+
+  def recv_data(client_ip,client_port,filename,file_handle)
+    
+    client_sock=UDPSocket.open
+    client_sock.connect(client_ip,client_port)
+    
+    log_msg("receiving #{filename} from #{client_ip}:#{client_port}")
+    got_last_block=false
+    sent_last_ack=false
+    block_number=0
+    until sent_last_ack do
+      packet=[4,block_number].pack("nn")
+      got_block=false
+      TFTP_MAX_RESENDS.times do |attempt_number|
+        log_msg("ACKing block #{block_number} of #{filename} from #{client_ip}:#{client_port} - attempt #{attempt_number+1}")
+        client_sock.send(packet,0,client_ip,client_port)
+        if got_last_block then
+          puts "last block received"
+          sent_last_ack=true
+          break
+        else
+          if (IO.select([client_sock], nil, nil, 1)) then
+            data,addr_info=client_sock.recvfrom(4096)
+            client_ip=addr_info[3]
+            client_port=addr_info[1]        
+            opcode=data[0,2].unpack("n")[0]          
+            opcode_description=TFTP_OPCODES[opcode]
+            if opcode==3  then
+              recved_block_number=data[2,2].unpack("n")[0]
+              block_data=data[4,512]
+              got_last_block=(block_data.length!=512)
+              log_msg "TFTP: DATA from #{client_ip}:#{client_port} - block #{recved_block_number} - #{block_data.length} bytes"
+              if recved_block_number==block_number+1
+                got_block=true 
+                file_handle<<block_data
+              end
+            else             
+              opcode_description="[UNK]" if opcode_description.nil?          
+              log_msg "TFTP: response from #{client_ip}:#{client_port} - opcode #{opcode} : #{opcode_description}"
+            end
+          end
+          break if got_block
+          if !got_block then
+            log_msg "TFTP: timed out waiting for DATA for block #{block_number+1} from #{client_ip}"
+            break
+          end            
+        end
+      end
+      block_number+=1
+    end
+    file_handle.close
+  end
+
   def start()
     log_msg "TFTP: serving #{bootfile_dir} on port #{port}"
     Socket.do_not_reverse_lookup = true
@@ -104,7 +157,7 @@ class Netboot65TFTPServer
            filename,mode=filename_and_mode.split(0.chr)
            log_msg "RRQ for #{filename} (#{mode})"
            if filename=~/^\./ || filename=~/\.\./ then #looks like something dodgy - either a dotfile or a directory traversal attempt
-            send_error(client_ip,client_port,1,"'#{filename}' invalid") 
+            send_error(client_ip,client_port,1,"'#{filename}' invalid filename") 
            else
              full_filename="#{bootfile_dir}/#{filename}"
              if File.file?(full_filename) then
@@ -114,12 +167,27 @@ class Netboot65TFTPServer
               send_error(client_ip,client_port,1,"'#{filename}' not found") 
             end
           end
+          when 2 : #WRITE REQUEST
+           opcode,filename_and_mode=data.unpack("nA*")
+           filename,mode=filename_and_mode.split(0.chr)
+           log_msg "WRQ for #{filename} (#{mode})"
+           if filename=~/^\./ || filename=~/\.\./ then #looks like something dodgy - either a dotfile or a directory traversal attempt
+            send_error(client_ip,client_port,1,"'#{filename}' invalid filename") 
+           else
+             full_filename="#{bootfile_dir}/#{filename}"
+             begin
+              file_handle=File.open(full_filename,"wb")
+              Thread.new {recv_data(client_ip,client_port,full_filename,file_handle)}
+             rescue Exception=>e
+              send_error(client_ip,client_port,2,"error writing to '#{filename}':#{e.to_s}") 
+            end
+          end          
           when 0x65 : #DIR REQUEST
            opcode,filemask_and_mode=data.unpack("nA*")
            filemask,mode=filemask_and_mode.split(0.chr)
            log_msg "DIR for #{filemask} (#{mode})"
            if filename=~/^\./ || filename=~/\.\./ then #looks like something dodgy - either a dotfile or a directory traversal attempt
-            send_error(client_ip,client_port,1,"'#{filemask}' invalid") 
+            send_error(client_ip,client_port,1,"'#{filemask}' invalid filename") 
            else
              data_to_send=""
              Dir.chdir(bootfile_dir) do