DEV: Cleanup caption endpoint and account for secure uploads (#478)

Utilizes the check for secure upload permissions from core PR
https://github.com/discourse/discourse/pull/25758 and cleans up
controller codes and spec code to reuse existing code and better
reflect reality.

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
-
-        allowed =
-          SiteSetting.ai_helper_allowed_groups_map.any? do |group_id|
-            user_group_ids.include?(group_id)
+        if !current_user.in_any_groups?(SiteSetting.ai_helper_allowed_groups_map)
+          raise Discourse::InvalidAccess
+        end
       end
 
-        raise Discourse::InvalidAccess if !allowed
+      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
+
+        UrlHelper.absolute(image_url)
       end
     end
   end

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