sqli_test.go

package libinjection

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestIsSQLi(t *testing.T) {
	result, fingerprint := IsSQLi("-1' and 1=1 union/* foo */select load_file('/etc/passwd')--")
	fmt.Println("=========result==========: ", result)
	fmt.Println("=======fingerprint=======: ", string(fingerprint))
}

const (
	fingerprints = "fingerprints"
	folding      = "folding"
	tokens       = "tokens"

	sectionTest     = "--TEST--"
	sectionInput    = "--INPUT--"
	sectionExpected = "--EXPECTED--"

	whitespace = " \t\n\r"
)

func printTokenString(t *sqliToken) string {
	out := ""
	if t.strOpen != 0 {
		out += string(t.strOpen)
	}
	out += string(t.val[:t.len])
	if t.strClose != 0 {
		out += string(t.strClose)
	}
	return out
}

func printToken(t *sqliToken) string {
	out := ""
	out += string(t.category)
	out += " "
	switch t.category {
	case 's':
		out += printTokenString(t)
	case 'v':
		vc := t.count
		if vc == 1 {
			out += "@"
		} else if vc == 2 {
			out += "@@"
		}
		out += printTokenString(t)
	default:
		out += string(t.val[:t.len])
	}
	return strings.TrimRight(out, "\n\r")
}

func getToken(state *sqliState, i int) *sqliToken {
	if i < 0 || i > maxTokens {
		panic("token got error!")
	}
	return &state.tokenVec[i]
}

// sections defines the expected order of test file sections.
var sections = [3]string{sectionTest, sectionInput, sectionExpected}

func readTestData(filename string) map[string]string {
	f, err := os.Open(filename)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	data := make(map[string]string)
	count := 0
	state := ""

	br := bufio.NewReaderSize(f, 8192)
	for {
		line, _, err := br.ReadLine()
		if err == io.EOF {
			break
		}
		if err != nil {
			panic(err)
		}

		str := string(line)
		if count < len(sections) && strings.TrimSpace(str) == sections[count] {
			state = sections[count]
			count++
			continue
		}
		if state == "" {
			panic(fmt.Sprintf("readTestData: unexpected content before first section in %s", filename))
		}
		data[state] += str + "\n"
	}
	if count != len(sections) {
		panic(fmt.Sprintf("readTestData: missing sections in %s (got %d/%d)", filename, count, len(sections)))
	}
	// Right-trim only (matching C's modp_rtrim), not left-trim
	for _, s := range sections {
		data[s] = strings.TrimRight(data[s], whitespace)
	}
	return data
}

func runSQLiTest(t testing.TB, data map[string]string, filename string, flag string, sqliFlag int) {
	t.Helper()

	var (
		actual = ""
		state  = new(sqliState)
	)

	sqliInit(state, data[sectionInput], sqliFlag)

	switch flag {
	case fingerprints:
		result, fingerprints := IsSQLi(data[sectionInput])
		if result {
			actual = string(fingerprints[:])
		}

	case folding:
		numTokens := state.fold()
		for i := 0; i < numTokens; i++ {
			actual += printToken(getToken(state, i)) + "\n"
		}

	case tokens:
		for state.tokenize() {
			actual += printToken(state.current) + "\n"
		}
	}

	actual = strings.TrimSpace(actual)
	if actual != data[sectionExpected] {
		t.Errorf("FILE: (%s)\nINPUT: (%s)\nEXPECTED: (%s)\nGOT: (%s)\n",
			filename, data[sectionInput], data[sectionExpected], actual)
	}
}

func TestSQLiDriver(t *testing.T) {
	baseDir := "tests"
	dir, err := os.ReadDir(baseDir)
	if err != nil {
		t.Fatal(err)
	}

	for _, fi := range dir {
		p := filepath.Join(baseDir, fi.Name())
		data := readTestData(p)
		switch {
		case strings.Contains(fi.Name(), "-sqli-"):
			t.Run(fi.Name(), func(t *testing.T) {
				runSQLiTest(t, data, p, fingerprints, 0)
			})
		case strings.Contains(fi.Name(), "-folding-"):
			t.Run(fi.Name(), func(t *testing.T) {
				runSQLiTest(t, data, p, folding, sqliFlagQuoteNone|sqliFlagSQLAnsi)
			})
		case strings.Contains(fi.Name(), "-tokens_mysql-"):
			t.Run(fi.Name(), func(t *testing.T) {
				runSQLiTest(t, data, p, tokens, sqliFlagQuoteNone|sqliFlagSQLMysql)
			})
		case strings.Contains(fi.Name(), "-tokens-"):
			t.Run(fi.Name(), func(t *testing.T) {
				runSQLiTest(t, data, p, tokens, sqliFlagQuoteNone|sqliFlagSQLAnsi)
			})
		}
	}
}

