Skip to content

Commit

Permalink
optional obfuscate positional params, booleans and null
Browse files Browse the repository at this point in the history
  • Loading branch information
lu-zhengda committed Sep 13, 2023
1 parent a3ef011 commit 4be19ca
Show file tree
Hide file tree
Showing 5 changed files with 96 additions and 12 deletions.
40 changes: 38 additions & 2 deletions obfuscator.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,11 @@ import (
)

type obfuscatorConfig struct {
ReplaceDigits bool
DollarQuotedFunc bool
DollarQuotedFunc bool
ReplaceDigits bool
ReplacePositionalParameter bool
ReplaceBoolean bool
ReplaceNull bool
}

type obfuscatorOption func(*obfuscatorConfig)
Expand All @@ -17,6 +20,24 @@ func WithReplaceDigits(replaceDigits bool) obfuscatorOption {
}
}

func WithReplacePositionalParameter(replacePositionalParameter bool) obfuscatorOption {
return func(c *obfuscatorConfig) {
c.ReplacePositionalParameter = replacePositionalParameter
}
}

func WithReplaceBoolean(replaceBoolean bool) obfuscatorOption {
return func(c *obfuscatorConfig) {
c.ReplaceBoolean = replaceBoolean
}
}

func WithReplaceNull(replaceNull bool) obfuscatorOption {
return func(c *obfuscatorConfig) {
c.ReplaceNull = replaceNull
}
}

