start implementing artifact security

Author: SamSaffron

app/controllers/discourse_ai/ai_bot/artifacts_controller.rb

@@ -4,14 +4,20 @@ module DiscourseAi
   module AiBot
     class ArtifactsController < ApplicationController
       requires_plugin DiscourseAi::PLUGIN_NAME
+      before_action :require_site_settings!
 
       skip_before_action :preload_json, :check_xhr, only: %i[show]
 
       def show
         artifact = AiArtifact.find(params[:id])
 
         post = Post.find_by(id: artifact.post_id)
-        raise Discourse::NotFound unless post && guardian.can_see?(post)
+        if artifact.metadata&.dig("public")
+          # no guardian needed
+        else
+          raise Discourse::NotFound if !post&.topic&.private_message?
+          raise Discourse::NotFound if !guardian.can_see?(post)
+        end
 
         # Prepare the HTML document
         html = <<~HTML
@@ -39,6 +45,16 @@ def show
         # Render the content
         render html: html.html_safe, layout: false, content_type: "text/html"
       end
+
+      private
+
+      def require_site_settings!
+        if !SiteSetting.discourse_ai_enabled ||
+            !SiteSetting.ai_artifact_security.in?(%w[lax strict])
+          raise Discourse::NotFound
+        end
+      end
+
     end
   end
 end

app/models/ai_artifact.rb

@@ -3,6 +3,32 @@
 class AiArtifact < ActiveRecord::Base
   belongs_to :user
   belongs_to :post
+
+  def self.iframe_for(id)
+    <<~HTML
+      <iframe sandbox="allow-scripts allow-forms" height="600px" src='#{url(id)}' frameborder="0" width="100%"></iframe>
+    HTML
+  end
+
+  def self.url(id)
+    Discourse.base_url + "/discourse-ai/ai-bot/artifacts/#{id}"
+  end
+
+  def self.share_publicly(id:, post:)
+    artifact = AiArtifact.find_by(id: id)
+    if artifact&.post&.topic&.id == post.topic.id
+      artifact.update!(metadata: { public: true })
+    end
+  end
+
+  def self.unshare_publicly(id:)
+    artifact = AiArtifact.find_by(id: id)
+    artifact&.update!(metadata: { public: false })
+  end
+
+  def url
+    self.class.url(id)
+  end
 end
 
 # == Schema Information

app/models/shared_ai_conversation.rb

@@ -34,6 +34,12 @@ def self.share_conversation(user, target, max_posts: DEFAULT_MAX_POSTS)
 
   def self.destroy_conversation(conversation)
     conversation.destroy
