ScottLogic / prompt-injection

Application which investigates defensive measures against prompt injection attacks on an LLM, with a focus on the exposure of external tools.
MIT License
15 stars 10 forks source link

Make api endpoints type safe at runtime #814

Open pmarsh-scottlogic opened 7 months ago

pmarsh-scottlogic commented 7 months ago

Since Typescript only operates at compile time, and our endpoints receive arbitrary JSON at runtime, a client could pass in JSON with all the wrong types and our code will happily consume it. Take handleAddToChatHistory as an example.

function handleAddToChatHistory(req: OpenAiAddHistoryRequest, res: Response) {
    const infoMessage = req.body.message;
    const chatMessageType = req.body.chatMessageType;
    const level = req.body.level;
    if (
        infoMessage &&
        chatMessageType &&
        level !== undefined &&
        level >= LEVEL_NAMES.LEVEL_1
    ) {
        req.session.levelState[level].chatHistory = pushMessageToHistory(
            req.session.levelState[level].chatHistory,
            {
                chatMessageType,
                infoMessage,
            } as ChatMessage
        );
        res.send();
    } else {
        res.status(400);
        res.send();
    }
}

where

type OpenAiAddHistoryRequest = Request<
    never,
    never,
    {
        chatMessageType?: CHAT_MESSAGE_TYPE;
        message?: string;
        level?: LEVEL_NAMES;
    },
    never,
    never
>;

At the moment we check that stuff exists, but we don't check the type. Here's a nonsense request body that would be happily consumed at runtime (resulting in a 500 error):

{
    "level": 1000,
    "message": true,
    "chatMessageType": "hello!"
}