add support for handling avatar with SSO login

changelog.d/13917.feature

@@ -0,0 +1 @@
+Adds support for handling avatar in SSO login. Contributed by @ashfame.

synapse/handlers/oidc.py

@@ -1123,6 +1123,7 @@ class UserAttributeDict(TypedDict):
     localpart: Optional[str]
     confirm_localpart: bool
     display_name: Optional[str]
+    picture: Optional[str]
     emails: List[str]
 
 
@@ -1211,6 +1212,7 @@ def jinja_finalize(thing: Any) -> Any:
 @attr.s(slots=True, frozen=True, auto_attribs=True)
 class JinjaOidcMappingConfig:
     subject_claim: str
+    picture_claim: str
     localpart_template: Optional[Template]
     display_name_template: Optional[Template]
     email_template: Optional[Template]
@@ -1230,6 +1232,7 @@ def __init__(self, config: JinjaOidcMappingConfig):
     @staticmethod
     def parse_config(config: dict) -> JinjaOidcMappingConfig:
         subject_claim = config.get("subject_claim", "sub")
+        picture_claim = config.get("picture_claim", "picture")
 
         def parse_template_config(option_name: str) -> Optional[Template]:
             if option_name not in config:
@@ -1263,6 +1266,7 @@ def parse_template_config(option_name: str) -> Optional[Template]:
 
         return JinjaOidcMappingConfig(
             subject_claim=subject_claim,
+            picture_claim=picture_claim,
             localpart_template=localpart_template,
             display_name_template=display_name_template,
             email_template=email_template,
@@ -1302,10 +1306,16 @@ def render_template_field(template: Optional[Template]) -> Optional[str]:
         if email:
             emails.append(email)
 
+        if "picture" in userinfo:
+            picture = userinfo["picture"]
+        else:
+            picture = ""
+
         return UserAttributeDict(
             localpart=localpart,
             display_name=display_name,
             emails=emails,
+            picture=picture,
             confirm_localpart=self._config.confirm_localpart,
         )
 

synapse/handlers/sso.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import abc
+import io
 import logging
 from typing import (
     TYPE_CHECKING,
@@ -31,6 +32,7 @@
 import attr
 from typing_extensions import NoReturn, Protocol
 
+from twisted.web.client import readBody
 from twisted.web.iweb import IRequest
 from twisted.web.server import Request
 
@@ -42,6 +44,7 @@
 from synapse.http import get_request_user_agent
 from synapse.http.server import respond_with_html, respond_with_redirect
 from synapse.http.site import SynapseRequest
+from synapse.logging.context import make_deferred_yieldable
 from synapse.types import (
     JsonDict,
     UserID,
@@ -137,6 +140,7 @@ class UserAttributes:
     localpart: Optional[str]
     confirm_localpart: bool = False
     display_name: Optional[str] = None
+    picture: Optional[str] = None
     emails: Collection[str] = attr.Factory(list)
 
 
@@ -194,6 +198,8 @@ def __init__(self, hs: "HomeServer"):
         self._error_template = hs.config.sso.sso_error_template
         self._bad_user_template = hs.config.sso.sso_auth_bad_user_template
         self._profile_handler = hs.get_profile_handler()
+        self.media_repo = hs.get_media_repository()
+        self._http_client = hs.get_proxied_http_client()
 
         # The following template is shown after a successful user interactive
         # authentication session. It tells the user they can close the window.
@@ -701,8 +707,69 @@ async def _register_mapped_user(
         await self._store.record_user_external_id(
             auth_provider_id, remote_user_id, registered_user_id
         )
+
+        # Set avatar, if available
+        if attributes.picture:
+            await self.set_avatar(registered_user_id, attributes.picture)
+
         return registered_user_id
 
+    async def set_avatar(self, user_id: str, picture_https_url: str) -> None:
+        try:
+            uid = UserID.from_string(user_id)
+
+            # HEAD request to find image size & mime type before download
+            response = await self._http_client.request("HEAD", picture_https_url)
+            if response.code != 200:
+                raise Exception("error sending HEAD request to get image size")
+
+            content_length = int(
+                response.headers.getRawHeaders(b"Content-Length")[0].decode("utf-8")
+            )
+            content_type = response.headers.getRawHeaders(b"Content-Type")[0].decode(
+                "utf-8"
+            )
+
+            # ensure picture size respects max_avatar_size defined in config
+            if self._profile_handler.max_avatar_size is not None:
+                if content_length > self._profile_handler.max_avatar_size:
+                    raise Exception("sso avatar too big in size")
+
+            # ensure picture is of allowed mime type
+            if (
+                self._profile_handler.allowed_avatar_mimetypes
+                and content_type not in self._profile_handler.allowed_avatar_mimetypes
+            ):
+                raise Exception("mime type not allowed for sso avatar")
+
+            # download picture
+            response = await self._http_client.request("GET", picture_https_url)
+            if response.code != 200:
+                raise Exception(
+                    "error sending GET request to provided sso avatar image"
+                )
+            image = await make_deferred_yieldable(readBody(response))
+
+            # store it in media repository
+            avatar_mxc_url = await self.media_repo.create_content(
+                str(content_type),
+                None,
+                io.BytesIO(image),  # convert image into BytesIO
+                content_length,
+                uid,
+            )
+
+            # save it as user avatar
+            await self._profile_handler.set_avatar_url(
+                uid,
+                create_requester(uid),
+                str(avatar_mxc_url),
+            )
+
+            logger.info("successfully saved the user avatar")
+        except Exception:
+            logger.exception("failed to save the user avatar")
+
     async def complete_sso_ui_auth_request(
         self,
         auth_provider_id: str,