Browse Source

fix child-chunk ownership validation (#24374)

Signed-off-by: kenwoodjw <blackxin55+@gmail.com>
tags/1.8.0
kenwoodjw 2 months ago
parent
commit
8a348bea21
No account linked to committer's email address

+ 12
- 2
api/controllers/console/datasets/datasets_segments.py View File

child_chunk_id = str(child_chunk_id) child_chunk_id = str(child_chunk_id)
child_chunk = ( child_chunk = (
db.session.query(ChildChunk) db.session.query(ChildChunk)
.where(ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id)
.where(
ChildChunk.id == str(child_chunk_id),
ChildChunk.tenant_id == current_user.current_tenant_id,
ChildChunk.segment_id == segment.id,
ChildChunk.document_id == document_id,
)
.first() .first()
) )
if not child_chunk: if not child_chunk:
child_chunk_id = str(child_chunk_id) child_chunk_id = str(child_chunk_id)
child_chunk = ( child_chunk = (
db.session.query(ChildChunk) db.session.query(ChildChunk)
.where(ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id)
.where(
ChildChunk.id == str(child_chunk_id),
ChildChunk.tenant_id == current_user.current_tenant_id,
ChildChunk.segment_id == segment.id,
ChildChunk.document_id == document_id,
)
.first() .first()
) )
if not child_chunk: if not child_chunk:

+ 16
- 0
api/controllers/service_api/dataset/segment.py View File

if not segment: if not segment:
raise NotFound("Segment not found.") raise NotFound("Segment not found.")


# validate segment belongs to the specified document
if segment.document_id != document_id:
raise NotFound("Document not found.")

# check child chunk # check child chunk
child_chunk_id = str(child_chunk_id) child_chunk_id = str(child_chunk_id)
child_chunk = SegmentService.get_child_chunk_by_id( child_chunk = SegmentService.get_child_chunk_by_id(
if not child_chunk: if not child_chunk:
raise NotFound("Child chunk not found.") raise NotFound("Child chunk not found.")


# validate child chunk belongs to the specified segment
if child_chunk.segment_id != segment.id:
raise NotFound("Child chunk not found.")

try: try:
SegmentService.delete_child_chunk(child_chunk, dataset) SegmentService.delete_child_chunk(child_chunk, dataset)
except ChildChunkDeleteIndexServiceError as e: except ChildChunkDeleteIndexServiceError as e:
if not segment: if not segment:
raise NotFound("Segment not found.") raise NotFound("Segment not found.")


# validate segment belongs to the specified document
if segment.document_id != document_id:
raise NotFound("Segment not found.")

# get child chunk # get child chunk
child_chunk = SegmentService.get_child_chunk_by_id( child_chunk = SegmentService.get_child_chunk_by_id(
child_chunk_id=child_chunk_id, tenant_id=current_user.current_tenant_id child_chunk_id=child_chunk_id, tenant_id=current_user.current_tenant_id
if not child_chunk: if not child_chunk:
raise NotFound("Child chunk not found.") raise NotFound("Child chunk not found.")


# validate child chunk belongs to the specified segment
if child_chunk.segment_id != segment.id:
raise NotFound("Child chunk not found.")

# validate args # validate args
parser = reqparse.RequestParser() parser = reqparse.RequestParser()
parser.add_argument("content", type=str, required=True, nullable=False, location="json") parser.add_argument("content", type=str, required=True, nullable=False, location="json")

Loading…
Cancel
Save