diff --git a/backend/app/routers/questions.py b/backend/app/routers/questions.py index 2ea800f..8739f41 100644 --- a/backend/app/routers/questions.py +++ b/backend/app/routers/questions.py @@ -26,7 +26,8 @@ from app.models.quiz import Quiz from app.models.user import User from app.models.favorite import Favorite from app.models.folder import QuestionFolderQuestion -from app.services import article_service +from app.services import article_service, file_intake, vision_service +from app.services.ai_service import get_configured_model from app.services.search_service import hybrid_ids, hybrid_question_ids, rerank_ids from app.services.question_figures import figures_for as _figures_for from app.services.prepared_session import prepare_session @@ -712,6 +713,31 @@ def build_test_from_description( return {**created, "matched": len(matched)} +def _image_reader(db: Session): + """How an uploaded image becomes words, or None if it cannot. + + The tool model is the one an administrator chose for jobs like this. With + none configured there is no reader, and `file_intake` says so in a sentence + that names what to upload instead — rather than a 500 from a model call + that was never going to happen. + """ + chosen = get_configured_model(db, "tool") + if not chosen or not chosen[0]: + return None + model_id, api_key = chosen + + def read_image(data: bytes, media_type: str) -> str: + image = vision_service.prepare(data, media_type, caption="an uploaded page") + if image is None: + raise file_intake.Rejected("That image could not be read.") + try: + return vision_service.describe([image], model_id, api_key)[0] + except vision_service.VisionUnavailable as unavailable: + raise file_intake.Rejected(str(unavailable)) + + return read_image + + @router.post("/builder/from-upload") def build_test_from_upload( file: UploadFile = File(...), @@ -730,26 +756,24 @@ def build_test_from_upload( if not 1 <= count <= MAX_MATCHED_QUESTIONS: raise HTTPException(400, f"Choose between 1 and {MAX_MATCHED_QUESTIONS} questions") - raw = file.file.read(MAX_UPLOAD_BYTES + 1) - if len(raw) > MAX_UPLOAD_BYTES: - raise HTTPException(413, f"Keep the file under {MAX_UPLOAD_BYTES // (1024 * 1024)} MB") - if not raw: - raise HTTPException(400, "That file is empty") + # Size, then type, then text — and the type is sniffed from the bytes, not + # taken from the name. `handout.pdf` holding something else is not a PDF, + # and the old branch here decided by extension and fell through to + # "decode whatever it is as UTF-8" for everything else. + try: + raw = file_intake.read(file) + except file_intake.Rejected as refusal: + raise HTTPException(413 if "over" in str(refusal) else 400, str(refusal)) - name = (file.filename or "").lower() - if name.endswith(".pdf"): - try: - import fitz + kind = file_intake.kind_of(raw) + if kind is None: + raise HTTPException(415, f"Upload {file_intake.ACCEPT_HUMAN}.") - with fitz.open(stream=raw, filetype="pdf") as document: - text = " ".join(page.get_text() for page in document)[:MAX_QUERY_CHARS] - except Exception: - raise HTTPException(400, "That PDF could not be read") - else: - try: - text = raw.decode("utf-8", errors="ignore")[:MAX_QUERY_CHARS] - except Exception: - raise HTTPException(400, "Upload a PDF or a text file") + try: + text = file_intake.text_from(raw, kind, MAX_QUERY_CHARS, + describe=_image_reader(db)) + except file_intake.Rejected as refusal: + raise HTTPException(400, str(refusal)) if len(text.strip()) < 40: raise HTTPException(400, "That file has too little text to match against") diff --git a/backend/app/services/file_intake.py b/backend/app/services/file_intake.py new file mode 100644 index 0000000..cdaad4d --- /dev/null +++ b/backend/app/services/file_intake.py @@ -0,0 +1,164 @@ +"""What a person may upload, decided by the bytes rather than by the name. + +Three questions, in this order, and all three before anything else touches the +file: + +1. **Is it small enough.** Read one byte past the cap and refuse if that byte + exists, so an enormous file is never held in memory to be measured. +2. **What is it, actually.** Sniffed from the leading bytes. An extension is a + claim the uploader makes, and `report.pdf` holding a shell script is not a + PDF; the type we act on is the one the file's own header states. +3. **What text is in it.** PDF and DOCX are parsed for their text; an image is + read by the tool model. Nothing is executed, ever, and none of it is stored. + +On code injection specifically, because it is the question people mean when +they ask about uploads: there is no path from an uploaded file to code that +runs. The extracted text is used as a *search query* — passed as a bound +parameter to Postgres, never concatenated into SQL — and as the content of a +message to a language model. It is never rendered as HTML (the reader escapes +raw HTML in Markdown), never written to disk here, never passed to a shell, and +never `eval`'d. A DOCX is a zip file, so it is also checked against the one +attack that shape does allow: a member that expands to far more than the cap. +""" +import io +import re +import zipfile + +#: Two megabytes. Not a technical limit — a PDF of a lecture is comfortably +#: under it and a scanned textbook chapter is not, and the second one is a +#: different feature. +MAX_BYTES = 2 * 1024 * 1024 + +#: How far a DOCX may expand once unzipped. A 40 KB zip that becomes a +#: gigabyte of XML is the only interesting thing about the format's shape. +MAX_UNZIPPED = 40 * 1024 * 1024 + +PDF = "pdf" +DOCX = "docx" +IMAGE = "image" + +IMAGE_TYPES = { + b"\x89PNG\r\n\x1a\n": "image/png", + b"\xff\xd8\xff": "image/jpeg", + b"GIF87a": "image/gif", + b"GIF89a": "image/gif", +} + +#: What the upload control offers, and what the error message names. Kept here +#: so the two cannot drift apart. +ACCEPT = ".pdf,.docx,.png,.jpg,.jpeg,.gif,.webp" +ACCEPT_HUMAN = "a PDF, a Word document (.docx), or an image" + + +class Rejected(ValueError): + """The file may not be taken, with the sentence to show the person.""" + + +def media_type(data: bytes) -> str | None: + """The image type these bytes are, or None if they are not an image.""" + for signature, mime in IMAGE_TYPES.items(): + if data.startswith(signature): + return mime + # RIFF....WEBP — the size sits between the two markers. + if data[:4] == b"RIFF" and data[8:12] == b"WEBP": + return "image/webp" + return None + + +def kind_of(data: bytes) -> str | None: + """pdf, docx, image — or None for anything else at all.""" + if data.startswith(b"%PDF-"): + return PDF + if media_type(data): + return IMAGE + # Every OOXML file is a zip whose first member is [Content_Types].xml. A + # plain zip, a jar and an xlsx all start the same way, so the manifest is + # what separates a Word document from them. + if data[:2] == b"PK" and b"word/document.xml" in data[:4096] + data[-4096:]: + return DOCX + if data[:2] == b"PK": + try: + with zipfile.ZipFile(io.BytesIO(data)) as bundle: + if "word/document.xml" in bundle.namelist(): + return DOCX + except zipfile.BadZipFile: + return None + return None + + +def read(upload) -> bytes: + """The file's bytes, refused if there are too many of them. + + One byte past the cap is read deliberately: its presence is the test, and + reading it costs one byte rather than the whole file. + """ + data = upload.file.read(MAX_BYTES + 1) + if len(data) > MAX_BYTES: + raise Rejected( + f"That file is over {MAX_BYTES // (1024 * 1024)} MB. Upload a " + "smaller one — a single chapter or handout rather than a whole book.") + if not data: + raise Rejected("That file is empty.") + return data + + +def _docx_text(data: bytes, limit: int) -> str: + try: + with zipfile.ZipFile(io.BytesIO(data)) as bundle: + total = sum(info.file_size for info in bundle.infolist()) + if total > MAX_UNZIPPED: + raise Rejected("That Word document unpacks to far more than its size suggests.") + xml = bundle.read("word/document.xml").decode("utf-8", errors="ignore") + except Rejected: + raise + except Exception: + raise Rejected("That Word document could not be read.") + # Paragraph and line breaks become newlines before the tags go, so the text + # does not run together into one sentence. + xml = re.sub(r"", "\n", xml) + xml = re.sub(r"]*/?>", "\n", xml) + xml = re.sub(r"<[^>]+>", "", xml) + return clean(xml)[:limit] + + +def _pdf_text(data: bytes, limit: int) -> str: + try: + import fitz + + with fitz.open(stream=data, filetype="pdf") as document: + return clean(" ".join(page.get_text() for page in document))[:limit] + except Exception: + raise Rejected("That PDF could not be read.") + + +def clean(text: str) -> str: + """Printable text, with the control characters gone. + + Not a security measure — nothing downstream interprets them — but a file + full of NULs makes a query that matches nothing and an error nobody can + read, and stripping them is one line. + """ + text = text.replace("\x00", " ") + text = re.sub(r"[\x01-\x08\x0b\x0c\x0e-\x1f\x7f]", " ", text) + return re.sub(r"[ \t]{2,}", " ", text).strip() + + +def text_from(data: bytes, kind: str, limit: int, *, describe=None) -> str: + """The words in this file. + + `describe` is how an image becomes text: a callable taking (bytes, mime) + and returning what the tool model read in it. Passed in rather than + imported so this module stays a file-handling module — and so the caller, + which knows whether a tool model is even configured, decides. + """ + if kind == PDF: + return _pdf_text(data, limit) + if kind == DOCX: + return _docx_text(data, limit) + if kind == IMAGE: + if describe is None: + raise Rejected( + "Reading an image needs a tool model, and no model is configured " + "for that job. Upload a PDF or a Word document instead.") + return clean(describe(data, media_type(data) or "image/png"))[:limit] + raise Rejected(f"Upload {ACCEPT_HUMAN}.") diff --git a/backend/tests/test_ai_mode_matching.py b/backend/tests/test_ai_mode_matching.py index dc19363..8be66b0 100644 --- a/backend/tests/test_ai_mode_matching.py +++ b/backend/tests/test_ai_mode_matching.py @@ -6,6 +6,7 @@ import os os.environ["DATABASE_URL"] = "sqlite:///:memory:" import io +import zipfile from datetime import datetime import unittest from unittest.mock import patch @@ -92,10 +93,30 @@ class AiModeMatchingTests(unittest.TestCase): self.assertEqual(self.client.post("/questions/builder/describe", json={"text": "short"}).status_code, 422) - def upload(self, content=b"febrile seizure in a toddler with fever, and jaundice in a newborn infant", - filename="notes.txt", **data): + NOTES = "febrile seizure in a toddler with fever, and jaundice in a newborn infant" + + @staticmethod + def as_docx(text: str) -> bytes: + """The smallest thing that is really a Word document. + + Built rather than fixtured because what is being tested is that the + endpoint reads the *bytes*: a .docx name over a text file is refused, + and this is what makes the difference visible. + """ + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as bundle: + bundle.writestr("[Content_Types].xml", "") + bundle.writestr("word/document.xml", + f"{text}" + "") + return buffer.getvalue() + + def upload(self, content=None, filename="handout.docx", **data): + if content is None: + content = self.as_docx(self.NOTES) return self.client.post("/questions/builder/from-upload", - files={"file": (filename, io.BytesIO(content), "text/plain")}, + files={"file": (filename, io.BytesIO(content), + "application/octet-stream")}, data={"count": "5", "mode": "learning", **data}) def test_an_upload_is_matched_against_the_bank_and_not_stored(self): @@ -109,10 +130,27 @@ class AiModeMatchingTests(unittest.TestCase): oversized = b"x" * (questions.MAX_UPLOAD_BYTES + 10) self.assertEqual(self.upload(content=oversized).status_code, 413) self.assertEqual(self.upload(content=b"").status_code, 400) - self.assertEqual(self.upload(content=b"too short").status_code, 400) + self.assertEqual(self.upload(content=self.as_docx("too short")).status_code, 400) self.assertEqual(self.upload(count="99").status_code, 400) self.assertEqual(self.upload(mode="nonsense").status_code, 400) + def test_only_the_three_kinds_are_taken_and_the_name_is_not_evidence(self): + """A file is what its bytes say it is. + + The endpoint used to decide by extension and fall through to "decode + whatever this is as UTF-8" for everything else, so a shell script, an + HTML page or a CSV all became a search query. Each is now refused with + 415, and renaming one to .pdf does not change that. + """ + for content in [b"#!/bin/sh\nrm -rf /", b"", + b"id,name\n1,two\n", self.NOTES.encode()]: + self.assertEqual(self.upload(content=content, filename="notes.pdf").status_code, 415) + # And an image with no tool model configured is told why, not 500'd. + png = b"\x89PNG\r\n\x1a\n" + b"0" * 64 + answer = self.upload(content=png, filename="slide.png") + self.assertEqual(answer.status_code, 400) + self.assertIn("tool model", answer.json()["detail"]) + def test_a_matched_test_is_capped(self): self.assertLessEqual(questions.MAX_MATCHED_QUESTIONS, 30) diff --git a/backend/tests/test_file_intake.py b/backend/tests/test_file_intake.py new file mode 100644 index 0000000..0be5ff6 --- /dev/null +++ b/backend/tests/test_file_intake.py @@ -0,0 +1,99 @@ +"""What may be uploaded, and what the bytes have to say about it. + +Run: DATABASE_URL=sqlite:// PYTHONPATH=backend python -m unittest discover -s backend/tests +""" +import os +os.environ.setdefault("DATABASE_URL", "sqlite://") + +import io +import unittest +import zipfile + +from app.services import file_intake + + +def docx(text: str = "Febrile seizures are common between six months and five years.") -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as bundle: + bundle.writestr("[Content_Types].xml", "") + bundle.writestr( + "word/document.xml", + f"{text}") + return buffer.getvalue() + + +class Upload: + """The two attributes the reader touches on a Starlette UploadFile.""" + + def __init__(self, data: bytes): + self.file = io.BytesIO(data) + + +PNG = b"\x89PNG\r\n\x1a\n" + b"0" * 40 + + +class SniffingTests(unittest.TestCase): + def test_the_bytes_decide_not_the_extension(self): + self.assertEqual(file_intake.kind_of(b"%PDF-1.7\nstuff"), file_intake.PDF) + self.assertEqual(file_intake.kind_of(docx()), file_intake.DOCX) + self.assertEqual(file_intake.kind_of(PNG), file_intake.IMAGE) + self.assertEqual(file_intake.kind_of(b"RIFF\x00\x00\x00\x00WEBPmore"), file_intake.IMAGE) + + def test_everything_else_is_refused(self): + # A shell script called report.pdf is not a PDF; a zip that is not a + # Word document is not one either. + self.assertIsNone(file_intake.kind_of(b"#!/bin/sh\nrm -rf /")) + self.assertIsNone(file_intake.kind_of(b"")) + self.assertIsNone(file_intake.kind_of(b"id,name\n1,two\n")) + plain_zip = io.BytesIO() + with zipfile.ZipFile(plain_zip, "w") as bundle: + bundle.writestr("notes.txt", "hello") + self.assertIsNone(file_intake.kind_of(plain_zip.getvalue())) + self.assertIsNone(file_intake.kind_of(b"")) + + +class SizeTests(unittest.TestCase): + def test_over_the_cap_is_refused_by_the_byte_past_it(self): + with self.assertRaises(file_intake.Rejected) as refusal: + file_intake.read(Upload(b"x" * (file_intake.MAX_BYTES + 1))) + self.assertIn("2 MB", str(refusal.exception)) + + def test_exactly_the_cap_is_allowed_and_empty_is_not(self): + self.assertEqual(len(file_intake.read(Upload(b"x" * file_intake.MAX_BYTES))), + file_intake.MAX_BYTES) + with self.assertRaises(file_intake.Rejected): + file_intake.read(Upload(b"")) + + +class TextTests(unittest.TestCase): + def test_a_word_document_gives_up_its_words(self): + text = file_intake.text_from(docx(), file_intake.DOCX, 4000) + self.assertIn("Febrile seizures are common", text) + # Tags never reach the text, so no markup can be smuggled through one. + self.assertNotIn("", text) + + def test_a_zip_bomb_is_refused_before_it_is_read(self): + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as bundle: + bundle.writestr("[Content_Types].xml", "") + bundle.writestr("word/document.xml", "" + ("a" * (file_intake.MAX_UNZIPPED + 10))) + with self.assertRaises(file_intake.Rejected): + file_intake.text_from(buffer.getvalue(), file_intake.DOCX, 4000) + + def test_control_characters_go_and_the_words_stay(self): + self.assertEqual(file_intake.clean("a\x00b\x07 c"), "a b c") + + def test_an_image_without_a_reader_says_so_rather_than_failing(self): + with self.assertRaises(file_intake.Rejected) as refusal: + file_intake.text_from(PNG, file_intake.IMAGE, 4000, describe=None) + self.assertIn("tool model", str(refusal.exception)) + + def test_an_image_with_a_reader_is_read(self): + text = file_intake.text_from( + PNG, file_intake.IMAGE, 4000, + describe=lambda data, mime: f"A slide about croup ({mime})") + self.assertEqual(text, "A slide about croup (image/png)") + + +if __name__ == "__main__": + unittest.main()