fix the SQL Injection (#330)

Author: haoel

global/global.go

@@ -311,3 +311,22 @@ func CommandLine(cmd string, args []string) string {
 	}
 	return result
 }
+
+// EscapeQuote escape the string the single quote, double quote, and backtick
+func EscapeQuote(str string) string {
+	type Escape struct {
+		From string
+		To   string
+	}
+	escape := []Escape{
+		{From: "`", To: ""}, // remove the backtick
+		{From: `\`, To: `\\`},
+		{From: `'`, To: `\'`},
+		{From: `"`, To: `\"`},
+	}
+
+	for _, e := range escape {
+		str = strings.ReplaceAll(str, e.From, e.To)
+	}
+	return str
+}

global/global_test.go

@@ -438,3 +438,12 @@ func TestCommandLine(t *testing.T) {
 	s = CommandLine("kubectl", []string{"get", "pod", "--all-namespaces", "-o", "json"})
 	assert.Equal(t, "kubectl get pod --all-namespaces -o json", s)
 }
+
+func TestEscape(t *testing.T) {
+	assert.Equal(t, "test", EscapeQuote("test"))
+	assert.Equal(t, "test", EscapeQuote("`test`"))
+	assert.Equal(t, `\'test\'`, EscapeQuote("'test'"))
+	assert.Equal(t, `\"test\"`, EscapeQuote(`"test"`))
+	assert.Equal(t, `\\test\\`, EscapeQuote(`\test\`))
+
+}

probe/client/mysql/mysql.go

@@ -106,50 +106,65 @@ func (r *MySQL) Probe() (bool, string) {
 
 	// Check if we need to query specific data
 	if len(r.Data) > 0 {
-		for k, v := range r.Data {
-			log.Debugf("[%s / %s / %s] - Verifying Data - [%s] : [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, k, v)
-			sql, err := r.getSQL(k)
-			if err != nil {
-				return false, err.Error()
-			}
-			log.Debugf("[%s / %s / %s] - SQL - [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, sql)
-			rows, err := db.Query(sql)
-			if err != nil {
-				return false, err.Error()
-			}
-			if !rows.Next() {
-				rows.Close()
-				return false, fmt.Sprintf("No data found for [%s]", k)
-			}
-			//check the value is equal to the value in data
-			var value string
-			if err := rows.Scan(&value); err != nil {
-				rows.Close()
-				return false, err.Error()
-			}
-			if value != v {
-				rows.Close()
-				return false, fmt.Sprintf("Value not match for [%s] expected [%s] got [%s] ", k, v, value)
-			}
-			rows.Close()
-			log.Debugf("[%s / %s / %s] - Data Verified Successfully! - [%s] : [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, k, v)
-		}
-	} else {
-		err = db.Ping()
-		if err != nil {
+		if err := r.ProbeWithDataVerification(db); err != nil {
 			return false, err.Error()
 		}
-		row, err := db.Query("show status like \"uptime\"") // run a SQL to test
-		if err != nil {
+	} else {
+		if err := r.ProbeWithPing(db); err != nil {
 			return false, err.Error()
 		}
-		defer row.Close()
 	}
 
 	return true, "Check MySQL Server Successfully!"
 
 }
 
+// ProbeWithPing do the health check with ping
+func (r *MySQL) ProbeWithPing(db *sql.DB) error {
+	if err := db.Ping(); err != nil {
+		return err
+	}
+	row, err := db.Query("show status like \"uptime\"") // run a SQL to test
+	if err != nil {
+		return err
+	}
+	defer row.Close()
+	return nil
+}
+
+// ProbeWithDataVerification do the health check with data verification
+func (r *MySQL) ProbeWithDataVerification(db *sql.DB) error {
+	for k, v := range r.Data {
+		log.Debugf("[%s / %s / %s] - Verifying Data - [%s] : [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, k, v)
+		sql, err := r.getSQL(k)
+		if err != nil {
+			return err
+		}
+		log.Debugf("[%s / %s / %s] - SQL - [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, sql)
+		rows, err := db.Query(sql)
+		if err != nil {
+			return err
+		}
+		if !rows.Next() {
+			rows.Close()
+			return fmt.Errorf("No data found for [%s]", k)
+		}
+		//check the value is equal to the value in data
+		var value string
+		if err := rows.Scan(&value); err != nil {
+			rows.Close()
+			return err
+		}
+		if value != v {
+			rows.Close()
+			return fmt.Errorf("Value not match for [%s] expected [%s] got [%s] ", k, v, value)
+		}
+		rows.Close()
+		log.Debugf("[%s / %s / %s] - Data Verified Successfully! - [%s] : [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, k, v)
+	}
+	return nil
+}
+
 // getSQL get the SQL statement
 // input: database:table:column:key:value
 // output: SELECT column FROM database.table WHERE key = value
@@ -161,16 +176,16 @@ func (r *MySQL) getSQL(str string) (string, error) {
 	if len(fields) != 5 {
 		return "", fmt.Errorf("Invalid SQL data - [%s]. (syntax: database:table:field:key:value)", str)
 	}
-	db := fields[0]
-	table := fields[1]
-	field := fields[2]
-	key := fields[3]
-	value := fields[4]
+	db := global.EscapeQuote(fields[0])
+	table := global.EscapeQuote(fields[1])
+	field := global.EscapeQuote(fields[2])
+	key := global.EscapeQuote(fields[3])
+	value := global.EscapeQuote(fields[4])
 	//check value is int or not
 	if _, err := strconv.Atoi(value); err != nil {
 		return "", fmt.Errorf("Invalid SQL data - [%s], the value must be int", str)
 	}
 
-	sql := fmt.Sprintf("SELECT %s FROM %s.%s WHERE %s = %s", field, db, table, key, value)
+	sql := fmt.Sprintf("SELECT `%s` FROM `%s`.`%s` WHERE `%s` = %s", field, db, table, key, value)
 	return sql, nil
 }

probe/client/postgres/postgres.go

@@ -190,16 +190,16 @@ func (r *PostgreSQL) getSQL(str string) (string, string, error) {
 	if len(fields) != 5 {
 		return "", "", fmt.Errorf("Invalid SQL data - [%s]. (syntax: database:table:field:key:value)", str)
 	}
-	db := fields[0]
-	table := fields[1]
-	field := fields[2]
-	key := fields[3]
-	value := fields[4]
+	db := global.EscapeQuote(fields[0])
+	table := global.EscapeQuote(fields[1])
+	field := global.EscapeQuote(fields[2])
+	key := global.EscapeQuote(fields[3])
+	value := global.EscapeQuote(fields[4])
 	//check value is int or not
 	if _, err := strconv.Atoi(value); err != nil {
 		return "", "", fmt.Errorf("Invalid SQL data - [%s], the value must be int", str)
 	}
 
-	sql := fmt.Sprintf("SELECT %s FROM %s WHERE %s = %s", field, table, key, value)
+	sql := fmt.Sprintf(`SELECT "%s" FROM "%s" WHERE "%s" = %s`, field, table, key, value)
 	return db, sql, nil
 }