From a9d1410160ba81d291557e6d4d86637d216ac944 Mon Sep 17 00:00:00 2001 From: Julien Schmidt Date: Fri, 15 Nov 2013 22:20:12 +0100 Subject: [PATCH 1/4] infile: lazy map init --- infile.go | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/infile.go b/infile.go index 952bd3d6d..9d3ec033e 100644 --- a/infile.go +++ b/infile.go @@ -21,11 +21,6 @@ var ( readerRegister map[string]func() io.Reader ) -func init() { - fileRegister = make(map[string]bool) - readerRegister = make(map[string]func() io.Reader) -} - // RegisterLocalFile adds the given file to the file whitelist, // so that it can be used by "LOAD DATA LOCAL INFILE ". // Alternatively you can allow the use of all local files with @@ -38,6 +33,11 @@ func init() { // ... // func RegisterLocalFile(filePath string) { + // lazy map init + if fileRegister == nil { + fileRegister = make(map[string]bool) + } + fileRegister[strings.Trim(filePath, `"`)] = true } @@ -62,6 +62,11 @@ func DeregisterLocalFile(filePath string) { // ... // func RegisterReaderHandler(name string, handler func() io.Reader) { + // lazy map init + if readerRegister == nil { + readerRegister = make(map[string]func() io.Reader) + } + readerRegister[name] = handler } From d503c4bfce6797e196f3de19683507e9ac77b5e1 Mon Sep 17 00:00:00 2001 From: Julien Schmidt Date: Fri, 15 Nov 2013 22:35:58 +0100 Subject: [PATCH 2/4] infile: better buffer size calculation --- infile.go | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/infile.go b/infile.go index 9d3ec033e..1594ede2f 100644 --- a/infile.go +++ b/infile.go @@ -78,25 +78,38 @@ func DeregisterReaderHandler(name string) { func (mc *mysqlConn) handleInFileRequest(name string) (err error) { var rdr io.Reader - data := make([]byte, 4+mc.maxWriteSize) + var data []byte if strings.HasPrefix(name, "Reader::") { // io.Reader name = name[8:] - handler, inMap := readerRegister[name] - if handler != nil { + if handler, inMap := readerRegister[name]; inMap { rdr = handler() - } - if rdr == nil { - if !inMap { - err = fmt.Errorf("Reader '%s' is not registered", name) + if rdr != nil { + data = make([]byte, 4+mc.maxWriteSize) } else { err = fmt.Errorf("Reader '%s' is ", name) } + } else { + err = fmt.Errorf("Reader '%s' is not registered", name) } } else { // File name = strings.Trim(name, `"`) if mc.cfg.allowAllFiles || fileRegister[name] { - rdr, err = os.Open(name) + var file *os.File + var fi os.FileInfo + + if file, err = os.Open(name); err == nil { + if fi, err = file.Stat(); err == nil { + rdr = file + if fileSize := int(fi.Size()); fileSize <= mc.maxWriteSize { + data = make([]byte, 4+fileSize) + } else if fileSize <= mc.maxPacketAllowed { + data = make([]byte, 4+mc.maxWriteSize) + } else { + err = fmt.Errorf("Local File '%s' too large: Size: %d, Max: %d", name, fileSize, mc.maxPacketAllowed) + } + } + } } else { err = fmt.Errorf("Local File '%s' is not registered. Use the DSN parameter 'allowAllFiles=true' to allow all files", name) } From d7e2ac4160f9d2b066d05f40c7213097cddfea9b Mon Sep 17 00:00:00 2001 From: Julien Schmidt Date: Fri, 15 Nov 2013 23:26:12 +0100 Subject: [PATCH 3/4] infile: refactoring --- infile.go | 55 ++++++++++++++++++++++++++++--------------------------- 1 file changed, 28 insertions(+), 27 deletions(-) diff --git a/infile.go b/infile.go index 1594ede2f..522b54eba 100644 --- a/infile.go +++ b/infile.go @@ -9,7 +9,6 @@ package mysql import ( - "database/sql/driver" "fmt" "io" "os" @@ -86,6 +85,16 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) { rdr = handler() if rdr != nil { data = make([]byte, 4+mc.maxWriteSize) + + if rdc, ok := rdr.(io.ReadCloser); ok { + defer func() { + if err == nil { + err = rdc.Close() + } else { + rdc.Close() + } + }() + } } else { err = fmt.Errorf("Reader '%s' is ", name) } @@ -99,6 +108,15 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) { var fi os.FileInfo if file, err = os.Open(name); err == nil { + defer func() { + if err == nil { + err = file.Close() + } else { + file.Close() + } + }() + + // get file size if fi, err = file.Stat(); err == nil { rdr = file if fileSize := int(fi.Size()); fileSize <= mc.maxWriteSize { @@ -115,45 +133,28 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) { } } - if rdc, ok := rdr.(io.ReadCloser); ok { - defer func() { - if err == nil { - err = rdc.Close() - } else { - rdc.Close() - } - }() - } - // send content packets - var ioErr error if err == nil { var n int - for err == nil && ioErr == nil { + for err == nil { n, err = rdr.Read(data[4:]) if n > 0 { - ioErr = mc.writePacket(data[:4+n]) + if ioErr := mc.writePacket(data[:4+n]); ioErr != nil { + return ioErr + } } } if err == io.EOF { err = nil } - if ioErr != nil { - errLog.Print(ioErr.Error()) - return driver.ErrBadConn - } } // send empty packet (termination) - ioErr = mc.writePacket([]byte{ - 0x00, - 0x00, - 0x00, - mc.sequence, - }) - if ioErr != nil { - errLog.Print(ioErr.Error()) - return driver.ErrBadConn + if data == nil { + data = make([]byte, 4) + } + if ioErr := mc.writePacket(data[:4]); ioErr != nil { + return ioErr } // read OK packet From 34f105cf7fda695e2101f55ac5b6420ba5137d6f Mon Sep 17 00:00:00 2001 From: Julien Schmidt Date: Sat, 16 Nov 2013 14:54:04 +0100 Subject: [PATCH 4/4] infile: new deferredClose func instead of lambda funcs --- infile.go | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/infile.go b/infile.go index 522b54eba..121a04c71 100644 --- a/infile.go +++ b/infile.go @@ -75,6 +75,13 @@ func DeregisterReaderHandler(name string) { delete(readerRegister, name) } +func deferredClose(err *error, closer io.Closer) { + closeErr := closer.Close() + if *err == nil { + *err = closeErr + } +} + func (mc *mysqlConn) handleInFileRequest(name string) (err error) { var rdr io.Reader var data []byte @@ -86,14 +93,8 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) { if rdr != nil { data = make([]byte, 4+mc.maxWriteSize) - if rdc, ok := rdr.(io.ReadCloser); ok { - defer func() { - if err == nil { - err = rdc.Close() - } else { - rdc.Close() - } - }() + if cl, ok := rdr.(io.Closer); ok { + defer deferredClose(&err, cl) } } else { err = fmt.Errorf("Reader '%s' is ", name) @@ -108,13 +109,7 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) { var fi os.FileInfo if file, err = os.Open(name); err == nil { - defer func() { - if err == nil { - err = file.Close() - } else { - file.Close() - } - }() + defer deferredClose(&err, file) // get file size if fi, err = file.Stat(); err == nil {