Merge pull request #22 from nixys/fix/19

fix(#19): Preserve MySQL `CREATE TABLE` statement
nixys · Jun 7, 2024 · 215a9db · 215a9db
2 parents 11afbc7 + 8e5cb1e
commit 215a9db
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 78 deletions.
diff --git a/modules/anonymizers/mysql/dh.go b/modules/anonymizers/mysql/dh.go
@@ -9,16 +9,16 @@ import (
 	"github.com/nixys/nxs-data-anonymizer/modules/filters/relfilter"
 )
 
-func dhSecurityCreateTable(usrCtx any, deferred, token []byte) ([]byte, error) {
+func dhSecurityInsertInto(usrCtx any, deferred, token []byte) ([]byte, error) {
 
 	uctx := usrCtx.(*userCtx)
 
-	uctx.security.tmpBuf = append(uctx.security.tmpBuf, token...)
+	uctx.security.tmpBuf = token
 
 	return deferred, nil
 }
 
-func dhSecurityCreateTableName(usrCtx any, deferred, token []byte) ([]byte, error) {
+func dhSecurityInsertIntoTableNameSearch(usrCtx any, deferred, token []byte) ([]byte, error) {
 
 	uctx := usrCtx.(*userCtx)
 
@@ -42,37 +42,14 @@ func dhSecurityNil(usrCtx any, deferred, token []byte) ([]byte, error) {
 func dhCreateTableName(usrCtx any, deferred, token []byte) ([]byte, error) {
 
 	uctx := usrCtx.(*userCtx)
+	uctx.filter.TableCreate(string(deferred))
 
-	tn := string(deferred)
-
-	// Check table pass through security rules
-	if !securityPolicyCheck(uctx, tn) {
-
-		// If not: table will be skipped from result dump
-
-		uctx.security.isSkip = true
-		uctx.security.tmpBuf = []byte{}
-		return []byte{}, nil
-	}
-
-	uctx.filter.TableCreate(tn)
-
-	d := append(uctx.security.tmpBuf, append(deferred, token...)...)
-
-	uctx.security.isSkip = false
-	uctx.security.tmpBuf = []byte{}
-
-	return d, nil
+	return append(deferred, token...), nil
 }
 
 func dhCreateTableFieldName(usrCtx any, deferred, token []byte) ([]byte, error) {
 
 	uctx := usrCtx.(*userCtx)
-
-	if uctx.security.isSkip == true {
-		return []byte{}, nil
-	}
-
 	uctx.column.name = string(deferred)
 
 	return append(deferred, token...), nil
@@ -82,10 +59,6 @@ func dhCreateTableColumnTypeAdd(usrCtx any, deferred, token []byte) ([]byte, err
 
 	uctx := usrCtx.(*userCtx)
 
-	if uctx.security.isSkip == true {
-		return []byte{}, nil
-	}
-
 	for k, v := range typeKeys {
 		if k == "generated" {
 			if k == string(token) || strings.ToUpper(k) == string(token) {
@@ -100,21 +73,13 @@ func dhCreateTableColumnTypeAdd(usrCtx any, deferred, token []byte) ([]byte, err
 		}
 	}
 
-	if uctx.column.columnType == "" {
-		fmt.Println("token:", token)
-	}
-
 	return append(deferred, token...), nil
 }
 
 func dhCreateTableColumnAdd(usrCtx any, deferred, token []byte) ([]byte, error) {
 
 	uctx := usrCtx.(*userCtx)
 
-	if uctx.security.isSkip == true {
-		return []byte{}, nil
-	}
-
 	if uctx.column.isSkip == false {
 		uctx.filter.ColumnAdd(uctx.column.name, uctx.column.columnType)
 	}
@@ -128,16 +93,29 @@ func dhInsertIntoTableName(usrCtx any, deferred, token []byte) ([]byte, error) {
 
 	uctx := usrCtx.(*userCtx)
 
-	if uctx.security.isSkip == true {
+	tn := string(deferred)
+
+	// Check table pass through security rules
+	if !securityPolicyCheck(uctx, tn) {
+
+		// If not: table will be skipped from result dump
+
+		uctx.security.isSkip = true
+		uctx.security.tmpBuf = []byte{}
 		return []byte{}, nil
 	}
 
+	d := append(uctx.security.tmpBuf, append(deferred, token...)...)
+
+	uctx.security.isSkip = false
+	uctx.security.tmpBuf = []byte{}
+
 	// Check insert into table name
-	if bytes.Compare([]byte(uctx.filter.TableNameGet()), deferred) != 0 {
-		return append(deferred, token...), fmt.Errorf("`create` and `insert into` table names are mismatch (create table: '%s', insert into table: '%s')", uctx.filter.TableNameGet(), string(deferred))
+	if tn != uctx.filter.TableNameGet() {
+		return d, fmt.Errorf("`create` and `insert into` table names are mismatch (create table: '%s', insert into table: '%s')", uctx.filter.TableNameGet(), tn)
 	}
 
-	return append(deferred, token...), nil
+	return d, nil
 }
 
 func dhCreateTableValues(usrCtx any, deferred, token []byte) ([]byte, error) {

diff --git a/modules/anonymizers/mysql/mysql.go b/modules/anonymizers/mysql/mysql.go
@@ -143,7 +143,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									R: []byte{' '},
 								},
 							},
-							DataHandler: dhSecurityCreateTable,
+							DataHandler: nil,
 						},
 					},
 				},
@@ -158,7 +158,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									R: []byte{' '},
 								},
 							},
-							DataHandler: dhSecurityCreateTableName,
+							DataHandler: nil,
 						},
 					},
 				},
@@ -169,7 +169,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 							Switch: fsm.Switch{
 								Trigger: []byte("`"),
 							},
-							DataHandler: dhSecurityCreateTableName,
+							DataHandler: nil,
 						},
 					},
 				},
