feat(tests): enhance knowledge and flowfile tests with containment barriers

Added new test cases to validate containment barriers in the `ResolvePulledStagedTarget` and `ZipRelativePaths` functions, ensuring that paths escaping the designated directories are properly rejected. Updated the `knowledge_test.go` to include scenarios for handling nil embedder and error propagation during document creation, improving overall test coverage and robustness.
This commit is contained in:
Dmitry Ng
2026-05-29 15:42:11 +03:00
parent ce06a0d26a
commit 8feb7bf311
3 changed files with 162 additions and 52 deletions
+105 -52
View File
@@ -27,6 +27,7 @@ import (
type mockDB struct {
database.Querier // nil — panics if an unexpected method is called
insertKnowledge func(ctx context.Context, arg database.InsertKnowledgeDocumentParams) (string, error)
getKnowledge func(ctx context.Context, uuid string) (database.GetKnowledgeDocumentRow, error)
getUserKnowledge func(ctx context.Context, arg database.GetUserKnowledgeDocumentParams) (database.GetUserKnowledgeDocumentRow, error)
listAll func(ctx context.Context) ([]database.ListAllKnowledgeDocumentsRow, error)
@@ -39,6 +40,9 @@ type mockDB struct {
searchUserKnowledge func(ctx context.Context, arg database.SearchUserKnowledgeDocumentsParams) ([]database.SearchUserKnowledgeDocumentsRow, error)
}
func (m *mockDB) InsertKnowledgeDocument(ctx context.Context, arg database.InsertKnowledgeDocumentParams) (string, error) {
return m.insertKnowledge(ctx, arg)
}
func (m *mockDB) GetKnowledgeDocument(ctx context.Context, uuid string) (database.GetKnowledgeDocumentRow, error) {
return m.getKnowledge(ctx, uuid)
}
@@ -1437,56 +1441,100 @@ func TestCreateDocument(t *testing.T) {
ctx := context.Background()
const userID = int64(11)
t.Run("store nil returns error immediately", func(t *testing.T) {
ks := &knowledgeStore{store: nil}
// insertOK returns a mockDB whose InsertKnowledgeDocument returns a fixed id.
insertOK := func(id string) *mockDB {
return &mockDB{
insertKnowledge: func(_ context.Context, _ database.InsertKnowledgeDocumentParams) (string, error) {
return id, nil
},
}
}
t.Run("embedder nil returns error", func(t *testing.T) {
ks := &knowledgeStore{db: insertOK("id"), embedder: nil}
_, err := ks.CreateDocument(ctx, userID, model.CreateKnowledgeDocumentInput{
DocType: model.KnowledgeDocTypeAnswer, Content: "x", Question: "q",
})
if err == nil {
t.Fatal("expected error when store is nil")
t.Fatal("expected error when embedder is nil")
}
})
t.Run("store AddDocuments error propagates", func(t *testing.T) {
vs := &mockVectorStore{
addDocumentsFn: func(_ context.Context, _ []schema.Document, _ ...vectorstores.Option) ([]string, error) {
return nil, errors.New("vector db error")
t.Run("embedder unavailable returns error", func(t *testing.T) {
ks := &knowledgeStore{db: insertOK("id"), embedder: &mockEmbedder{available: false}}
_, err := ks.CreateDocument(ctx, userID, model.CreateKnowledgeDocumentInput{
DocType: model.KnowledgeDocTypeAnswer, Content: "x", Question: "q",
})
if err == nil {
t.Fatal("expected error when embedder unavailable")
}
})
t.Run("EmbedDocuments error propagates", func(t *testing.T) {
ks := &knowledgeStore{
db: insertOK("id"),
embedder: &mockEmbedder{
available: true,
embedDocumentsFn: func(_ context.Context, _ []string) ([][]float32, error) {
return nil, errors.New("embed error")
},
},
}
ks := &knowledgeStore{store: vs, newKnp: newPublisherFactory(&mockPublisher{})}
_, err := ks.CreateDocument(ctx, userID, model.CreateKnowledgeDocumentInput{
DocType: model.KnowledgeDocTypeGuide, Content: "c", Question: "q",
})
if err == nil {
t.Fatal("expected error from AddDocuments")
t.Fatal("expected error from EmbedDocuments")
}
})
t.Run("AddDocuments returning no IDs is an error", func(t *testing.T) {
vs := &mockVectorStore{
addDocumentsFn: func(_ context.Context, _ []schema.Document, _ ...vectorstores.Option) ([]string, error) {
return []string{}, nil // empty slice
t.Run("embedder returning empty vectors is an error", func(t *testing.T) {
ks := &knowledgeStore{
db: insertOK("id"),
embedder: &mockEmbedder{
available: true,
embedDocumentsFn: func(_ context.Context, _ []string) ([][]float32, error) {
return [][]float32{}, nil // empty slice
},
},
}
ks := &knowledgeStore{store: vs, newKnp: newPublisherFactory(&mockPublisher{})}
_, err := ks.CreateDocument(ctx, userID, model.CreateKnowledgeDocumentInput{
DocType: model.KnowledgeDocTypeAnswer, Content: "c", Question: "q",
})
if err == nil {
t.Fatal("expected error for empty IDs")
t.Fatal("expected error for empty vectors")
}
})
t.Run("db insert error propagates", func(t *testing.T) {
db := &mockDB{
insertKnowledge: func(_ context.Context, _ database.InsertKnowledgeDocumentParams) (string, error) {
return "", errors.New("constraint error")
},
}
ks := &knowledgeStore{
db: db,
embedder: &mockEmbedder{available: true},
newKnp: newPublisherFactory(&mockPublisher{}),
}
_, err := ks.CreateDocument(ctx, userID, model.CreateKnowledgeDocumentInput{
DocType: model.KnowledgeDocTypeAnswer, Content: "c", Question: "q",
})
if err == nil {
t.Fatal("expected error from db insert")
}
})
t.Run("success: all fields set, manual=true, user_id present, event published", func(t *testing.T) {
pub := &mockPublisher{}
var capturedDocs []schema.Document
vs := &mockVectorStore{
addDocumentsFn: func(_ context.Context, docs []schema.Document, _ ...vectorstores.Option) ([]string, error) {
capturedDocs = docs
return []string{"new-uuid"}, nil
var gotParams database.InsertKnowledgeDocumentParams
db := &mockDB{
insertKnowledge: func(_ context.Context, arg database.InsertKnowledgeDocumentParams) (string, error) {
gotParams = arg
return "new-uuid", nil
},
}
ks := &knowledgeStore{store: vs, newKnp: newPublisherFactory(pub)}
ks := &knowledgeStore{db: db, embedder: &mockEmbedder{available: true}, newKnp: newPublisherFactory(pub)}
input := model.CreateKnowledgeDocumentInput{
DocType: model.KnowledgeDocTypeCode,
@@ -1510,26 +1558,36 @@ func TestCreateDocument(t *testing.T) {
if !doc.Manual {
t.Fatal("manual must be true for manually created docs")
}
// Metadata stored in pgvector
if len(capturedDocs) != 1 {
t.Fatal("expected exactly one document passed to AddDocuments")
}
meta := capturedDocs[0].Metadata
if meta["user_id"] != userID {
t.Fatalf("user_id in metadata: want %d, got %v", userID, meta["user_id"])
}
if meta["manual"] != true {
t.Fatal("manual flag must be true in metadata")
}
if meta["code_lang"] != "go" {
t.Fatal("code_lang missing from metadata")
}
// Returned model must carry the creator's UserID
if doc.UserID != userID {
t.Fatalf("doc.UserID: want %d, got %d", userID, doc.UserID)
}
if doc.CodeLang == nil || *doc.CodeLang != "go" {
t.Fatal("code_lang missing from returned doc")
}
// Persisted document text (trimmed) and embedding literal
if gotParams.Document.String != "func main() {}" {
t.Fatalf("persisted document not trimmed: %q", gotParams.Document.String)
}
emb, ok := gotParams.Embedding.(string)
if !ok || emb == "" || emb[0] != '[' {
t.Fatalf("embedding must be a vector literal, got %v", gotParams.Embedding)
}
// Metadata stored in pgvector cmetadata
meta := parseMeta(string(gotParams.Cmetadata))
if meta.UserID != userID {
t.Fatalf("user_id in metadata: want %d, got %d", userID, meta.UserID)
}
if !meta.Manual {
t.Fatal("manual flag must be true in metadata")
}
if meta.CodeLang != "go" {
t.Fatal("code_lang missing from metadata")
}
if meta.Description != "a Go main" {
t.Fatalf("description missing from metadata: %q", meta.Description)
}
// Event
if len(pub.createdDocs) != 1 || pub.createdDocs[0].ID != "new-uuid" {
@@ -1541,32 +1599,27 @@ func TestCreateDocument(t *testing.T) {
})
t.Run("content is trimmed of whitespace", func(t *testing.T) {
var capturedContent string
vs := &mockVectorStore{
addDocumentsFn: func(_ context.Context, docs []schema.Document, _ ...vectorstores.Option) ([]string, error) {
capturedContent = docs[0].PageContent
return []string{"id"}, nil
var gotParams database.InsertKnowledgeDocumentParams
db := &mockDB{
insertKnowledge: func(_ context.Context, arg database.InsertKnowledgeDocumentParams) (string, error) {
gotParams = arg
return "id", nil
},
}
ks := &knowledgeStore{store: vs, newKnp: newPublisherFactory(&mockPublisher{})}
ks := &knowledgeStore{db: db, embedder: &mockEmbedder{available: true}, newKnp: newPublisherFactory(&mockPublisher{})}
_, err := ks.CreateDocument(ctx, userID, model.CreateKnowledgeDocumentInput{
DocType: model.KnowledgeDocTypeAnswer, Content: " trimmed ", Question: "q",
})
if err != nil {
t.Fatal(err)
}
if capturedContent != "trimmed" {
t.Fatalf("expected trimmed content, got %q", capturedContent)
if gotParams.Document.String != "trimmed" {
t.Fatalf("expected trimmed content, got %q", gotParams.Document.String)
}
})
t.Run("optional fields absent means nil in model", func(t *testing.T) {
vs := &mockVectorStore{
addDocumentsFn: func(_ context.Context, _ []schema.Document, _ ...vectorstores.Option) ([]string, error) {
return []string{"id"}, nil
},
}
ks := &knowledgeStore{store: vs, newKnp: newPublisherFactory(&mockPublisher{})}
ks := &knowledgeStore{db: insertOK("id"), embedder: &mockEmbedder{available: true}, newKnp: newPublisherFactory(&mockPublisher{})}
doc, err := ks.CreateDocument(ctx, userID, model.CreateKnowledgeDocumentInput{
DocType: model.KnowledgeDocTypeAnswer, Content: "c", Question: "q",
})
+13
View File
@@ -354,6 +354,12 @@ func ResolvePulledStagedTarget(stagingDir, cacheRelPath string) string {
}
for _, candidate := range candidates {
// Defense-in-depth containment barrier: ignore any candidate that
// resolves outside the staging directory. cacheRelPath is sanitized by
// the caller today, but this keeps the function safe under refactoring.
if !IsWithinDir(candidate, stagingDir) {
continue
}
if _, err := os.Lstat(candidate); err == nil {
return candidate
}
@@ -470,6 +476,13 @@ func ZipRelativePaths(w io.Writer, baseDir string, relPaths []string) error {
for _, relPath := range relPaths {
localPath := filepath.Join(baseDir, filepath.FromSlash(relPath))
// Defense-in-depth containment barrier: callers are expected to pass
// pre-validated cache-relative paths, but never operate on a path that
// resolves outside baseDir regardless of caller behaviour.
if !IsWithinDir(localPath, baseDir) {
continue
}
info, err := os.Lstat(localPath)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
+44
View File
@@ -304,6 +304,19 @@ func TestResolvePulledStagedTarget(t *testing.T) {
assert.Equal(t, target, ResolvePulledStagedTarget(stagingDir, "etc/nginx/nginx.conf"))
})
t.Run("escaping cache path is rejected", func(t *testing.T) {
root := t.TempDir()
stagingDir := filepath.Join(root, "staging")
require.NoError(t, os.MkdirAll(stagingDir, 0755))
// A file outside the staging dir that an escaping path would resolve to.
require.NoError(t, os.WriteFile(filepath.Join(root, "evil.conf"), []byte("evil"), 0644))
// The first candidate (Join(stagingDir, "../evil.conf")) escapes the
// staging dir and must be ignored by the containment barrier; the
// second candidate (basename "evil.conf") does not exist in staging.
assert.Equal(t, "", ResolvePulledStagedTarget(stagingDir, "../evil.conf"))
})
}
func TestWriteUploadsTar(t *testing.T) {
@@ -683,6 +696,37 @@ func TestZipRelativePaths(t *testing.T) {
require.NoError(t, err)
assert.Empty(t, zr.File)
})
t.Run("escaping relPath is skipped", func(t *testing.T) {
root := t.TempDir()
base := filepath.Join(root, "flow-1-data")
require.NoError(t, os.MkdirAll(filepath.Join(base, "uploads"), 0755))
require.NoError(t, os.WriteFile(filepath.Join(base, "uploads", "a.txt"), []byte("alpha"), 0644))
// Secret file outside baseDir that an escaping relPath would resolve to.
require.NoError(t, os.WriteFile(filepath.Join(root, "secret.txt"), []byte("secret"), 0644))
var buf bytes.Buffer
err := ZipRelativePaths(&buf, base, []string{
"uploads/a.txt",
"../secret.txt", // escapes baseDir: must be skipped by the barrier
})
require.NoError(t, err)
zr, err := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
require.NoError(t, err)
contents := map[string]string{}
for _, f := range zr.File {
rc, err := f.Open()
require.NoError(t, err)
data, _ := io.ReadAll(rc)
rc.Close()
contents[f.Name] = string(data)
}
assert.Equal(t, "alpha", contents["uploads/a.txt"])
assert.NotContains(t, contents, "../secret.txt", "escaping paths must be excluded")
assert.NotContains(t, contents, "secret.txt", "escaping paths must be excluded")
})
}
func TestZipDirectory(t *testing.T) {