support OPT_LOAD_DATA_LOCAL_DIR

Incompatible change: raise error if the server requires a local file
whiout client's OPT_LOCAL_INFILE option or if the filename specified
by the server does not match OPT_LOAD_DATA_LOCAL_DIR.

Author: tmtm

README.rdoc

@@ -42,7 +42,8 @@ MySQL connector for Ruby.
 * Mysql#options でサポートしているオプションは次のものだけです:
   Mysql::INIT_COMMAND, Mysql::OPT_CONNECT_TIMEOUT,
   Mysql::OPT_LOCAL_INFILE, Mysql::OPT_READ_TIMEOUT,
-  Mysql::OPT_WRITE_TIMEOUT, Mysql::SET_CHARSET_NAME.
+  Mysql::OPT_WRITE_TIMEOUT, Mysql::SET_CHARSET_NAME,
+  Mysql::OPT_LOAD_DATA_LOCAL_DIR.
   これら以外を指定すると "option not implementted" という warning が標準エラー出力に出力されます。
 
 * Mysql#use_result は Mysql#store_result と同じです。つまりサーバーから一気に結果セットを読み込みます。
@@ -64,5 +65,5 @@ MySQL connector for Ruby.
 == Copyright
 
 Author    :: TOMITA Masahiro <[email protected]>
-Copyright :: Copyright (c) 2009-2012 TOMITA Masahiro
+Copyright :: Copyright (c) 2008 TOMITA Masahiro
 License   :: Ruby's

lib/mysql.rb

@@ -1,5 +1,5 @@
 # coding: ascii-8bit
-# Copyright (C) 2008-2012 TOMITA Masahiro
+# Copyright (C) 2008 TOMITA Masahiro
 # mailto:[email protected]
 
 # MySQL connection class.
