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.
52 lines
1.3 KiB
Go
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)
|
|
}
|
|
})
|
|
}
|
|
}
|