Repository: pingcap/tidb-lightning
Pull request #414 (2020-10-01 16:17:43)

mydump: fix infinite loop in ExportStatement when Read() returns an error but is not EOF
<!--
Thank you for contributing to TiDB-Lightning! Please read the [CONTRIBUTING](https://github.com/pingcap/tidb-lightning/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->

In the `ExportStatement()` function, if the reader errored which is not EOF, it will cause the function go into an infinite loop because the error is ignored.

### What is changed and how it works?

Handle the error.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test

Side effects

Related changes


glorv: LGTM

Patch:
diff --git a/lightning/mydump/reader.go b/lightning/mydump/reader.go
index f129d8a..08cf063 100644
--- a/lightning/mydump/reader.go
+++ b/lightning/mydump/reader.go
@@ -80,17 +80,21 @@ func ExportStatement(ctx context.Context, store storage.ExternalStorage, sqlFile
 	data := make([]byte, 0, sqlFile.Size+1)
 	buffer := make([]byte, 0, sqlFile.Size+1)
 	for {
-		line, err := br.ReadString('\n')
-		if errors.Cause(err) == io.EOF && len(line) == 0 { // it will return EOF if there is no trailing new line.
-			break
+		line, err := br.ReadBytes('\n')
+		if errors.Cause(err) == io.EOF {
+			if len(line) == 0 { // it will return EOF if there is no trailing new line.
+				break
+			}
+		} else if err != nil {
+			return nil, errors.Trace(err)
 		}
 
-		line = strings.TrimSpace(line)
+		line = bytes.TrimSpace(line)
 		if len(line) == 0 {
 			continue
 		}
 
-		buffer = append(buffer, []byte(line)...)
+		buffer = append(buffer, line...)
 		if buffer[len(buffer)-1] == ';' {
 			statement := string(buffer)
 			if !(strings.HasPrefix(statement, "/*") && strings.HasSuffix(statement, "*/;")) {
diff --git a/mock/importer.go b/mock/importer.go
index 50af4d4..a58c0d0 100644
--- a/mock/importer.go
+++ b/mock/importer.go
@@ -1,6 +1,8 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/pingcap/kvproto/pkg/import_kvpb (interfaces: ImportKVClient,ImportKV_WriteEngineClient)
 
+// $ mockgen -package mock github.com/pingcap/kvproto/pkg/import_kvpb ImportKVClient,ImportKV_WriteEngineClient
+
 // Package mock is a generated GoMock package.
 package mock
 
diff --git a/mock/storage.go b/mock/storage.go
new file mode 100644
index 0000000..dc1c49f
--- /dev/null
+++ b/mock/storage.go
@@ -0,0 +1,125 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/pingcap/br/pkg/storage (interfaces: ExternalStorage)
+
+// $ mockgen -package mock github.com/pingcap/br/pkg/storage ExternalStorage
+
+// Package mock is a generated GoMock package.
+package mock
+
+import (
+	context "context"
+	gomock "github.com/golang/mock/gomock"
+	storage "github.com/pingcap/br/pkg/storage"
+	reflect "reflect"
+)
+
+// MockExternalStorage is a mock of ExternalStorage interface
+type MockExternalStorage struct {
+	ctrl     *gomock.Controller
+	recorder *MockExternalStorageMockRecorder
+}
+
+// MockExternalStorageMockRecorder is the mock recorder for MockExternalStorage
+type MockExternalStorageMockRecorder struct {
+	mock *MockExternalStorage
+}
+
+// NewMockExternalStorage creates a new mock instance
+func NewMockExternalStorage(ctrl *gomock.Controller) *MockExternalStorage {
+	mock := &MockExternalStorage{ctrl: ctrl}
+	mock.recorder = &MockExternalStorageMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockExternalStorage) EXPECT() *MockExternalStorageMockRecorder {
+	return m.recorder
+}
+
+// CreateUploader mocks base method
+func (m *MockExternalStorage) CreateUploader(arg0 context.Context, arg1 string) (storage.Uploader, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CreateUploader", arg0, arg1)
+	ret0, _ := ret[0].(storage.Uploader)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// CreateUploader indicates an expected call of CreateUploader
+func (mr *MockExternalStorageMockRecorder) CreateUploader(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateUploader", reflect.TypeOf((*MockExternalStorage)(nil).CreateUploader), arg0, arg1)
+}
+
+// FileExists mocks base method
+func (m *MockExternalStorage) FileExists(arg0 context.Context, arg1 string) (bool, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "FileExists", arg0, arg1)
+	ret0, _ := ret[0].(bool)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// FileExists indicates an expected call of FileExists
+func (mr *MockExternalStorageMockRecorder) FileExists(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FileExists", reflect.TypeOf((*MockExternalStorage)(nil).FileExists), arg0, arg1)
+}
+
+// Open mocks base method
+func (m *MockExternalStorage) Open(arg0 context.Context, arg1 string) (storage.ReadSeekCloser, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Open", arg0, arg1)
+	ret0, _ := ret[0].(storage.ReadSeekCloser)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Open indicates an expected call of Open
+func (mr *MockExternalStorageMockRecorder) Open(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Open", reflect.TypeOf((*MockExternalStorage)(nil).Open), arg0, arg1)
+}
+
+// Read mocks base method
+func (m *MockExternalStorage) Read(arg0 context.Context, arg1 string) ([]byte, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Read", arg0, arg1)
+	ret0, _ := ret[0].([]byte)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Read indicates an expected call of Read
+func (mr *MockExternalStorageMockRecorder) Read(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Read", reflect.TypeOf((*MockExternalStorage)(nil).Read), arg0, arg1)
+}
+
+// WalkDir mocks base method
+func (m *MockExternalStorage) WalkDir(arg0 context.Context, arg1 *storage.WalkOption, arg2 func(string, int64) error) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "WalkDir", arg0, arg1, arg2)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// WalkDir indicates an expected call of WalkDir
+func (mr *MockExternalStorageMockRecorder) WalkDir(arg0, arg1, arg2 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WalkDir", reflect.TypeOf((*MockExternalStorage)(nil).WalkDir), arg0, arg1, arg2)
+}
+
+// Write mocks base method
+func (m *MockExternalStorage) Write(arg0 context.Context, arg1 string, arg2 []byte) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Write", arg0, arg1, arg2)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Write indicates an expected call of Write
+func (mr *MockExternalStorageMockRecorder) Write(arg0, arg1, arg2 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Write", reflect.TypeOf((*MockExternalStorage)(nil).Write), arg0, arg1, arg2)
+}


Test patch:
diff --git a/lightning/mydump/reader_test.go b/lightning/mydump/reader_test.go
index f115f7c..2336ae8 100644
--- a/lightning/mydump/reader_test.go
+++ b/lightning/mydump/reader_test.go
@@ -15,13 +15,16 @@ package mydump_test
 
 import (
 	"context"
+	"errors"
 	"io/ioutil"
 	"os"
 
+	"github.com/golang/mock/gomock"
 	"github.com/pingcap/br/pkg/storage"
-
 	. "github.com/pingcap/check"
+
 	. "github.com/pingcap/tidb-lightning/lightning/mydump"
+	"github.com/pingcap/tidb-lightning/mock"
 )
 
 //////////////////////////////////////////////////////////
@@ -159,6 +162,34 @@ func (s *testMydumpReaderSuite) TestExportStatementGibberishError(c *C) {
 
 	f := FileInfo{FileMeta: SourceFileMeta{Path: stat.Name()}, Size: stat.Size()}
 	data, err := ExportStatement(context.TODO(), store, f, "auto")
-	c.Assert(data, IsNil)
-	c.Assert(err, NotNil)
+	c.Assert(data, HasLen, 0)
+	c.Assert(err, ErrorMatches, `failed to decode \w* as auto: invalid schema encoding`)
+}
+
+type AlwaysErrorReadSeekCloser struct{}
+
+func (AlwaysErrorReadSeekCloser) Read([]byte) (int, error) {
+	return 0, errors.New("read error!")
+}
+func (AlwaysErrorReadSeekCloser) Seek(int64, int) (int64, error) {
+	return 0, errors.New("seek error!")
+}
+func (AlwaysErrorReadSeekCloser) Close() error {
+	return nil
+}
+
+func (s *testMydumpReaderSuite) TestExportStatementHandleNonEOFError(c *C) {
+	controller := gomock.NewController(c)
+	defer controller.Finish()
+
+	ctx := context.TODO()
+
+	mockStorage := mock.NewMockExternalStorage(controller)
+	mockStorage.EXPECT().
+		Open(ctx, "no-perm-file").
+		Return(AlwaysErrorReadSeekCloser{}, nil)
+
+	f := FileInfo{FileMeta: SourceFileMeta{Path: "no-perm-file"}, Size: 1}
+	_, err := ExportStatement(ctx, mockStorage, f, "auto")
+	c.Assert(err, ErrorMatches, "read error!")
 }