@@ -112,8 +112,8 @@ def connect(host=nil, user=nil, passwd=nil, db=nil, port=nil, socket=nil, flag=0
       warn 'unsupported flag: CLIENT_COMPRESS' if $VERBOSE
       flag &= ~CLIENT_COMPRESS
     end
-    @protocol = Protocol.new host, port, socket, @connect_timeout, @read_timeout, @write_timeout
-    @protocol.authenticate user, passwd, db, (@local_infile ? CLIENT_LOCAL_FILES : 0) | flag, @charset
+    @protocol = Protocol.new host, port, socket, @connect_timeout, @read_timeout, @write_timeout, @local_infile
+    @protocol.authenticate user, passwd, db, flag, @charset
     @charset ||= @protocol.charset
     @host_info = (host.nil? || host == "localhost") ? 'Localhost via UNIX socket' : "#{host} via TCP/IP"
     query @init_command if @init_command
@@ -158,6 +158,8 @@ def options(opt, value=nil)
       @connect_timeout = value
 #    when Mysql::GUESS_CONNECTION
     when Mysql::OPT_LOCAL_INFILE
+      @local_infile = value ? '' : nil
+    when Mysql::OPT_LOAD_DATA_LOCAL_DIR
       @local_infile = value
 #    when Mysql::OPT_NAMED_PIPE
 #    when Mysql::OPT_PROTOCOL

lib/mysql/protocol.rb

@@ -1,5 +1,5 @@
 # coding: ascii-8bit
-# Copyright (C) 2008-2012 TOMITA Masahiro
+# Copyright (C) 2008 TOMITA Masahiro
 # mailto:[email protected]
 
 require "socket"
@@ -134,15 +134,17 @@ def self.value2net(v)
     # conn_timeout :: [Integer] connect timeout (sec).
     # read_timeout :: [Integer] read timeout (sec).
     # write_timeout :: [Integer] write timeout (sec).
+    # local_infile :: [String] local infile path
     # === Exception
     # [ClientError] :: connection timeout
-    def initialize(host, port, socket, conn_timeout, read_timeout, write_timeout)
+    def initialize(host, port, socket, conn_timeout, read_timeout, write_timeout, local_infile)
       @insert_id = 0
       @warning_count = 0
       @gc_stmt_queue = []   # stmt id list which GC destroy.
       set_state :INIT
       @read_timeout = read_timeout
       @write_timeout = write_timeout
+      @local_infile = local_infile
       begin
         Timeout.timeout conn_timeout do
           if host.nil? or host.empty? or host == "localhost"
@@ -180,6 +182,7 @@ def authenticate(user, passwd, db, flag, charset)
       @server_version = init_packet.server_version.split(/\D/)[0,3].inject{|a,b|a.to_i*100+b.to_i}
       @thread_id = init_packet.thread_id
       client_flags = CLIENT_LONG_PASSWORD | CLIENT_LONG_FLAG | CLIENT_TRANSACTIONS | CLIENT_PROTOCOL_41 | CLIENT_SECURE_CONNECTION
+      client_flags |= CLIENT_LOCAL_FILES if @local_infile
       client_flags |= CLIENT_CONNECT_WITH_DB if db
       client_flags |= flag
       @charset = charset
@@ -230,10 +233,7 @@ def get_result
           return res_packet.field_count
         end
         if res_packet.field_count.nil?      # LOAD DATA LOCAL INFILE
-          filename = res_packet.message
-          File.open(filename){|f| write f}
-          write nil  # EOF mark
-          read
+          send_local_file(res_packet.message)
         end
         @affected_rows, @insert_id, @server_status, @warning_count, @message =
           res_packet.affected_rows, res_packet.insert_id, res_packet.server_status, res_packet.warning_count, res_packet.message
@@ -245,6 +245,19 @@ def get_result
       end
     end
 
+    # send local file to server
+    def send_local_file(filename)
+      filename = File.absolute_path(filename)
+      if filename.start_with? @local_infile
+        File.open(filename){|f| write f}
+      else
+        raise ClientError::LoadDataLocalInfileRejected, 'LOAD DATA LOCAL INFILE file request rejected due to restrictions on access.'
+      end
+    ensure
+      write nil # EOF mark
+      read
+    end
+
     # Retrieve n fields
     # === Argument
     # n :: [Integer] number of fields

test/test_mysql.rb

@@ -160,6 +160,40 @@ class TestMysql < Test::Unit::TestCase
       @m.query("load data local infile '#{tmpf.path}' into table t")
       assert{ @m.query('select * from t').fetch_row == ['123','abc'] }
     end
+    test 'OPT_LOAD_DATA_LOCAL_DIR: client can execute LOAD DATA LOCAL INFILE query with specified directory' do
+      require 'tempfile'
+      tmpf = Tempfile.new 'mysql_spec'
+      tmpf.puts "123\tabc\n"
+      tmpf.close
+      assert{ @m.options(Mysql::OPT_LOAD_DATA_LOCAL_DIR, File.dirname(tmpf.path)) == @m }
+      @m.connect(MYSQL_SERVER, MYSQL_USER, MYSQL_PASSWORD, MYSQL_DATABASE, MYSQL_PORT, MYSQL_SOCKET)
+      @m.query('create temporary table t (i int, c char(10))')
+      @m.query("load data local infile '#{tmpf.path}' into table t")
+      assert{ @m.query('select * from t').fetch_row == ['123','abc'] }
+    end
+    test 'OPT_LOAD_DATA_LOCAL_DIR: client cannot execute LOAD DATA LOCAL INFILE query without specified directory' do
+      require 'tempfile'
+      tmpf = Tempfile.new 'mysql_spec'
+      tmpf.puts "123\tabc\n"
+      tmpf.close
+      assert{ @m.options(Mysql::OPT_LOAD_DATA_LOCAL_DIR, '/hoge') == @m }
+      @m.connect(MYSQL_SERVER, MYSQL_USER, MYSQL_PASSWORD, MYSQL_DATABASE, MYSQL_PORT, MYSQL_SOCKET)
+      @m.query('create temporary table t (i int, c char(10))')
+      assert_raise Mysql::ClientError::LoadDataLocalInfileRejected, 'LOAD DATA LOCAL INFILE file request rejected due to restrictions on access.' do
+        @m.query("load data local infile '#{tmpf.path}' into table t")
+      end
+    end
+    test 'without OPT_LOCAL_INFILE and OPT_LOAD_DATA_LOCAL_DIR: client cannot execute LOAD DATA LOCAL INFILE query' do
+      require 'tempfile'
+      tmpf = Tempfile.new 'mysql_spec'
+      tmpf.puts "123\tabc\n"
+      tmpf.close
+      @m.connect(MYSQL_SERVER, MYSQL_USER, MYSQL_PASSWORD, MYSQL_DATABASE, MYSQL_PORT, MYSQL_SOCKET)
+      @m.query('create temporary table t (i int, c char(10))')
+      assert_raise Mysql::ServerError::NotAllowedCommand, 'The used command is not allowed with this MySQL version' do
+        @m.query("load data local infile '#{tmpf.path}' into table t")
+      end
+    end
     test 'OPT_READ_TIMEOUT: set timeout for reading packet' do
       assert{ @m.options(Mysql::OPT_READ_TIMEOUT, 10) == @m }
     end