func WithDollarQuotedFunc(dollarQuotedFunc bool) obfuscatorOption {
return func(c *obfuscatorConfig) {
c.DollarQuotedFunc = dollarQuotedFunc
Expand Down Expand Up @@ -69,7 +90,22 @@ func (o *Obfuscator) Obfuscate(input string, lexerOpts ...lexerOption) string {
}
case STRING, INCOMPLETE_STRING, DOLLAR_QUOTED_STRING:
obfuscatedSQL.WriteString(StringPlaceholder)
case POSITIONAL_PARAMETER:
if o.config.ReplacePositionalParameter {
obfuscatedSQL.WriteString(StringPlaceholder)
} else {
obfuscatedSQL.WriteString(token.Value)
}
case IDENT:
if o.config.ReplaceBoolean && isBoolean(token.Value) {
obfuscatedSQL.WriteString(StringPlaceholder)
continue
}
if o.config.ReplaceNull && isNull(token.Value) {
obfuscatedSQL.WriteString(StringPlaceholder)
continue
}

if o.config.ReplaceDigits {
obfuscatedSQL.WriteString(replaceDigits(token.Value, "?"))
} else {
Expand Down
50 changes: 45 additions & 5 deletions obfuscator_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,14 @@ import (

func TestObfuscator(t *testing.T) {
tests := []struct {
input string
expected string
replaceDigits bool
dollarQuotedFunc bool
dbms DBMSType
input string
expected string
replaceDigits bool
replacePositionalParameter bool
replaceBoolean bool
replaceNull bool
dollarQuotedFunc bool
dbms DBMSType
}{
{
input: "SELECT * FROM users where id = 1",
Expand Down Expand Up @@ -55,6 +58,30 @@ func TestObfuscator(t *testing.T) {
expected: "SELECT * FROM users? where id = ?",
replaceDigits: true,
},
{
input: "SELECT * FROM users where id is NULL and is_active = TRUE and is_admin = FALSE",
expected: "SELECT * FROM users where id is NULL and is_active = TRUE and is_admin = FALSE",
replaceBoolean: false,
replaceNull: false,
},
{
input: "SELECT * FROM users where id is NULL and is_active = TRUE and is_admin = FALSE",
expected: "SELECT * FROM users where id is NULL and is_active = ? and is_admin = ?",
replaceBoolean: true,
replaceNull: false,
},
{
input: "SELECT * FROM users where id is NULL and is_active = TRUE and is_admin = FALSE",
expected: "SELECT * FROM users where id is ? and is_active = TRUE and is_admin = FALSE",
replaceBoolean: false,
replaceNull: true,
},
{
input: "SELECT * FROM users where id is NULL and is_active = TRUE and is_admin = FALSE",
expected: "SELECT * FROM users where id is ? and is_active = ? and is_admin = ?",
replaceBoolean: true,
replaceNull: true,
},
{
input: "SELECT * FROM users where id = 1 -- this is a comment",
expected: "SELECT * FROM users where id = ? -- this is a comment",
Expand Down Expand Up @@ -206,6 +233,16 @@ func TestObfuscator(t *testing.T) {
expected: `USING - SELECT`,
replaceDigits: true,
},
{
input: "SELECT * FROM users where id = $1",
expected: `SELECT * FROM users where id = $1`,
replacePositionalParameter: false,
},
{
input: "SELECT * FROM users where id = $1",
expected: `SELECT * FROM users where id = ?`,
replacePositionalParameter: true,
},
{
input: `SELECT * FROM "public"."users" where id = 2`,
expected: `SELECT * FROM "public"."users" where id = ?`,
Expand Down Expand Up @@ -335,6 +372,9 @@ func TestObfuscator(t *testing.T) {
t.Run("", func(t *testing.T) {
obfuscator := NewObfuscator(
WithReplaceDigits(tt.replaceDigits),
WithReplacePositionalParameter(tt.replacePositionalParameter),
WithReplaceBoolean(tt.replaceBoolean),
WithReplaceNull(tt.replaceNull),
WithDollarQuotedFunc(tt.dollarQuotedFunc),
)
got := obfuscator.Obfuscate(tt.input, WithDBMS(tt.dbms))
Expand Down
8 changes: 4 additions & 4 deletions sqllexer.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ const (
PUNCTUATION // punctuation
DOLLAR_QUOTED_FUNCTION // dollar quoted function
DOLLAR_QUOTED_STRING // dollar quoted string
NUMBERED_PARAMETER // numbered parameter
POSITIONAL_PARAMETER // numbered parameter
UNKNOWN // unknown token
)

Expand Down Expand Up @@ -121,7 +121,7 @@ func (s *Lexer) Scan() Token {
case ch == '$':
if isDigit(s.lookAhead(1)) {
// if the dollar sign is followed by a digit, then it's a numbered parameter
return s.scanNumberedParameter()
return s.scanPositionalParameter()
}
if s.config.DBMS == DBMSSQLServer && isLetter(s.lookAhead(1)) {
return s.scanIdentifier()
Expand Down Expand Up @@ -393,7 +393,7 @@ func (s *Lexer) scanDollarQuotedString() Token {
return Token{ERROR, s.src[s.start:s.cursor]}
}

func (s *Lexer) scanNumberedParameter() Token {
func (s *Lexer) scanPositionalParameter() Token {
s.start = s.cursor
ch := s.nextBy(2) // consume the dollar sign and the number
for {
Expand All @@ -402,7 +402,7 @@ func (s *Lexer) scanNumberedParameter() Token {
}
ch = s.next()
}
return Token{NUMBERED_PARAMETER, s.src[s.start:s.cursor]}
return Token{POSITIONAL_PARAMETER, s.src[s.start:s.cursor]}
}

func (s *Lexer) scanUnknown() Token {
Expand Down
2 changes: 1 addition & 1 deletion sqllexer_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -312,7 +312,7 @@ func TestLexer(t *testing.T) {
{WS, " "},
{OPERATOR, "="},
{WS, " "},
{NUMBERED_PARAMETER, "$1"},
{POSITIONAL_PARAMETER, "$1"},
},
},
{
Expand Down
8 changes: 8 additions & 0 deletions sqllexer_utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,14 @@ func isSQLKeyword(token Token) bool {
return token.Type == IDENT && keywordsRegex.MatchString(token.Value)
}

func isBoolean(ident string) bool {
return strings.ToUpper(ident) == "TRUE" || strings.ToUpper(ident) == "FALSE"
}

func isNull(ident string) bool {
return strings.ToUpper(ident) == "NULL"
}

func replaceDigits(input string, placeholder string) string {
var builder strings.Builder

Expand Down

0 comments on commit 4be19ca

Please sign in to comment.