backend: fix panic when reading result set fails (#486)

Author: djshow832

pkg/proxy/backend/cmd_processor_query.go

@@ -39,8 +39,9 @@ func (cp *CmdProcessor) query(packetIO *pnet.PacketIO, sql string) (result *mysq
 		err = errors.WithStack(mysql.ErrMalformPacket)
 	default:
 		var rs *mysql.Result
-		rs, err = cp.readResultSet(packetIO, response)
-		result = rs.Resultset
+		if rs, err = cp.readResultSet(packetIO, response); err == nil {
+			result = rs.Resultset
+		}
 	}
 	return
 }

pkg/proxy/backend/cmd_processor_test.go

@@ -1068,3 +1068,35 @@ func TestNetworkError(t *testing.T) {
 	ts.query(t, proxyErrChecker)
 	clean()
 }
+
+// Test that TiProxy won't panic when query encounters an error.
+func TestQueryError(t *testing.T) {
+	tc := newTCPConnSuite(t)
+	tests := []struct {
+		cfg cfgOverrider
+		c   checker
+	}{
+		{
+			cfg: func(cfg *testConfig) {
+				cfg.backendConfig.abnormalExit = true
+			},
+		},
+		{
+			cfg: func(cfg *testConfig) {
+				cfg.backendConfig.respondType = responseTypeErr
+			},
+		},
+		{
+			cfg: func(cfg *testConfig) {
+				cfg.backendConfig.respondType = responseTypeResultSet
+				cfg.backendConfig.columns = 1
+				cfg.backendConfig.exitInResult = true
+			},
+		},
+	}
+	for _, test := range tests {
+		ts, clean := newTestSuite(t, tc, test.cfg)
+		ts.query(t, func(t *testing.T, ts *testSuite) {})
+		clean()
+	}
+}

pkg/proxy/backend/mock_backend_test.go

@@ -28,6 +28,7 @@ type backendConfig struct {
 	proxyProtocol bool
 	authSucceed   bool
 	abnormalExit  bool
+	exitInResult  bool
 }
 
 func newBackendConfig() *backendConfig {
@@ -293,6 +294,12 @@ func (mb *mockBackend) writeResultSet(packetIO *pnet.PacketIO, names []string, v
 				return err
 			}
 		}
+		if mb.exitInResult {
+			if err := packetIO.Flush(); err != nil {
+				return err
+			}
+			return packetIO.Close()
+		}
 
 		if status&mysql.SERVER_STATUS_CURSOR_EXISTS == 0 {
 			if mb.capability&pnet.ClientDeprecateEOF == 0 {