CVE-2024-5128: Improper Access Control (CWE-284)
value: An Insecure Direct Object Reference (IDOR) vulnerability was identified in lunary-ai/lunary, affecting versions up to and including 1.2.2. This vulnerability allows unauthorized users to view, update, or delete any dataset_prompt or dataset_prompt_variation within any dataset or project. The issue stems from improper access control checks in the dataset management endpoints, where direct references to object IDs are not adequately secured against unauthorized access. This vulnerability was fixed in version 1.2.25.
CWE description: The product does not restrict or incorrectly restricts access to a resource from an unauthorized actor.
Published: 2024-06-06T19:16Z
Repository: https://github.com/lunary-ai/lunary
Language: TypeScript
Files changed: packages/backend/src/api/v1/datasets/index.ts
Commit date: 2024-05-19 15:03:37 -0300
Commit message: fix: api security (#330)


Diff: diff --git a/packages/backend/src/api/v1/datasets/index.ts b/packages/backend/src/api/v1/datasets/index.ts
index 476b7a39..2474af23 100644
--- a/packages/backend/src/api/v1/datasets/index.ts
+++ b/packages/backend/src/api/v1/datasets/index.ts
@@ -19,37 +19,32 @@ datasets.get("/", checkAccess("datasets", "list"), async (ctx: Context) => {
     await sql`select * from dataset d where project_id = ${projectId} order by created_at desc`
 
   ctx.body = rows
 })
 
-// Can be public route
-datasets.get(
-  "/:identifier",
-  // checkAccess("datasets", "read"),
-  async (ctx: Context) => {
-    const { projectId } = ctx.state
-    const { identifier } = ctx.params
+datasets.get("/:identifier", async (ctx: Context) => {
+  const { projectId } = ctx.state
+  const { identifier } = ctx.params
 
-    const isUUID = validateUUID(identifier)
+  const isUUID = validateUUID(identifier)
 
-    if (isUUID) {
-      // For frontend
-      const datasetId = identifier
-      const dataset = await getDatasetById(datasetId, projectId)
+  if (isUUID) {
+    // For frontend
+    const datasetId = identifier
+    const dataset = await getDatasetById(datasetId, projectId)
 
-      ctx.body = dataset
-      return
-    } else {
-      // For SDK
-      const slug = identifier
-      const dataset = await getDatasetBySlug(slug, projectId)
+    ctx.body = dataset
+    return
+  } else {
+    // For SDK
+    const slug = identifier
+    const dataset = await getDatasetBySlug(slug, projectId)
 
-      ctx.body = dataset
-      return
-    }
-  },
-)
+    ctx.body = dataset
+    return
+  }
+})
 
 const DEFAULT_PROMPT = {
   chat: [
     {
       role: "system",
@@ -79,11 +74,10 @@ datasets.post("/", checkAccess("datasets", "create"), async (ctx: Context) => {
       ownerId: userId,
       projectId,
     })} returning *
   `
 
-  // insert 1 prompt and 1 variation
   const [prompt] = await sql`insert into dataset_prompt
     ${sql({
       datasetId: dataset.id,
       messages: DEFAULT_PROMPT[format],
     })}
@@ -122,20 +116,19 @@ datasets.patch(
 
 datasets.delete(
   "/:id",
   checkAccess("datasets", "delete"),
   async (ctx: Context) => {
-    const { id } = ctx.params
+    const { id: datasetId } = ctx.params
     const { projectId } = ctx.state
 
-    await sql`delete from dataset where id = ${id} and project_id = ${projectId}`
+    await sql`delete from dataset where id = ${datasetId} and project_id = ${projectId}`
 
     ctx.status = 200
   },
 )
 
-// Create prompt
 datasets.post(
   "/prompts",
   checkAccess("datasets", "update"),
   async (ctx: Context) => {
     const { projectId } = ctx.state
@@ -155,60 +148,73 @@ datasets.post(
       messages: messages || DEFAULT_PROMPT[format],
     })}
     returning *
   `
 
-    // insert 1 variation
-    await sql`insert into dataset_prompt_variation
-    ${sql({
-      promptId: prompt.id,
-      variables: {},
-      idealOutput: idealOutput ? lastMsg(idealOutput) : "",
-    })}
-    returning *
-  `
+    await sql`
+      insert into dataset_prompt_variation
+        ${sql({
+          promptId: prompt.id,
+          variables: {},
+          idealOutput: idealOutput ? lastMsg(idealOutput) : "",
+        })}
+      returning *
+    `
 
     ctx.body = prompt
   },
 )
 
-// return array of prompts and variations
 datasets.get(
   "/prompts/:id",
   checkAccess("datasets", "read"),
   async (ctx: Context) => {
     const { id } = ctx.params as { id: string }
 
     const [prompt] = await sql`
-    select * from dataset_prompt where id = ${id} order by created_at asc
+      select * from dataset_prompt where id = ${id} order by created_at asc
     `
 
     const variations = await sql`
-    select * from dataset_prompt_variation where prompt_id = ${id} order by created_at asc
+      select * from dataset_prompt_variation where prompt_id = ${id} order by created_at asc
     `
 
     prompt.variations = variations
 
     ctx.body = prompt
   },
 )
 
-// add route ot to delete dataset_prompt and dataset_prompt_variation
 datasets.delete(
   "/prompts/:id",
   checkAccess("datasets", "update"),
   async (ctx: Context) => {
-    const { id } = ctx.params
+    const { id: promptId } = ctx.params
+
+    const [datasetPrompt] = await sql`
+      select
+        *
+      from
+        dataset_prompt dp
+        left join dataset d on dp.dataset_id = d.id
+        left join project p on d.project_id = p.id
+      where
+        p.org_id = ${ctx.state.orgId} 
+        and dp.id = ${promptId}
+    `
 
-    await sql`delete from dataset_prompt where id = ${id}`
-    await sql`delete from dataset_prompt_variation where prompt_id = ${id}`
+    if (!datasetPrompt) {
+      ctx.throw(401, "You do not have access to this ressource.")
+    }
+
+    await sql`delete from dataset_prompt where id = ${promptId}`
+    await sql`delete from dataset_prompt_variation where prompt_id = ${promptId}`
 
     ctx.status = 200
   },
 )
 
-// Update prompt
 datasets.patch(
   "/prompts/:id",
   checkAccess("datasets", "update"),
   async (ctx: Context) => {
     const { id } = ctx.params
@@ -243,19 +249,34 @@ datasets.get(
 
 datasets.delete(
   "/variations/:id",
   checkAccess("datasets", "update"),
   async (ctx: Context) => {
-    const { id } = ctx.params
+    const { id: variationId } = ctx.params
+
+    const [promptVariation] = await sql`
+      select
+        *
+      from
+        dataset_prompt_variation dpv
+        left join dataset_prompt dp on dpv.prompt_id = dp.id
+        left join dataset d on dp.dataset_id = d.id
+        left join project p on d.project_id = p.id
+      where
+        p.org_id = ${ctx.state.orgId} 
+        and dpv.id = ${variationId}
+    `
+    if (!promptVariation) {
+      ctx.throw(401, "You do not have access to this ressource.")
+    }
 
-    await sql`delete from dataset_prompt_variation where id = ${id}`
+    await sql`delete from dataset_prompt_variation where id = ${variationId}`
 
     ctx.status = 200
   },
 )
 
-// Update variation
 datasets.patch(
   "/variations/:variationId",
   checkAccess("datasets", "update"),
   async (ctx: Context) => {
     const { variationId } = ctx.params
@@ -277,11 +298,10 @@ datasets.patch(
 
     ctx.body = variation
   },
 )
 
-// Create variation
 datasets.post(
   "/variations",
   checkAccess("datasets", "update"),
   async (ctx: Context) => {
     const { promptId, variables, idealOutput } = ctx.request.body as {
@@ -289,19 +309,19 @@ datasets.post(
       variables: any
       idealOutput: string
     }
 
     const [variation] = await sql`insert into dataset_prompt_variation
-  ${sql(
-    clearUndefined({
-      promptId,
-      variables,
-      idealOutput,
-    }),
-  )}
-  returning *
-  `
+      ${sql(
+        clearUndefined({
+          promptId,
+          variables,
+          idealOutput,
+        }),
+      )}
+      returning *
+    `
 
     ctx.body = variation
   },
 )
 

Vulnerable code: // Can be public route
datasets.get(
  "/:identifier",
  // checkAccess("datasets", "read"),
  async (ctx: Context) => {
    const { projectId } = ctx.state
    const { identifier } = ctx.params
    const isUUID = validateUUID(identifier)
    if (isUUID) {
      // For frontend
      const datasetId = identifier
      const dataset = await getDatasetById(datasetId, projectId)
      ctx.body = dataset
      return
    } else {
      // For SDK
      const slug = identifier
      const dataset = await getDatasetBySlug(slug, projectId)
      ctx.body = dataset
      return
    }
  },
)
  // insert 1 prompt and 1 variation
    const { id } = ctx.params
    await sql`delete from dataset where id = ${id} and project_id = ${projectId}`
// Create prompt
    // insert 1 variation
    await sql`insert into dataset_prompt_variation
    ${sql({
      promptId: prompt.id,
      variables: {},
      idealOutput: idealOutput ? lastMsg(idealOutput) : "",
    })}
    returning *
  `
// return array of prompts and variations
    select * from dataset_prompt where id = ${id} order by created_at asc
    select * from dataset_prompt_variation where prompt_id = ${id} order by created_at asc
// add route ot to delete dataset_prompt and dataset_prompt_variation
    const { id } = ctx.params
    await sql`delete from dataset_prompt where id = ${id}`
    await sql`delete from dataset_prompt_variation where prompt_id = ${id}`
// Update prompt
    const { id } = ctx.params
    await sql`delete from dataset_prompt_variation where id = ${id}`
// Update variation
// Create variation
  ${sql(
    clearUndefined({
      promptId,
      variables,
      idealOutput,
    }),
  )}
  returning *
  `
Fixed code: datasets.get("/:identifier", async (ctx: Context) => {
  const { projectId } = ctx.state
  const { identifier } = ctx.params
  const isUUID = validateUUID(identifier)
  if (isUUID) {
    // For frontend
    const datasetId = identifier
    const dataset = await getDatasetById(datasetId, projectId)
    ctx.body = dataset
    return
  } else {
    // For SDK
    const slug = identifier
    const dataset = await getDatasetBySlug(slug, projectId)
    ctx.body = dataset
    return
  }
})
    const { id: datasetId } = ctx.params
    await sql`delete from dataset where id = ${datasetId} and project_id = ${projectId}`
    await sql`
      insert into dataset_prompt_variation
        ${sql({
          promptId: prompt.id,
          variables: {},
          idealOutput: idealOutput ? lastMsg(idealOutput) : "",
        })}
      returning *
    `
      select * from dataset_prompt where id = ${id} order by created_at asc
      select * from dataset_prompt_variation where prompt_id = ${id} order by created_at asc
    const { id: promptId } = ctx.params

    const [datasetPrompt] = await sql`
      select
        *
      from
        dataset_prompt dp
        left join dataset d on dp.dataset_id = d.id
        left join project p on d.project_id = p.id
      where
        p.org_id = ${ctx.state.orgId} 
        and dp.id = ${promptId}
    `
    if (!datasetPrompt) {
      ctx.throw(401, "You do not have access to this ressource.")
    }

    await sql`delete from dataset_prompt where id = ${promptId}`
    await sql`delete from dataset_prompt_variation where prompt_id = ${promptId}`
    const { id: variationId } = ctx.params

    const [promptVariation] = await sql`
      select
        *
      from
        dataset_prompt_variation dpv
        left join dataset_prompt dp on dpv.prompt_id = dp.id
        left join dataset d on dp.dataset_id = d.id
        left join project p on d.project_id = p.id
      where
        p.org_id = ${ctx.state.orgId} 
        and dpv.id = ${variationId}
    `
    if (!promptVariation) {
      ctx.throw(401, "You do not have access to this ressource.")
    }
    await sql`delete from dataset_prompt_variation where id = ${variationId}`
      ${sql(
        clearUndefined({
          promptId,
          variables,
          idealOutput,
        }),
      )}
      returning *
    `