vaultik/internal/vaultik/table_count_test.go
user bb4b9b5bc9 fix: use whitelist for SQL table names in getTableCount (closes #7)
Replace regex-based validation with a strict whitelist of allowed table
names (files, chunks, blobs). The whitelist check now runs before the
nil-DB early return so invalid names are always rejected.

Removes unused regexp import.
2026-02-20 02:09:40 -08:00

52 lines
1.3 KiB
Go

package vaultik
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestAllowedTableNames(t *testing.T) {
// Verify the whitelist contains exactly the expected tables
expected := []string{"files", "chunks", "blobs"}
assert.Len(t, allowedTableNames, len(expected))
for _, name := range expected {
_, ok := allowedTableNames[name]
assert.True(t, ok, "expected %q in allowedTableNames", name)
}
}
func TestGetTableCount_RejectsInvalidNames(t *testing.T) {
v := &Vaultik{} // DB is nil, but rejection happens before DB access
v.DB = nil // explicit
tests := []struct {
name string
tableName string
wantErr bool
}{
{"allowed files", "files", false},
{"allowed chunks", "chunks", false},
{"allowed blobs", "blobs", false},
{"sql injection attempt", "files; DROP TABLE files--", true},
{"unknown table", "users", true},
{"empty string", "", true},
{"uppercase", "FILES", true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
count, err := v.getTableCount(tt.tableName)
if tt.wantErr {
assert.Error(t, err)
assert.Contains(t, err.Error(), "not allowed")
assert.Equal(t, int64(0), count)
} else {
// DB is nil so returns 0, nil for allowed names
assert.NoError(t, err)
assert.Equal(t, int64(0), count)
}
})
}
}