+
+    maybe_topic = conversation.target
+    if maybe_topic.is_a?(Topic)
+      AiArtifact.where(post: maybe_topic.posts).update_all(metadata: { public: false })
+    end
+
     ::Jobs.enqueue(
       :shared_conversation_adjust_upload_security,
       target_id: conversation.target_id,
@@ -165,7 +171,7 @@ def self.build_conversation_data(topic, max_posts: DEFAULT_MAX_POSTS, include_us
             id: post.id,
             user_id: post.user_id,
             created_at: post.created_at,
-            cooked: post.cooked,
+            cooked: cook_artifacts(post),
           }
 
           mapped[:persona] = persona if ai_bot_participant&.id == post.user_id
@@ -175,6 +181,22 @@ def self.build_conversation_data(topic, max_posts: DEFAULT_MAX_POSTS, include_us
     }
   end
 
+  def self.cook_artifacts(post)
+    html = post.cooked
+    return html if !["lax", "strict"].include?(SiteSetting.ai_artifact_security)
+
+    doc = Nokogiri::HTML5.fragment(html)
+    doc.css("div.ai-artifact").each do |node|
+      id = node["data-ai-artifact-id"].to_i
+      if id > 0
+        AiArtifact.share_publicly(id: id, post: post)
+        node.replace(AiArtifact.iframe_for(id))
+      end
+    end
+
+    doc.to_s
+  end
+
   private
 
   def populate_user_info!(posts)

assets/stylesheets/modules/ai-bot/common/ai-artifact.scss

@@ -86,6 +86,7 @@ html.ai-artifact-expanded {
     position: fixed;
     top: 0;
     height: 100%;
+    max-height: 100%;
     left: 0;
     right: 0;
     bottom: 0;

config/locales/server.en.yml

@@ -17,6 +17,7 @@ en:
         description: "Periodic report based on a large language model"
   site_settings:
     discourse_ai_enabled: "Enable the discourse AI plugin."
+    ai_artifact_security: "The AI artifact system generates IFRAMEs with runnable code. Strict mode disables sharing and forces an extra click to run code. Lax mode allows sharing of artifacts and runs code directly. Disabled mode disables the artifact system."
     ai_toxicity_enabled: "Enable the toxicity module."
     ai_toxicity_inference_service_api_endpoint: "URL where the API is running for the toxicity module"
     ai_toxicity_inference_service_api_key: "API key for the toxicity API"
@@ -79,7 +80,7 @@ en:
     ai_embeddings_semantic_related_include_closed_topics: "Include closed topics in semantic search results"
     ai_embeddings_semantic_search_hyde_model: "Model used to expand keywords to get better results during a semantic search"
     ai_embeddings_per_post_enabled: Generate embeddings for each post
-    
+
     ai_summarization_enabled: "Enable the topic summarization module."
     ai_summarization_model: "Model to use for summarization."
     ai_custom_summarization_allowed_groups: "Groups allowed to use create new summaries."

config/settings.yml

@@ -2,7 +2,13 @@ discourse_ai:
   discourse_ai_enabled:
     default: true
     client: true
-
+  ai_artifact_security:
+    type: enum
+    default: "strict"
+    choices:
+      - "disabled"
+      - "lax"
+      - "strict"
   ai_toxicity_enabled:
     default: false
     client: true

lib/completions/dialects/dialect.rb

@@ -168,7 +168,7 @@ def system_msg(msg)
           raise NotImplemented
         end
 
-        def assistant_msg(msg)
+        def model_msg(msg)
           raise NotImplemented
         end
 
@@ -177,11 +177,15 @@ def user_msg(msg)
         end
 
         def tool_call_msg(msg)
-          { role: "assistant", content: tools_dialect.from_raw_tool_call(msg) }
+          new_content = tools_dialect.from_raw_tool_call(msg)
+          msg = msg.merge(content: new_content)
+          model_msg(msg)
         end
 
         def tool_msg(msg)
-          { role: "user", content: tools_dialect.from_raw_tool(msg) }
+          new_content = tools_dialect.from_raw_tool(msg)
+          msg = msg.merge(content: new_content)
+          user_msg(msg)
         end
       end
     end

spec/requests/ai_bot/shared_ai_conversations_spec.rb

@@ -86,6 +86,78 @@ def share_error(key)
         expect(response).to have_http_status(:success)
       end
 
+      context "when ai artifacts are in lax mode" do
+        before do
+          SiteSetting.ai_artifact_security = "lax"
+        end
+
+        it "properly shares artifacts" do
+          first_post = user_pm_share.posts.first
+
+          artifact_not_allowed = AiArtifact.create!(
+            user: bot_user,
+            post: Fabricate(:private_message_post),
+            name: "test",
+            html: "<div>test</div>",
+          )
+
+          artifact = AiArtifact.create!(
+            user: bot_user,
+            post: first_post,
+            name: "test",
+            html: "<div>test</div>",
+          )
+
+          # lets log out and see we can not access the artifacts
+          delete "/session/#{user.id}"
+
+          get artifact.url
+          expect(response).to have_http_status(:not_found)
+
+          get artifact_not_allowed.url
+          expect(response).to have_http_status(:not_found)
+
+          sign_in(user)
+
+          first_post.update!(raw: <<~RAW)
+            This is a post with an artifact
+
+            <div class="ai-artifact" data-ai-artifact-id="#{artifact.id}"></div>
+            <div class="ai-artifact" data-ai-artifact-id="#{artifact_not_allowed.id}"></div>
+          RAW
+
+          post "#{path}.json", params: { topic_id: user_pm_share.id }
+          expect(response).to have_http_status(:success)
+
+          key = response.parsed_body["share_key"]
+
+          get "#{path}/#{key}"
+          expect(response).to have_http_status(:success)
+
+
+          expect(response.body).to include(artifact.url)
+          expect(response.body).to include(artifact_not_allowed.url)
+
+          # lets log out and see we can not access the artifacts
+          delete "/session/#{user.id}"
+
+          get artifact.url
+          expect(response).to have_http_status(:success)
+
+          get artifact_not_allowed.url
+          expect(response).to have_http_status(:not_found)
+
+          sign_in(user)
+          delete "#{path}/#{key}.json"
+          expect(response).to have_http_status(:success)
+
+          # we can not longer see it...
+          delete "/session/#{user.id}"
+          get artifact.url
+          expect(response).to have_http_status(:not_found)
+        end
+      end
+
       context "when secure uploads are enabled" do
         let(:upload_1) { Fabricate(:s3_image_upload, user: bot_user, secure: true) }
         let(:upload_2) { Fabricate(:s3_image_upload, user: bot_user, secure: true) }