fix: dynamically update image URLs in agent_thoughts with valid signatures#31955
fix: dynamically update image URLs in agent_thoughts with valid signatures#31955jackjin1997 wants to merge 2 commits intolanggenius:mainfrom
Conversation
dosubot
bot
added
the
size:M
Summary of ChangesHello @jackjin1997, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves an issue where image URLs embedded in agent thoughts and observations would expire due to static signatures. It introduces a robust solution to dynamically re-sign these URLs upon access, ensuring their continued validity. This change aligns the handling of these fields with existing practices for other signed URLs within the system, improving the reliability of displayed content. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request aims to resolve expiring image URLs in agent thoughts by implementing dynamic re-signing logic. However, the _re_sign_file_url method in api/models/model.py introduces a critical security vulnerability: an 'Authorization Oracle'. This flaw allows re-signing of file URLs without verifying their original signature, enabling attackers to bypass access controls and obtain validly signed URLs for arbitrary files through prompt injection. It is crucial to verify the original signature before generating a new one. Furthermore, the _re_sign_file_url method also requires refactoring to correct filename parsing, reduce code duplication, and improve overall maintainability by extracting shared helper functions.
| def _re_sign_file_url(self, text: str | None) -> str | None: | ||
| if not text: | ||
| return text | ||
|
|
||
| pattern = r"\[!?.*?\]\((((http|https):\/\/.+)?\/files\/(tools\/)?[\w-]+.*?timestamp=.*&nonce=.*&sign=.*)\)" | ||
| matches = re.findall(pattern, text) | ||
|
|
||
| if not matches: | ||
| return text | ||
|
|
||
| urls = [match[0] for match in matches] | ||
|
|
||
| # remove duplicate urls | ||
| urls = list(set(urls)) | ||
|
|
||
| if not urls: | ||
| return text | ||
|
|
||
| re_signed_text = text | ||
| for url in urls: | ||
| if "files/tools" in url: | ||
| # get tool file id | ||
| tool_file_id_pattern = r"\/files\/tools\/([\.\w-]+)?\?timestamp=" | ||
| result = re.search(tool_file_id_pattern, url) | ||
| if not result: | ||
| continue | ||
|
|
||
| tool_file_id = result.group(1) | ||
|
|
||
| # get extension | ||
| if "." in tool_file_id: | ||
| split_result = tool_file_id.split(".") | ||
| extension = f".{split_result[-1]}" | ||
| if len(extension) > 10: | ||
| extension = ".bin" | ||
| tool_file_id = split_result[0] | ||
| else: | ||
| extension = ".bin" | ||
|
|
||
| if not tool_file_id: | ||
| continue | ||
|
|
||
| sign_url = sign_tool_file(tool_file_id=tool_file_id, extension=extension) | ||
| elif "file-preview" in url: | ||
| # get upload file id | ||
| upload_file_id_pattern = r"\/files\/([\w-]+)\/file-preview\?timestamp=" | ||
| result = re.search(upload_file_id_pattern, url) | ||
| if not result: | ||
| continue | ||
|
|
||
| upload_file_id = result.group(1) | ||
| if not upload_file_id: | ||
| continue | ||
| sign_url = file_helpers.get_signed_file_url(upload_file_id) | ||
| elif "image-preview" in url: | ||
| # image-preview is deprecated, use file-preview instead | ||
| upload_file_id_pattern = r"\/files\/([\w-]+)\/image-preview\?timestamp=" | ||
| result = re.search(upload_file_id_pattern, url) | ||
| if not result: | ||
| continue | ||
| upload_file_id = result.group(1) | ||
| if not upload_file_id: | ||
| continue | ||
| sign_url = file_helpers.get_signed_file_url(upload_file_id) | ||
| else: | ||
| continue | ||
| # if as_attachment is in the url, add it to the sign_url. | ||
| if "as_attachment" in url: | ||
| sign_url += "&as_attachment=true" | ||
| re_signed_text = re_signed_text.replace(url, sign_url) | ||
|
|
||
| return re_signed_text |
There was a problem hiding this comment.
The _re_sign_file_url method in api/models/model.py presents a critical security vulnerability, functioning as an 'Authorization Oracle'. It re-signs file URLs without verifying the original signature, which means an attacker can inject malicious markdown links via prompt injection to obtain validly signed URLs for arbitrary files, thereby bypassing access controls. It is imperative to implement robust signature verification before re-signing any URL.
Beyond the security concern, this method also has several code quality issues:
- Code Duplication: The logic is very similar to
re_sign_file_url_answerin theMessagemodel and should be extracted into a shared helper function. - Filename Parsing Bug: Using
split('.')is incorrect for filenames with multiple dots;rsplit('.', 1)is the correct approach. - Redundant Logic: The handling for
file-previewandimage-previewis duplicated and can be consolidated.
The provided code suggestion addresses the code quality issues, but the fundamental security flaw of missing signature verification must be prioritized and resolved.
def _re_sign_file_url(self, text: str | None) -> str | None:
if not text:
return text
pattern = r"[!?.*]((((http|https):\/\/.+)?\/files\/(tools\/)?[:w-]+.*?timestamp=.*&nonce=.*&sign=.*))"
matches = re.findall(pattern, text)
if not matches:
return text
urls = list(set(match[0] for match in matches))
if not urls:
return text
re_signed_text = text
for url in urls:
if "files/tools" in url:
# get tool file id
tool_file_id_pattern = r"\/files\/tools\/([\.:w-]+)?\?timestamp="
result = re.search(tool_file_id_pattern, url)
if not result:
continue
tool_file_id = result.group(1)
# get extension
if "." in tool_file_id:
split_result = tool_file_id.rsplit(".", 1)
tool_file_id = split_result[0]
extension = f".{split_result[-1]}"
if len(extension) > 10:
extension = ".bin"
else:
extension = ".bin"
if not tool_file_id:
continue
sign_url = sign_tool_file(tool_file_id=tool_file_id, extension=extension)
elif "file-preview" in url or "image-preview" in url:
# get upload file id
upload_file_id_pattern = r"\/files\/([:w-]+)\/(?:file-preview|image-preview)\?timestamp="
result = re.search(upload_file_id_pattern, url)
if not result:
continue
upload_file_id = result.group(1)
if not upload_file_id:
continue
sign_url = file_helpers.get_signed_file_url(upload_file_id)
else:
continue
# if as_attachment is in the url, add it to the sign_url.
if "as_attachment" in url:
sign_url += "&as_attachment=true"
re_signed_text = re_signed_text.replace(url, sign_url)
return re_signed_text1 participant
Fixes #31907.
Description
Currently, image URLs in
agent_thoughts.thoughtandagent_thoughts.observationhave static signatures that expire after ~10 minutes. This PR adds dynamic re-signing logic to these fields, consistent with how theanswerfield is handled in theMessagemodel.Changes
_re_sign_file_urlhelper and related properties (re_sign_file_url_thought,re_sign_file_url_observation) toMessageAgentThoughtmodel.AgentThoughtPydantic schema to use these properties viavalidation_alias.This is a more lightweight approach compared to other proposed fixes, focusing on the core issue of signature expiration in API responses.
Fixes #31907