fix: Preperation for customer encryption at rest (#671)

crates/db/authz.rs

@@ -33,12 +33,7 @@ pub async fn get_permissions(
         setup_user_if_not_already_registered(transaction, authentication).await?
     };
 
-    transaction
-        .query(
-            &format!("SET LOCAL row_level_security.user_id = {}", user_id),
-            &[],
-        )
-        .await?;
+    set_rls_and_encryption_keys(transaction, user_id).await?;
 
     let permissions = queries::users::get_permissions()
         .bind(transaction, &current_team_id)
@@ -57,6 +52,26 @@ pub async fn get_permissions(
     Ok(rbac)
 }
 
+async fn set_rls_and_encryption_keys(
+    transaction: &Transaction<'_>,
+    user_id: i32,
+) -> Result<(), crate::TokioPostgresError> {
+    transaction
+        .query(
+            &format!("SET LOCAL row_level_security.user_id = {}", user_id),
+            &[],
+        )
+        .await?;
+
+    if let Some(key) = crate::customer_keys::get_customer_key() {
+        transaction
+            .query(&format!("SET LOCAL encryption.root_key = '{}'", key), &[])
+            .await?;
+    }
+
+    Ok(())
+}
+
 pub async fn set_row_level_security_user_id(
     transaction: &Transaction<'_>,
     user_id: String,
@@ -66,12 +81,7 @@ pub async fn set_row_level_security_user_id(
         .one()
         .await?;
 
-    transaction
-        .query(
-            &format!("SET LOCAL row_level_security.user_id = {}", user.id),
-            &[],
-        )
-        .await?;
+    set_rls_and_encryption_keys(transaction, user.id).await?;
 
     Ok(user.id)
 }
@@ -86,12 +96,7 @@ pub async fn setup_user_if_not_already_registered(
         .one()
         .await?;
 
-    transaction
-        .query(
-            &format!("SET LOCAL row_level_security.user_id = {}", user_id),
-            &[],
-        )
-        .await?;
+    set_rls_and_encryption_keys(transaction, user_id).await?;
 
     let inserted_org_id = queries::teams::insert_team()
         .bind(transaction)

crates/db/customer_keys.rs

@@ -0,0 +1,6 @@
+use std::env;
+
+// export CUSTOMER_KEY='190a5bf4b3cbb6c0991967ab1c48ab30790af876720f1835cbbf3820f4f5d949'
+pub fn get_customer_key() -> Option<String> {
+    env::var("CUSTOMER_KEY").ok()
+}

crates/db/lib.rs

@@ -1,4 +1,5 @@
 pub mod authz;
+pub mod customer_keys;
 pub mod vector_search;
 
 use std::str::FromStr;

crates/db/migrations/20241030095419_encryption.sql

@@ -0,0 +1,55 @@
+-- migrate:up
+CREATE EXTENSION IF NOT EXISTS pgcrypto;
+
+-- Encrypt function
+CREATE OR REPLACE FUNCTION encrypt_text(data text) RETURNS text AS $$
+DECLARE
+    key text;
+    encrypted text;
+BEGIN
+    -- Retrieve the 'app.root_key' parameter, allowing it to be missing (returns NULL if not set)
+    key := current_setting('encryption.root_key', true);
+
+    IF key IS NULL THEN
+        RETURN data;
+    ELSE
+        BEGIN
+            -- Encrypt the data using pgp_sym_encrypt with 'armor' option
+            encrypted := pgp_sym_encrypt(data, key, 'compress-algo=1, cipher-algo=aes256');
+            RETURN encrypted;
+        EXCEPTION WHEN others THEN
+            -- Return the exception message if encryption fails
+            RETURN SQLERRM;
+        END;
+    END IF;
+END;
+$$ LANGUAGE plpgsql;
+
+
+-- Decrypt function
+
+CREATE OR REPLACE FUNCTION decrypt_text(data text) RETURNS text AS $$
+DECLARE
+    key text;
+    decrypted text;
+BEGIN
+    -- Retrieve the 'app.root_key' parameter, allowing it to be missing (returns NULL if not set)
+    key := current_setting('encryption.root_key', true);
+
+    IF key IS NULL THEN
+        RETURN data;
+    ELSE
+        BEGIN
+            -- Decrypt the data using pgp_sym_decrypt
+            decrypted := pgp_sym_decrypt(data::bytea, key);
+            RETURN decrypted;
+        EXCEPTION WHEN others THEN
+            -- Return the exception message if decryption fails
+            RETURN data;
+        END;
+    END IF;
+END;
+$$ LANGUAGE plpgsql;
+
+-- migrate:down
+

crates/db/queries/prompts.sql

@@ -135,7 +135,7 @@ WHERE
             (p.visibility = 'Private' AND created_by = current_app_user()) 
     )
     AND p.prompt_type = :prompt_type
-ORDER BY updated_at;
+ORDER BY updated_at DESC;
 
 --! prompt : SinglePrompt
 SELECT

crates/k8s-operator/config/grafana.yaml

@@ -24,6 +24,7 @@ data:
     auth_url = $HOSTNAME_URL/oidc/realms/observability/protocol/openid-connect/auth
     token_url = $HOSTNAME_URL/oidc/realms/observability/protocol/openid-connect/token
     api_url = $HOSTNAME_URL/oidc/realms/observability/protocol/openid-connect/userinfo
+    role_attribute_path = contains(roles[*], 'admin') && 'Admin' || contains(roles[*], 'editor') && 'Editor' || 'Viewer'
   datasources.yaml: |
     apiVersion: 1
     datasources:

crates/web-pages/app_layout.rs

@@ -24,7 +24,7 @@ pub enum SideBar {
     RateLimits,
     Switch,
     Team,
-    Training,
+    Security,
 }
 
 impl std::fmt::Display for SideBar {
@@ -168,6 +168,13 @@ pub fn Layout(props: LayoutProps) -> Element {
                                 icon: limits_svg.name,
                                 title: "Rate Limits"
                             }
+                            /***NavItem {
+                                id: SideBar::Security.to_string(),
+                                selected_item_id: props.selected_item.to_string(),
+                                href: super::security::routes::Index { team_id: props.team_id },
+                                icon: security_svg.name,
+                                title: "Security & Encryption"
+                            }**/
                         )
                     }
                 }

