mirror of
https://github.com/vxcontrol/pentagi.git
synced 2026-08-26 21:16:45 +00:00
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:
@@ -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",
|
||||
})
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user