type testCaseSQLI struct {
	name string
	data map[string]string
}

func BenchmarkSQLiDriver(b *testing.B) {
	baseDir := "./tests/"
	dir, err := os.ReadDir(baseDir)
	if err != nil {
		b.Fatal(err)
	}

	cases := struct {
		sqli        []testCaseSQLI
		folding     []testCaseSQLI
		tokensMySQL []testCaseSQLI
		tokens      []testCaseSQLI
	}{}

	for _, fi := range dir {
		p := filepath.Join(baseDir, fi.Name())
		data := readTestData(p)
		tc := testCaseSQLI{
			name: fi.Name(),
			data: data,
		}
		switch {
		case strings.Contains(fi.Name(), "-sqli-"):
			cases.sqli = append(cases.sqli, tc)
		case strings.Contains(fi.Name(), "-folding-"):
			cases.folding = append(cases.folding, tc)
		case strings.Contains(fi.Name(), "-tokens-"):
			cases.tokens = append(cases.tokens, tc)
		}
	}

	b.Run("sqli", func(b *testing.B) {
		b.ReportAllocs()
		for i := 0; i < b.N; i++ {
			for _, tc := range cases.sqli {
				tt := tc
				runSQLiTest(b, tt.data, tt.name, fingerprints, 0)
			}
		}
	})

	b.Run("folding", func(b *testing.B) {
		b.ReportAllocs()
		for i := 0; i < b.N; i++ {
			for _, tc := range cases.folding {
				tt := tc
				runSQLiTest(b, tt.data, tt.name, folding, sqliFlagQuoteNone|sqliFlagSQLAnsi)
			}
		}
	})

	b.Run("tokens", func(b *testing.B) {
		b.ReportAllocs()
		for i := 0; i < b.N; i++ {
			for _, tc := range cases.tokens {
				tt := tc
				runSQLiTest(b, tt.data, tt.name, tokens, sqliFlagQuoteNone|sqliFlagSQLAnsi)
			}
		}
	})
}

// TestFlag2Delimiter exercises all three return paths of flag2Delimiter.
func TestFlag2Delimiter(t *testing.T) {
	tests := []struct {
		name string
		flag int
		want byte
	}{
		{name: "single quote flag", flag: sqliFlagQuoteSingle, want: byteSingle},
		{name: "double quote flag", flag: sqliFlagQuoteDouble, want: byteDouble},
		{name: "no quote flag", flag: sqliFlagQuoteNone, want: byteNull},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := flag2Delimiter(tt.flag); got != tt.want {
				t.Errorf("flag2Delimiter(%d) = %v, want %v", tt.flag, got, tt.want)
			}
		})
	}
}

// TestResetWithZeroFlags exercises the flags==0 branch in reset.
func TestResetWithZeroFlags(t *testing.T) {
	s := new(sqliState)
	sqliInit(s, "SELECT 1", sqliFlagQuoteNone|sqliFlagSQLAnsi)
	// Calling reset(0) should default to sqliFlagQuoteNone | sqliFlagSQLAnsi.
	s.reset(0)
	if s.flags != sqliFlagQuoteNone|sqliFlagSQLAnsi {
		t.Errorf("reset(0) flags = %d, want %d", s.flags, sqliFlagQuoteNone|sqliFlagSQLAnsi)
	}
}

// TestParseBStringEarlyReturn exercises the early-return branches in
// parseBString.
func TestParseBStringEarlyReturn(t *testing.T) {
	tests := []struct {
		name  string
		input string
	}{
		// First early return: fewer than 3 chars or second byte is not a single quote.
		{name: "B with no following quote", input: "SELECT B0101"},
		{name: "b at end of input", input: "SELECT b"},
		// Second early return: B' followed by binary digits that are not closed by '.
		{name: "b-string unclosed with non-binary char", input: "b'0x"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// IsSQLi exercises parseBString internally; verify these inputs are
			// not classified as SQLi (i.e., early-return / fallback behavior is sane).
			got, fingerprint := IsSQLi(tt.input)
			if got {
				t.Errorf("IsSQLi(%q) = true, fingerprint=%q; want false", tt.input, fingerprint)
			}
		})
	}
}

