From 8feb7bf311d6310fc124d44b8151dc3789db4d1a Mon Sep 17 00:00:00 2001 From: Dmitry Ng <19asdek91@gmail.com> Date: Fri, 29 May 2026 15:42:11 +0300 Subject: [PATCH] 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. --- .../pkg/database/knowledge/knowledge_test.go | 157 ++++++++++++------ backend/pkg/flowfiles/files.go | 13 ++ backend/pkg/flowfiles/files_test.go | 44 +++++ 3 files changed, 162 insertions(+), 52 deletions(-) diff --git a/backend/pkg/database/knowledge/knowledge_test.go b/backend/pkg/database/knowledge/knowledge_test.go index 2e2e9c0f..4e8be017 100644 --- a/backend/pkg/database/knowledge/knowledge_test.go +++ b/backend/pkg/database/knowledge/knowledge_test.go @@ -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", }) diff --git a/backend/pkg/flowfiles/files.go b/backend/pkg/flowfiles/files.go index f1766664..80b2b466 100644 --- a/backend/pkg/flowfiles/files.go +++ b/backend/pkg/flowfiles/files.go @@ -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) { diff --git a/backend/pkg/flowfiles/files_test.go b/backend/pkg/flowfiles/files_test.go index fe18ecec..b82d99e1 100644 --- a/backend/pkg/flowfiles/files_test.go +++ b/backend/pkg/flowfiles/files_test.go @@ -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) {