fix(security): escape generated table-name Go string literal (#84)

Author: keroxp

generator.go

@@ -7,6 +7,7 @@ import (
 	"go/format"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"text/template"
 
@@ -26,15 +27,16 @@ type GenerateOptions struct {
 }
 
 type templateData struct {
-	Imports       string
-	Model         string
-	ModelLower    string
-	M             string
-	Package       string
-	Fields        string
-	UpdaterFields string
-	ScannedFields string
-	TableName     string
+	Imports            string
+	Model              string
+	ModelLower         string
+	M                  string
+	Package            string
+	Fields             string
+	UpdaterFields      string
+	ScannedFields      string
+	TableName          string
+	TableNameGoLiteral string
 }
 
 func NewGenerator(db *sql.DB) Generator {
@@ -120,15 +122,16 @@ func (d *generator) generateModelFile(tableName string, opt *GenerateOptions) er
 		}
 	}
 	data := &templateData{
-		Imports:       strings.Join(imports, "\n"),
-		Model:         strcase.ToCamel(table.TableName),
-		ModelLower:    strcase.ToLowerCamel(table.TableName),
-		M:             table.TableName[0:1],
-		UpdaterFields: updateFields.String(),
-		Package:       opt.Package,
-		Fields:        fields.String(),
-		TableName:     tableName,
-		ScannedFields: scannedFields.String(),
+		Imports:            strings.Join(imports, "\n"),
+		Model:              strcase.ToCamel(table.TableName),
+		ModelLower:         strcase.ToLowerCamel(table.TableName),
+		M:                  table.TableName[0:1],
+		UpdaterFields:      updateFields.String(),
+		Package:            opt.Package,
+		Fields:             fields.String(),
+		TableName:          tableName,
+		TableNameGoLiteral: strconv.Quote(tableName),
+		ScannedFields:      scannedFields.String(),
 	}
 	outFile := filepath.Join(
 		opt.OutDir,
@@ -167,5 +170,5 @@ func ({{.M}} *Update{{.Model}}) UpdateTableName() string {
 	return {{.Model}}TableName
 }
 
-const {{.Model}}TableName = "{{.TableName}}"
+const {{.Model}}TableName = {{.TableNameGoLiteral}}
 `

generator_test.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"regexp"
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
@@ -108,3 +109,30 @@ func TestGenerator_Generate(t *testing.T) {
 			}), "columns err")
 	})
 }
+
+func TestGenerator_Generate_EscapesTableNameGoLiteral(t *testing.T) {
+	mockDb, mock, err := sqlmock.New()
+	assert.NoError(t, err)
+	defer mockDb.Close()
+
+	table := `x";func init(){panic(1)};var _="`
+	mock.ExpectQuery(`show tables`).WillReturnRows(sqlmock.NewRows([]string{"tables"}).AddRow(table))
+	mock.ExpectQuery(regexp.QuoteMeta(fmt.Sprintf("show columns from `%s`", table))).WillReturnRows(
+		sqlmock.NewRows([]string{"Field", "Type", "Null", "Key", "Default", "Extra"}).
+			AddRow("id", "int(11)", "NO", "PRI", nil, ""),
+	)
+
+	dir := t.TempDir()
+	err = NewGenerator(mockDb).Generate(&GenerateOptions{OutDir: dir, Package: "dist"})
+	assert.NoError(t, err)
+
+	entries, err := os.ReadDir(dir)
+	assert.NoError(t, err)
+	if assert.Len(t, entries, 1) {
+		content, err := os.ReadFile(filepath.Join(dir, entries[0].Name()))
+		assert.NoError(t, err)
+		assert.Contains(t, string(content), `const XfuncInitpanic1varTableName = "x\";func init(){panic(1)};var _=\""`)
+		assert.NotContains(t, string(content), "\nfunc init()")
+	}
+	assert.NoError(t, mock.ExpectationsWereMet())
+}