// TestParseQStringCoreDelimiters exercises the { and < delimiter cases in
// parseQStringCore and the ch < 33 early-return case.
func TestParseQStringCoreDelimiters(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		wantSQLi bool
	}{
		// { delimiter maps to }; not detected as SQLi
		{name: "Q-string with brace delimiter", input: "Q'{hello}'", wantSQLi: false},
		// < delimiter maps to >; in single-quote mode the token boundary makes it look like "sos"
		{name: "Q-string with angle-bracket delimiter", input: "Q'<hello>'", wantSQLi: true},
		{name: "q-string with brace delimiter", input: "q'{hello}'", wantSQLi: false},
		{name: "q-string with angle-bracket delimiter", input: "q'<hello>'", wantSQLi: true},
		// ch < 33: Q' followed by a control character (ASCII < 33) falls back to parseWord
		{name: "Q-string with control-char delimiter", input: "Q'\x01hello'", wantSQLi: false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, fingerprint := IsSQLi(tt.input)
			if got != tt.wantSQLi {
				t.Errorf("IsSQLi(%q) = %v (fingerprint=%q), want %v", tt.input, got, fingerprint, tt.wantSQLi)
			}
		})
	}
}

// TestNotWhitelistEdgeCases exercises rarely-hit branches in notWhitelist
// to improve coverage of the false-positive-reduction logic.
func TestNotWhitelistEdgeCases(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		wantSQLi bool
	}{
		// "1 UNION" with exactly 2 stats tokens should not be SQLi
		{name: "1 union not sqli", input: "1 UNION", wantSQLi: false},
		// sp_password in input with comment fingerprint should be SQLi
		{name: "sp_password bypass", input: "1 -- sp_password", wantSQLi: true},
		// "1c" fingerprint with folding (statsTokens > 2) triggers detection
		{name: "1c with folding", input: "1+1/*comment*/", wantSQLi: true},
		// "1c" with whitespace before block comment (ch <= 32 branch)
		{name: "1c with whitespace before block comment", input: "1 /*comment*/", wantSQLi: true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, _ := IsSQLi(tt.input)
			if got != tt.wantSQLi {
				t.Errorf("IsSQLi(%q) = %v, want %v", tt.input, got, tt.wantSQLi)
			}
		})
	}
}

// TestNotWhitelistDirectState exercises the rarely-reached branches in the
// Number+Comment block of notWhitelist that cannot be triggered via IsSQLi.
// We call notWhitelist directly with a crafted sqliState.
func TestNotWhitelistDirectState(t *testing.T) {
	// Construct a synthetic "1c" fingerprint state where:
	//   - tokenVec[0] is a number of length 1
	//   - tokenVec[1] is a block-style comment (val starts with '/')
	//   - statsTokens == 2 (no folding)
	// The char at s.input[1] is '-', which bypasses the '/' and whitespace checks
	// and reaches the double-dash check or the final return-false path.
	makeState := func(input string) *sqliState {
		s := new(sqliState)
		sqliInit(s, input, sqliFlagQuoteNone|sqliFlagSQLAnsi)
		s.fingerprint = string([]byte{sqliTokenTypeNumber, sqliTokenTypeComment})
		s.statsTokens = 2
		s.tokenVec[0].category = sqliTokenTypeNumber
		s.tokenVec[0].len = 1
		s.tokenVec[1].category = sqliTokenTypeComment
		s.tokenVec[1].val = "/"
		s.tokenVec[1].len = 1
		return s
	}

	t.Run("double-dash pattern is SQLi", func(t *testing.T) {
		// input[1]=='-' and input[2]=='-': ch=='-' and next=='-' → return true
		s := makeState("1--")
		if !s.notWhitelist() {
			t.Error("expected notWhitelist to return true for dash-dash pattern")
		}
	})

	t.Run("single-dash followed by non-dash is not SQLi", func(t *testing.T) {
		// input[1]=='-' and input[2]!='-': ch=='-' but next!='-' → return false
		s := makeState("1-/")
		if s.notWhitelist() {
			t.Error("expected notWhitelist to return false for non-SQLi pattern")
		}
	})
}