@@ -191,7 +191,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 							Switch: fsm.Switch{
 								Trigger: []byte("("),
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: nil,
 						},
 					},
 				},
@@ -207,7 +207,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									R: []byte{' '},
 								},
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: nil,
 						},
 						{
 							// Skip table keys description
@@ -219,7 +219,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									R: []byte{' '},
 								},
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: nil,
 						},
 						{
 							// Skip table keys description
@@ -231,7 +231,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									R: []byte{' '},
 								},
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: nil,
 						},
 						{
 							// Skip table keys description
@@ -243,7 +243,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									R: []byte{' '},
 								},
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: nil,
 						},
 						{
 							// Skip table keys description
@@ -255,14 +255,14 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									R: []byte{' '},
 								},
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: nil,
 						},
 						{
 							Name: stateFieldsDescriptionName,
 							Switch: fsm.Switch{
 								Trigger: []byte("`"),
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: nil,
 						},
 					},
 				},
@@ -276,7 +276,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									R: []byte{'\n'},
 								},
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: nil,
 						},
 						{
 							Name: statefFieldsDescriptionBlockEnd,
@@ -286,7 +286,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									L: []byte{'\n'},
 								},
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: nil,
 						},
 					},
 				},
@@ -363,7 +363,7 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									R: []byte{'\n'},
 								},
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: nil,
 						},
 					},
 				},
@@ -378,45 +378,30 @@ func Init(ctx context.Context, r io.Reader, s InitSettings) io.Reader {
 									R: []byte{' '},
 								},
 							},
-							DataHandler: dhSecurityCreateTable,
+							DataHandler: nil,
 						},
 						{
-							Name: stateInsertInto,
+							Name: stateInsertIntoTableNameSearch,
 							Switch: fsm.Switch{
-								Trigger: []byte("INSERT"),
+								Trigger: []byte("INSERT INTO"),
 								Delimiters: fsm.Delimiters{
 									L: []byte{'\n'},
 									R: []byte{' '},
 								},
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: dhSecurityInsertInto,
 						},
 					},
 				},
 
-				stateInsertInto: {
-					NextStates: []fsm.NextState{
-						{
-							Name: stateInsertIntoTableNameSearch,
-							Switch: fsm.Switch{
-								Trigger: []byte("INTO"),
-								Delimiters: fsm.Delimiters{
-									L: []byte{' '},
-									R: []byte{' '},
-								},
-							},
-							DataHandler: dhSecurityNil,
-						},
-					},
-				},
 				stateInsertIntoTableNameSearch: {
 					NextStates: []fsm.NextState{
 						{
 							Name: stateInsertIntoTableName,
 							Switch: fsm.Switch{
 								Trigger: []byte("`"),
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: dhSecurityInsertIntoTableNameSearch,
 						},
 					},
 				},

diff --git a/modules/anonymizers/mysql/states.go b/modules/anonymizers/mysql/states.go
@@ -14,7 +14,6 @@ var (
 	stateFieldDescriptionTailSkip   = fsm.StateName("fields description tail skip")
 	statefFieldsDescriptionBlockEnd = fsm.StateName("fields description block end")
 	stateSomeIntermediateState      = fsm.StateName("some intermediate state")
-	stateInsertInto                 = fsm.StateName("insert into")
 	stateInsertIntoTableNameSearch  = fsm.StateName("insert into table name search")
 	stateInsertIntoTableName        = fsm.StateName("insert into table name")
 	stateValuesSearch               = fsm.StateName("values search")