Handle privileges more strictly

2026-07-07 14:10:43 +00:00 · 2024-10-15 11:47:55 +03:00
parent 4e15a68ffd
commit c6c29f86da
7 changed files with 170 additions and 11 deletions
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -38,6 +38,7 @@ use crate::checker::create_availability_check_data;
 use crate::local_proxy;
 use crate::logger::inlinify;
 use crate::pg_helpers::*;
+use crate::privilege::Privilege;
 use crate::spec::*;
 use crate::sync_sk::{check_if_synced, ping_safekeeper};
 use crate::{config, extension_server};
@@ -1372,7 +1373,7 @@ LIMIT 100",
        &self,
        db_name: &str,
        schema_name: &str,
-        privilege: &str,
+        privileges: &[Privilege],
        role_name: &str,
    ) -> Result<()> {
        let mut conf = Config::from_str(self.connstr.as_str()).unwrap();
@@ -1384,7 +1385,18 @@ LIMIT 100",

        let query = "GRANT $1 ON SCHEMA $2 TO $3";
        db_client
-            .execute(query, &[&schema_name, &privilege, &role_name])
+            .execute(
+                query,
+                &[
+                    &privileges
+                        .iter()
+                        .map(|p| p.as_str())
+                        .collect::<Vec<&str>>()
+                        .join(", "),
+                    &schema_name,
+                    &role_name,
+                ],
+            )
            .context(format!("Failed to execute query: {}", query))?;

        Ok(())
--- a/compute_tools/src/http/api.rs
+++ b/compute_tools/src/http/api.rs
@@ -2,6 +2,7 @@ use std::convert::Infallible;
 use std::net::IpAddr;
 use std::net::Ipv6Addr;
 use std::net::SocketAddr;
+use std::str::FromStr;
 use std::sync::Arc;
 use std::thread;

@@ -9,9 +10,10 @@ use crate::catalog::SchemaDumpError;
 use crate::catalog::{get_database_schema, get_dbs_and_roles};
 use crate::compute::forward_termination_signal;
 use crate::compute::{ComputeNode, ComputeState, ParsedSpec};
+use crate::privilege::Privilege;
 use compute_api::requests::{ConfigurationRequest, SetRoleGrantsRequest};
 use compute_api::responses::{
-    ComputeStatus, ComputeStatusResponse, GenericAPIError, SetRoleGrantsResult,
+    ComputeStatus, ComputeStatusResponse, GenericAPIError, SetRoleGrantsResponse,
 };

 use anyhow::Result;
@@ -181,14 +183,37 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body

            let request = hyper::body::to_bytes(req.into_body()).await.unwrap();
            let request = serde_json::from_slice::<SetRoleGrantsRequest>(&request).unwrap();
+
+            let privileges: Result<Vec<Privilege>, _> = request
+                .privileges
+                .iter()
+                .map(|p| Privilege::from_str(p.as_str()))
+                .collect();
+            let privileges = match privileges {
+                Ok(privs) => privs,
+                Err(_) => {
+                    let msg = format!("Invalid privilege in request: {:?}", &request.privileges);
+                    error!(msg);
+                    return Response::new(Body::from(msg));
+                }
+            };
+
            let res = compute.set_role_grants(
                &request.database,
                &request.schema,
-                &request.privilege,
+                &privileges,
                &request.role,
            );
            match res {
-                Ok(_) => Response::new(Body::from("true")),
+                Ok(_) => render_json(Body::from(
+                    serde_json::to_string(&SetRoleGrantsResponse {
+                        database: request.database,
+                        schema: request.schema,
+                        role: request.role,
+                        privileges: request.privileges,
+                    })
+                    .unwrap(),
+                )),
                Err(e) => {
                    error!("set_role_grants failed: {}", e);
                    Response::new(Body::from(e.to_string()))
--- a/compute_tools/src/http/openapi_spec.yaml
+++ b/compute_tools/src/http/openapi_spec.yaml
@@ -127,6 +127,34 @@ paths:
              schema:
                $ref: "#/components/schemas/GenericError"

+  /grants:
+    post:
+      tags:
+        - Grants
+      summary: Apply grants to the database.
+      description: ""
+      operationId: setRoleGrants
+      requestBody:
+        description: Grants request.
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: SetRoleGrantsRequest
+      responses:
+        200:
+          description: Grants applied.
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/SetRoleGrantsResponse"
+        500:
+          description: Error occurred during grants application.
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/GenericError"
+
  /check_writability:
    post:
      tags:
@@ -431,23 +459,55 @@ components:
      type: object
      required:
        - database
+        - schema
+        - privileges
        - role
-        - grants
      properties:
        database:
          type: string
          description: Database name.
          example: "neondb"
+        schema:
+          type: string
+          description: Schema name.
+          example: "public"
+        privileges:
+          type: array
+          items:
+            type: string
+          description: List of privileges to set.
+          example: ["SELECT", "INSERT"]
        role:
          type: string
          description: Role name.
          example: "neon"
-        grants:
+
+    SetRoleGrantsResponse:
+      type: object
+      required:
+        - database
+        - schema
+        - privileges
+        - role
+      properties:
+        database:
+          type: string
+          description: Database name.
+          example: "neondb"
+        schema:
+          type: string
+          description: Schema name.
+          example: "public"
+        privileges:
          type: array
          items:
            type: string
-          description: List of grants to set.
+          description: List of privileges set.
          example: ["SELECT", "INSERT"]
+        role:
+          type: string
+          description: Role name.
+          example: "neon"

    #
    # Errors
--- a/compute_tools/src/lib.rs
+++ b/compute_tools/src/lib.rs
@@ -22,6 +22,7 @@ mod migration;
 pub mod monitor;
 pub mod params;
 pub mod pg_helpers;
+pub mod privilege;
 pub mod spec;
 pub mod swap;
 pub mod sync_sk;
--- a/compute_tools/src/privilege.rs
+++ b/compute_tools/src/privilege.rs
@@ -0,0 +1,58 @@
+use std::str::FromStr;
+
+#[derive(Debug)]
+pub enum Privilege {
+    Select,
+    Insert,
+    Update,
+    Delete,
+    Truncate,
+    References,
+    Trigger,
+    Usage,
+    Create,
+    Connect,
+    Temporary,
+    Execute,
+}
+
+impl FromStr for Privilege {
+    type Err = &'static str;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_uppercase().as_str() {
+            "SELECT" => Ok(Privilege::Select),
+            "INSERT" => Ok(Privilege::Insert),
+            "UPDATE" => Ok(Privilege::Update),
+            "DELETE" => Ok(Privilege::Delete),
+            "TRUNCATE" => Ok(Privilege::Truncate),
+            "REFERENCES" => Ok(Privilege::References),
+            "TRIGGER" => Ok(Privilege::Trigger),
+            "USAGE" => Ok(Privilege::Usage),
+            "CREATE" => Ok(Privilege::Create),
+            "CONNECT" => Ok(Privilege::Connect),
+            "TEMPORARY" => Ok(Privilege::Temporary),
+            "EXECUTE" => Ok(Privilege::Execute),
+            _ => Err("Invalid privilege"),
+        }
+    }
+}
+
+impl Privilege {
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            Privilege::Select => "SELECT",
+            Privilege::Insert => "INSERT",
+            Privilege::Update => "UPDATE",
+            Privilege::Delete => "DELETE",
+            Privilege::Truncate => "TRUNCATE",
+            Privilege::References => "REFERENCES",
+            Privilege::Trigger => "TRIGGER",
+            Privilege::Usage => "USAGE",
+            Privilege::Create => "CREATE",
+            Privilege::Connect => "CONNECT",
+            Privilege::Temporary => "TEMPORARY",
+            Privilege::Execute => "EXECUTE",
+        }
+    }
+}
--- a/libs/compute_api/src/requests.rs
+++ b/libs/compute_api/src/requests.rs
@@ -17,6 +17,6 @@ pub struct ConfigurationRequest {
 pub struct SetRoleGrantsRequest {
    pub database: String,
    pub schema: String,
-    pub privilege: String,
+    pub privileges: Vec<String>,
    pub role: String,
 }
--- a/libs/compute_api/src/responses.rs
+++ b/libs/compute_api/src/responses.rs
@@ -170,6 +170,9 @@ pub struct InstalledExtensions {
 }

 #[derive(Clone, Debug, Default, Serialize)]
-pub struct SetRoleGrantsResult {
-    pub extension: String,
+pub struct SetRoleGrantsResponse {
+    pub database: String,
+    pub schema: String,
+    pub privileges: Vec<String>,
+    pub role: String,
 }