crates/web-pages/console/index.rs

@@ -20,6 +20,7 @@ pub fn NewConversation(
             prompt: prompt.clone(),
             title: "AI Chat Console",
             selected_item: SideBar::Console,
+            chats_with_chunks: vec![],
             is_tts_disabled: true,
             lock_console: false,
             header: rsx!(

crates/web-pages/console/layout.rs

@@ -11,15 +11,14 @@ pub fn ConsoleLayout(
     team_id: i32,
     conversation_id: Option<i64>,
     rbac: Rbac,
-    chats_with_chunks: Option<Vec<ChatWithChunks>>,
+    chats_with_chunks: Vec<ChatWithChunks>,
     prompt: SinglePrompt,
     selected_item: SideBar,
     title: String,
     header: Element,
     is_tts_disabled: bool,
     lock_console: bool,
 ) -> Element {
-    // Rerverse it because that's how we display it.
     rsx! {
         Layout {
             section_class: "console flex flex-col justify-start h-[calc(100%-79px)]",
@@ -31,7 +30,7 @@ pub fn ConsoleLayout(
             div {
                 id: "console-panel",
                 class: "h-full flex flex-col",
-                if let Some(chats_with_chunks) = chats_with_chunks {
+                if ! chats_with_chunks.is_empty() {
                     super::console_stream::ConsoleStream {
                         team_id: team_id,
                         chats_with_chunks: chats_with_chunks,

crates/web-pages/lib.rs

@@ -11,6 +11,7 @@ pub mod documents;
 pub mod hero;
 pub mod history;
 pub mod logout_form;
+pub mod security;
 //pub mod model_form;
 pub mod models;
 pub mod pipelines;

crates/web-pages/prompts/prompt_card.rs

@@ -7,10 +7,15 @@ use dioxus::prelude::*;
 
 #[component]
 pub fn PromptCard(team_id: i32, rbac: Rbac, prompt: Prompt) -> Element {
+    let description: String = prompt
+        .description
+        .chars()
+        .filter(|&c| c != '\n' && c != '\t' && c != '\r')
+        .collect();
     rsx! {
         Box {
             class: "cursor-pointer hover:bg-base-200 w-full",
-            drawer_trigger: format!("view-trigger-{}-{}", prompt.id, team_id),
+            modal_trigger: format!("view-trigger-{}-{}", prompt.id, team_id),
             BoxHeader {
                 class: "truncate ellipses flex justify-between p-2",
                 title: "{prompt.name}",
@@ -23,9 +28,10 @@ pub fn PromptCard(team_id: i32, rbac: Rbac, prompt: Prompt) -> Element {
                 div {
                     class: "flex w-full",
                     if let Some(object_id) = prompt.image_icon_object_id {
-                        Avatar {
-                            avatar_size: AvatarSize::Large,
-                            image_src: Image { team_id, id: object_id }.to_string()
+                        img {
+                            width: "96",
+                            height: "96",
+                            src: Image { team_id, id: object_id }.to_string()
                         }
                     } else {
                         Avatar {
@@ -34,13 +40,14 @@ pub fn PromptCard(team_id: i32, rbac: Rbac, prompt: Prompt) -> Element {
                         }
                     }
                     div {
+                        class: "ml-6 flex flex-col space-between",
                         p {
-                            class: "ml-8 text-sm",
-                            "{prompt.description}"
+                            class: "text-sm line-clamp-3",
+                            "{description}"
                         }
                         div {
-                            class: "ml-8 mt-3 text-xs flex justify-center gap-1",
-                            "Last update",
+                            class: "text-xs",
+                            "Last update ",
                             RelativeTime {
                                 format: RelativeTimeFormat::Relative,
                                 datetime: "{prompt.updated_at}"

crates/web-pages/prompts/view_prompt.rs

@@ -14,10 +14,9 @@ pub fn ViewDrawer(team_id: i32, prompt: Prompt, trigger_id: String) -> Element {
     ];
     let has_some_examples = examples.iter().any(|e| e.is_some());
     rsx! {
-        Drawer {
-            label: "{prompt.name}",
+        Modal {
             trigger_id,
-            DrawerBody {
+            ModalBody {
                 div {
                     class: "text-center",
                     if let Some(object_id) = prompt.image_icon_object_id {
@@ -50,7 +49,7 @@ pub fn ViewDrawer(team_id: i32, prompt: Prompt, trigger_id: String) -> Element {
                         "Conversation Starters"
                     }
                     div {
-                        class: "flex flex-col gap-4",
+                        class: "grid grid-cols-2 gap-x-1.5 gap-y-2",
                         for example in examples {
                             if let Some(example) = example {
                                 if ! example.is_empty() {
@@ -64,12 +63,12 @@ pub fn ViewDrawer(team_id: i32, prompt: Prompt, trigger_id: String) -> Element {
                         }
                     }
                 }
-            }
-            DrawerFooter {
-                a {
-                    class: "btn btn-primary btn-sm w-full",
-                    href: crate::routes::prompts::NewChat{team_id, prompt_id: prompt.id}.to_string(),
-                    "Start a Chat"
+                ModalAction {
+                    a {
+                        class: "btn btn-primary btn-sm w-full",
+                        href: crate::routes::prompts::NewChat{team_id, prompt_id: prompt.id}.to_string(),
+                        "Start a Chat"
+                    }
                 }
             }
         }