refactor: sync mysql changes acc to latest code (keploy#1809)

* refactor: sync mysql changes acc to latest code Signed-off-by: Sarthak160 <[email protected]> * Fixed lint issues Signed-off-by: Chinmay <[email protected]> * fix: Added an argument to DecodeMySQL function to find out the mode Signed-off-by: Chinmay <[email protected]> * refactor: remove set mocks function form the memdb interface Signed-off-by: Sarthak160 <[email protected]> --------- Signed-off-by: Sarthak160 <[email protected]> Signed-off-by: Chinmay <[email protected]> Co-authored-by: Chinmay <[email protected]>
Yaxhveer · May 13, 2024 · b1efc78 · b1efc78
1 parent b9673e3
commit b1efc78
Show file tree

Hide file tree

Showing 30 changed files with 665 additions and 480 deletions.
diff --git a/pkg/core/proxy/integrations/mysql/authMoreData.go b/pkg/core/proxy/integrations/mysql/authMoreData.go
@@ -6,7 +6,7 @@ import (
 )
 
 type NextAuthPacket struct {
-	PluginData byte `json:"plugin_data,omitempty" yaml:"plugin_data,omitempty"`
+	PluginData byte `yaml:"plugin_data"`
 }
 
 func decodeAuthMoreData(data []byte) (*NextAuthPacket, error) {
@@ -17,3 +17,11 @@ func decodeAuthMoreData(data []byte) (*NextAuthPacket, error) {
 		PluginData: data[0],
 	}, nil
 }
+
+// Encode function for Next Authentication method Packet
+//func encodeAuthMoreData(packet *NextAuthPacket) ([]byte, error) {
+//	if packet.PluginData != 0x02 {
+//		return nil, errors.New("invalid PluginData value for NextAuthPacket")
+//	}
+//	return []byte{packet.PluginData}, nil
+//}
diff --git a/pkg/core/proxy/integrations/mysql/authSwitchRequestPacket.go b/pkg/core/proxy/integrations/mysql/authSwitchRequestPacket.go
@@ -8,9 +8,9 @@ import (
 )
 
 type AuthSwitchRequestPacket struct {
-	StatusTag      byte   `json:"status_tag,omitempty" yaml:"status_tag,omitempty"`
-	PluginName     string `json:"plugin_name,omitempty" yaml:"plugin_name,omitempty"`
-	PluginAuthData string `json:"plugin_authdata,omitempty" yaml:"plugin_authdata,omitempty"`
+	StatusTag      byte   `yaml:"status_tag"`
+	PluginName     string `yaml:"plugin_name"`
+	PluginAuthData string `yaml:"plugin_authdata"`
 }
 
 func decodeAuthSwitchRequest(data []byte) (*AuthSwitchRequestPacket, error) {

diff --git a/pkg/core/proxy/integrations/mysql/authSwitchResponsePacket.go b/pkg/core/proxy/integrations/mysql/authSwitchResponsePacket.go
@@ -5,7 +5,7 @@ import (
 )
 
 type AuthSwitchResponsePacket struct {
-	AuthResponseData string `json:"auth_response_data,omitempty" yaml:"auth_response_data,omitempty"`
+	AuthResponseData string `yaml:"auth_response_data"`
 }
 
 func decodeAuthSwitchResponse(data []byte) (*AuthSwitchResponsePacket, error) {

diff --git a/pkg/core/proxy/integrations/mysql/comChangeUserPacket.go b/pkg/core/proxy/integrations/mysql/comChangeUserPacket.go
@@ -1,17 +1,16 @@
 package mysql
 
 import (
-	"encoding/base64"
 	"errors"
 	"strings"
 )
 
 type ComChangeUserPacket struct {
-	User         string `json:"user,omitempty" yaml:"user,omitempty,flow"`
-	Auth         string `json:"auth,omitempty" yaml:"auth,omitempty,flow"`
-	Db           string `json:"db,omitempty" yaml:"db,omitempty,flow"`
-	CharacterSet uint8  `json:"character_set,omitempty" yaml:"character_set,omitempty,flow"`
-	AuthPlugin   string `json:"auth_plugin,omitempty" yaml:"auth_plugin,omitempty,flow"`
+	User         string `yaml:"user"`
+	Auth         []byte `yaml:"auth,omitempty,flow"`
+	Db           string `yaml:"db"`
+	CharacterSet uint8  `yaml:"character_set"`
+	AuthPlugin   string `yaml:"auth_plugin"`
 }
 
 func decodeComChangeUser(data []byte) (ComChangeUserPacket, error) {
@@ -33,7 +32,7 @@ func decodeComChangeUser(data []byte) (ComChangeUserPacket, error) {
 
 	return ComChangeUserPacket{
 		User:         user,
-		Auth:         base64.StdEncoding.EncodeToString(auth),
+		Auth:         auth,
 		Db:           db,
 		CharacterSet: characterSet,
 		AuthPlugin:   authPlugin,

diff --git a/pkg/core/proxy/integrations/mysql/comFetchPacket.go b/pkg/core/proxy/integrations/mysql/comFetchPacket.go
@@ -6,9 +6,9 @@ import (
 )
 
 type ComStmtFetchPacket struct {
-	StatementID uint32 `json:"statement_id,omitempty" yaml:"statement_id,omitempty"`
-	RowCount    uint32 `json:"row_count,omitempty" yaml:"row_count,omitempty"`
-	Info        string `json:"info,omitempty" yaml:"info,omitempty"`
+	StatementID uint32 `yaml:"statement_id"`
+	RowCount    uint32 `yaml:"row_count"`
+	Info        string `yaml:"info"`
 }
 
 func decodeComStmtFetch(data []byte) (ComStmtFetchPacket, error) {

diff --git a/pkg/core/proxy/integrations/mysql/comInitDb.go b/pkg/core/proxy/integrations/mysql/comInitDb.go
@@ -1,6 +1,20 @@
 package mysql
 
 type ComInitDbPacket struct {
-	Status byte   `json:"status,omitempty" yaml:"status,omitempty"`
-	DbName string `json:"db_name,omitempty" yaml:"db_name,omitempty"`
+	Status byte
+	DbName string
 }
+
+//func decodeComInitDb(data []byte) (*ComInitDbPacket, error) {
+//	if len(data) < 2 {
+//		return nil, errors.New("data too short for COM_INIT_DB")
+//	}
+//	status := data[0]
+//
+//	// The rest of the packet after the command byte is the database name
+//	dbName := string(data[1:])
+//	return &ComInitDbPacket{
+//		Status: status,
+//		DbName: dbName,
+//	}, nil
+//}
diff --git a/pkg/core/proxy/integrations/mysql/comStmtCloseMoreData.go b/pkg/core/proxy/integrations/mysql/comStmtCloseMoreData.go
@@ -7,13 +7,13 @@ import (
 )
 
 type ComStmtPreparePacket1 struct {
-	Header []byte `json:"header,omitempty" yaml:"header,omitempty"`
-	Query  string `json:"query,omitempty" yaml:"query,omitempty"`
+	Header []byte
+	Query  string
 }
 
 type ComStmtCloseAndPrepare struct {
-	StmtClose   ComStmtClosePacket    `json:"stmt_close,omitempty" yaml:"stmt_close,omitempty"`
-	StmtPrepare ComStmtPreparePacket1 `json:"stmt_prepare,omitempty" yaml:"stmt_prepare,omitempty"`
+	StmtClose   ComStmtClosePacket
+	StmtPrepare ComStmtPreparePacket1
 }
 
 func decodeComStmtCloseMoreData(data []byte) (*ComStmtCloseAndPrepare, error) {

diff --git a/pkg/core/proxy/integrations/mysql/comStmtClosePacket.go b/pkg/core/proxy/integrations/mysql/comStmtClosePacket.go
@@ -6,8 +6,8 @@ import (
 )
 
 type ComStmtClosePacket struct {
-	Status      byte   `json:"status,omitempty" yaml:"status,omitempty"`
-	StatementID uint32 `json:"statement_id,omitempty" yaml:"statement_id,omitempty"`
+	Status      byte
+	StatementID uint32
 }
 
 func decodeComStmtClose(data []byte) (*ComStmtClosePacket, error) {

diff --git a/pkg/core/proxy/integrations/mysql/comStmtPrepareOk.go b/pkg/core/proxy/integrations/mysql/comStmtPrepareOk.go
@@ -9,13 +9,13 @@ import (
 )
 
 type StmtPrepareOk struct {
-	Status       byte               `json:"status,omitempty" yaml:"status,omitempty,flow"`
-	StatementID  uint32             `json:"statement_id,omitempty" yaml:"statement_id,omitempty,flow"`
-	NumColumns   uint16             `json:"num_columns,omitempty" yaml:"num_columns,omitempty,flow"`
-	NumParams    uint16             `json:"num_params,omitempty" yaml:"num_params,omitempty,flow"`
-	WarningCount uint16             `json:"warning_count,omitempty" yaml:"warning_count,omitempty,flow"`
-	ColumnDefs   []ColumnDefinition `json:"column_definitions,omitempty" yaml:"column_definitions,omitempty,flow"`
-	ParamDefs    []ColumnDefinition `json:"param_definitions,omitempty" yaml:"param_definitions,omitempty,flow"`
+	Status       byte               `yaml:"status"`
+	StatementID  uint32             `yaml:"statement_id"`
+	NumColumns   uint16             `yaml:"num_columns"`
+	NumParams    uint16             `yaml:"num_params"`
+	WarningCount uint16             `yaml:"warning_count"`
+	ColumnDefs   []ColumnDefinition `yaml:"column_definitions,omitempty,flow"`
+	ParamDefs    []ColumnDefinition `yaml:"param_definitions,omitempty,flow"`
 }
 
 func decodeComStmtPrepareOk(data []byte) (*StmtPrepareOk, error) {
@@ -191,24 +191,12 @@ func encodeStmtPrepareOk(packet *models.MySQLStmtPrepareOk) ([]byte, error) {
 
 func encodeColumnDefinition(buf *bytes.Buffer, column *models.ColumnDefinition, seqNum *byte) error {
 	tmpBuf := &bytes.Buffer{}
-	if err := writeLengthEncodedString(tmpBuf, column.Catalog); err != nil {
-		return err
-	}
-	if err := writeLengthEncodedString(tmpBuf, column.Schema); err != nil {
-		return err
-	}
-	if err := writeLengthEncodedString(tmpBuf, column.Table); err != nil {
-		return err
-	}
-	if err := writeLengthEncodedString(tmpBuf, column.OrgTable); err != nil {
-		return err
-	}
-	if err := writeLengthEncodedString(tmpBuf, column.Name); err != nil {
-		return err
-	}
-	if err := writeLengthEncodedString(tmpBuf, column.OrgName); err != nil {
-		return err
-	}
+	writeLengthEncodedString(tmpBuf, column.Catalog)
+	writeLengthEncodedString(tmpBuf, column.Schema)
+	writeLengthEncodedString(tmpBuf, column.Table)
+	writeLengthEncodedString(tmpBuf, column.OrgTable)
+	writeLengthEncodedString(tmpBuf, column.Name)
+	writeLengthEncodedString(tmpBuf, column.OrgName)
 	tmpBuf.WriteByte(0x0C)
 	if err := binary.Write(tmpBuf, binary.LittleEndian, column.CharacterSet); err != nil {
 		return err

diff --git a/pkg/core/proxy/integrations/mysql/comStmtPreparePacket.go b/pkg/core/proxy/integrations/mysql/comStmtPreparePacket.go
@@ -6,7 +6,7 @@ import (
 )
 
 type ComStmtPreparePacket struct {
-	Query string `json:"query,omitempty" yaml:"query,omitempty,flow"`
+	Query string
 }
 
 func decodeComStmtPrepare(data []byte) (*ComStmtPreparePacket, error) {

diff --git a/pkg/core/proxy/integrations/mysql/comStmtResetPacket.go b/pkg/core/proxy/integrations/mysql/comStmtResetPacket.go
@@ -6,7 +6,7 @@ import (
 )
 
 type COM_STMT_RESET struct {
-	StatementID uint32 `json:"statement_id,omitempty" yaml:"statement_id,omitempty,flow"`
+	StatementID uint32 `yaml:"statement_id"`
 }
 
 func decodeComStmtReset(packet []byte) (*COM_STMT_RESET, error) {

diff --git a/pkg/core/proxy/integrations/mysql/comStmtSendLongDataPacket.go b/pkg/core/proxy/integrations/mysql/comStmtSendLongDataPacket.go
@@ -1,15 +1,14 @@
 package mysql
 
 import (
-	"encoding/base64"
 	"encoding/binary"
 	"fmt"
 )
 
 type COM_STMT_SEND_LONG_DATA struct {
-	StatementID uint32 `json:"statement_id,omitempty" yaml:"statement_id,omitempty,flow"`
-	ParameterID uint16 `json:"parameter_id,omitempty" yaml:"parameter_id,omitempty,flow"`
-	Data        string `json:"data,omitempty" yaml:"data,omitempty,flow"`
+	StatementID uint32 `yaml:"statement_id"`
+	ParameterID uint16 `yaml:"parameter_id"`
+	Data        []byte `yaml:"data,omitempty,flow"`
 }
 
 func decodeComStmtSendLongData(packet []byte) (COM_STMT_SEND_LONG_DATA, error) {
@@ -22,6 +21,6 @@ func decodeComStmtSendLongData(packet []byte) (COM_STMT_SEND_LONG_DATA, error) {
 	return COM_STMT_SEND_LONG_DATA{
 		StatementID: stmtID,
 		ParameterID: paramID,
-		Data:        base64.StdEncoding.EncodeToString(data),
+		Data:        data,
 	}, nil
 }
diff --git a/pkg/core/proxy/integrations/mysql/decode.go b/pkg/core/proxy/integrations/mysql/decode.go
@@ -43,7 +43,7 @@ func decodeMySQL(ctx context.Context, logger *zap.Logger, clientConn net.Conn, d
 					errCh <- err
 					return
 				}
-				sqlMock, found := getFirstSQLMock(configMocks)
+				sqlMock, matchedIndex, found := getFirstSQLMock(configMocks)
 				if !found {
 					logger.Debug("No SQL mock found")
 					errCh <- err
@@ -59,7 +59,6 @@ func decodeMySQL(ctx context.Context, logger *zap.Logger, clientConn net.Conn, d
 					errCh <- err
 					return
 				}
-
 				_, err = clientConn.Write(binaryPacket)
 				if err != nil {
 					if ctx.Err() != nil {
@@ -69,17 +68,17 @@ func decodeMySQL(ctx context.Context, logger *zap.Logger, clientConn net.Conn, d
 					errCh <- err
 					return
 				}
-				matchedIndex := 0
 				matchedReqIndex := 0
 				configMocks[matchedIndex].Spec.MySQLResponses = append(configMocks[matchedIndex].Spec.MySQLResponses[:matchedReqIndex], configMocks[matchedIndex].Spec.MySQLResponses[matchedReqIndex+1:]...)
 				if len(configMocks[matchedIndex].Spec.MySQLResponses) == 0 {
-					configMocks = append(configMocks[:matchedIndex], configMocks[matchedIndex+1:]...)
-					err = mockDb.FlagMockAsUsed(configMocks[matchedIndex])
-					if err != nil {
-						utils.LogError(logger, err, "Failed to flag mock as used")
-						errCh <- err
-						return
-					}
+					//configMocks = append(configMocks[:matchedIndex], configMocks[matchedIndex+1:]...)
+					//err = mockDb.FlagMockAsUsed(configMocks[matchedIndex])
+					//if err != nil {
+					//	utils.LogError(logger, err, "Failed to flag mock as used")
+					//	errCh <- err
+					//	return
+					//}
+					mockDb.DeleteUnFilteredMock(configMocks[matchedIndex])
 				}
 				//h.SetConfigMocks(configMocks)
 				firstLoop = false
@@ -131,7 +130,7 @@ func decodeMySQL(ctx context.Context, logger *zap.Logger, clientConn net.Conn, d
 					expectingHandshakeResponseTest = true
 				}
 
-				oprRequest, requestHeader, decodedRequest, err := DecodeMySQLPacket(logger, bytesToMySQLPacket(requestBuffer))
+				oprRequest, requestHeader, decodedRequest, err := DecodeMySQLPacket(logger, bytesToMySQLPacket(requestBuffer), clientConn, models.MODE_TEST)
 				if err != nil {
 					utils.LogError(logger, err, "Failed to decode MySQL packet")
 					errCh <- err
@@ -206,7 +205,6 @@ func decodeMySQL(ctx context.Context, logger *zap.Logger, clientConn net.Conn, d
 					errCh <- err
 					return
 				}
-
 				_, err = clientConn.Write(responseBinary)
 				if err != nil {
 					if ctx.Err() != nil {
@@ -228,11 +226,11 @@ func decodeMySQL(ctx context.Context, logger *zap.Logger, clientConn net.Conn, d
 	}
 }
 
-func getFirstSQLMock(configMocks []*models.Mock) (*models.Mock, bool) {
-	for _, mock := range configMocks {
+func getFirstSQLMock(configMocks []*models.Mock) (*models.Mock, int, bool) {
+	for index, mock := range configMocks {
 		if len(mock.Spec.MySQLResponses) > 0 && mock.Kind == "SQL" && mock.Spec.MySQLResponses[0].Header.PacketType == "MySQLHandshakeV10" {
-			return mock, true
+			return mock, index, true
 		}
 	}
-	return nil, false
+	return nil, 0, false
 }