diff --git a/app/controllers/discourse_ai/ai_helper/assistant_controller.rb b/app/controllers/discourse_ai/ai_helper/assistant_controller.rb index 684f6ee0..2df162df 100644 --- a/app/controllers/discourse_ai/ai_helper/assistant_controller.rb +++ b/app/controllers/discourse_ai/ai_helper/assistant_controller.rb @@ -8,6 +8,8 @@ module DiscourseAi before_action :ensure_can_request_suggestions before_action :rate_limiter_performed!, except: %i[prompts] + include SecureUploadEndpointHelpers + def suggest input = get_text_param! @@ -109,16 +111,16 @@ module DiscourseAi image_url = params[:image_url] raise Discourse::InvalidParameters.new(:image_url) if !image_url - image = Upload.find_by(sha1: Upload.sha1_from_long_url(image_url)) - - if image&.secure? - url = Upload.signed_url_from_secure_uploads_url(image_url) - else - url = UrlHelper.absolute(image_url) - end + image = upload_from_full_url(image_url) + raise Discourse::NotFound if image.blank? + final_image_url = get_caption_url(image, image_url) hijack do - caption = DiscourseAi::AiHelper::Assistant.new.generate_image_caption(url, current_user) + caption = + DiscourseAi::AiHelper::Assistant.new.generate_image_caption( + final_image_url, + current_user, + ) render json: { caption: caption }, status: 200 end rescue DiscourseAi::Completions::Endpoints::Base::CompletionFailed, Net::HTTPBadResponse @@ -141,14 +143,18 @@ module DiscourseAi end def ensure_can_request_suggestions - user_group_ids = current_user.group_ids + if !current_user.in_any_groups?(SiteSetting.ai_helper_allowed_groups_map) + raise Discourse::InvalidAccess + end + end - allowed = - SiteSetting.ai_helper_allowed_groups_map.any? do |group_id| - user_group_ids.include?(group_id) - end + def get_caption_url(image_upload, image_url) + if image_upload.secure? + check_secure_upload_permission(image_upload) + return Discourse.store.url_for(image_upload) + end - raise Discourse::InvalidAccess if !allowed + UrlHelper.absolute(image_url) end end end diff --git a/spec/requests/ai_helper/assistant_controller_spec.rb b/spec/requests/ai_helper/assistant_controller_spec.rb index 0e28715c..752b1144 100644 --- a/spec/requests/ai_helper/assistant_controller_spec.rb +++ b/spec/requests/ai_helper/assistant_controller_spec.rb @@ -109,15 +109,15 @@ RSpec.describe DiscourseAi::AiHelper::AssistantController do end describe "#caption_image" do - let(:image_url) { "https://example.com/image.jpg" } + fab!(:upload) { Fabricate(:upload) } + let(:image_url) { "#{Discourse.base_url}#{upload.url}" } let(:caption) { "A picture of a cat sitting on a table" } context "when logged in as an allowed user" do - fab!(:user) + fab!(:user) { Fabricate(:user, refresh_auto_groups: true) } before do sign_in(user) - user.group_ids = [Group::AUTO_GROUPS[:trust_level_1]] SiteSetting.ai_helper_allowed_groups = Group::AUTO_GROUPS[:trust_level_1] SiteSetting.ai_llava_endpoint = "https://example.com" @@ -147,6 +147,50 @@ RSpec.describe DiscourseAi::AiHelper::AssistantController do expect(response.status).to eq(400) end + + it "returns a 404 error if no upload is found" do + post "/discourse-ai/ai-helper/caption_image", + params: { + image_url: "http://blah.com/img.jpeg", + } + + expect(response.status).to eq(404) + end + + context "for secure uploads" do + fab!(:group) { Fabricate(:group) } + fab!(:private_category) { Fabricate(:private_category, group: group) } + fab!(:post_in_secure_context) do + Fabricate(:post, topic: Fabricate(:topic, category: private_category)) + end + fab!(:upload) { Fabricate(:secure_upload, access_control_post: post_in_secure_context) } + let(:image_url) { "#{Discourse.base_url}/secure-uploads/#{upload.url}" } + + before { enable_secure_uploads } + + it "returns a 403 error if the user cannot access the secure upload" do + post "/discourse-ai/ai-helper/caption_image", params: { image_url: image_url } + expect(response.status).to eq(403) + end + + it "returns a 200 message and caption if user can access the secure upload" do + group.add(user) + post "/discourse-ai/ai-helper/caption_image", params: { image_url: image_url } + expect(response.status).to eq(200) + expect(response.parsed_body["caption"]).to eq(caption) + end + + context "if the input URL is for a secure upload but not on the secure-uploads path" do + let(:image_url) { "#{Discourse.base_url}#{upload.url}" } + + it "creates a signed URL properly and makes the caption" do + group.add(user) + post "/discourse-ai/ai-helper/caption_image", params: { image_url: image_url } + expect(response.status).to eq(200) + expect(response.parsed_body["caption"]).to eq(caption) + end + end + end end end end