|
|
@@ -13,11 +13,14 @@ import (
|
|
|
"io"
|
|
|
"os"
|
|
|
"strings"
|
|
|
+ "sync"
|
|
|
)
|
|
|
|
|
|
var (
|
|
|
- fileRegister map[string]bool
|
|
|
- readerRegister map[string]func() io.Reader
|
|
|
+ fileRegister map[string]bool
|
|
|
+ fileRegisterLock sync.RWMutex
|
|
|
+ readerRegister map[string]func() io.Reader
|
|
|
+ readerRegisterLock sync.RWMutex
|
|
|
)
|
|
|
|
|
|
// RegisterLocalFile adds the given file to the file whitelist,
|
|
|
@@ -32,17 +35,21 @@ var (
|
|
|
// ...
|
|
|
//
|
|
|
func RegisterLocalFile(filePath string) {
|
|
|
+ fileRegisterLock.Lock()
|
|
|
// lazy map init
|
|
|
if fileRegister == nil {
|
|
|
fileRegister = make(map[string]bool)
|
|
|
}
|
|
|
|
|
|
fileRegister[strings.Trim(filePath, `"`)] = true
|
|
|
+ fileRegisterLock.Unlock()
|
|
|
}
|
|
|
|
|
|
// DeregisterLocalFile removes the given filepath from the whitelist.
|
|
|
func DeregisterLocalFile(filePath string) {
|
|
|
+ fileRegisterLock.Lock()
|
|
|
delete(fileRegister, strings.Trim(filePath, `"`))
|
|
|
+ fileRegisterLock.Unlock()
|
|
|
}
|
|
|
|
|
|
// RegisterReaderHandler registers a handler function which is used
|
|
|
@@ -61,18 +68,22 @@ func DeregisterLocalFile(filePath string) {
|
|
|
// ...
|
|
|
//
|
|
|
func RegisterReaderHandler(name string, handler func() io.Reader) {
|
|
|
+ readerRegisterLock.Lock()
|
|
|
// lazy map init
|
|
|
if readerRegister == nil {
|
|
|
readerRegister = make(map[string]func() io.Reader)
|
|
|
}
|
|
|
|
|
|
readerRegister[name] = handler
|
|
|
+ readerRegisterLock.Unlock()
|
|
|
}
|
|
|
|
|
|
// DeregisterReaderHandler removes the ReaderHandler function with
|
|
|
// the given name from the registry.
|
|
|
func DeregisterReaderHandler(name string) {
|
|
|
+ readerRegisterLock.Lock()
|
|
|
delete(readerRegister, name)
|
|
|
+ readerRegisterLock.Unlock()
|
|
|
}
|
|
|
|
|
|
func deferredClose(err *error, closer io.Closer) {
|
|
|
@@ -90,7 +101,11 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
|
|
|
// The server might return an an absolute path. See issue #355.
|
|
|
name = name[idx+8:]
|
|
|
|
|
|
- if handler, inMap := readerRegister[name]; inMap {
|
|
|
+ readerRegisterLock.RLock()
|
|
|
+ handler, inMap := readerRegister[name]
|
|
|
+ readerRegisterLock.RUnlock()
|
|
|
+
|
|
|
+ if inMap {
|
|
|
rdr = handler()
|
|
|
if rdr != nil {
|
|
|
data = make([]byte, 4+mc.maxWriteSize)
|
|
|
@@ -106,7 +121,10 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
|
|
|
}
|
|
|
} else { // File
|
|
|
name = strings.Trim(name, `"`)
|
|
|
- if mc.cfg.allowAllFiles || fileRegister[name] {
|
|
|
+ fileRegisterLock.RLock()
|
|
|
+ fr := fileRegister[name]
|
|
|
+ fileRegisterLock.RUnlock()
|
|
|
+ if mc.cfg.allowAllFiles || fr {
|
|
|
var file *os.File
|
|
|
var fi os.FileInfo
|
|
|
|