/
/
/
1"""Helpers for setting up a aiohttp session (and related)."""
2
3from __future__ import annotations
4
5import asyncio
6import socket
7import sys
8from contextlib import suppress
9from functools import cache
10from ssl import SSLContext
11from types import MappingProxyType
12from typing import TYPE_CHECKING, Any, Self
13
14import aiohttp
15from aiohttp import web
16from aiohttp.hdrs import USER_AGENT
17from aiohttp_asyncmdnsresolver.api import AsyncDualMDNSResolver
18from aiohttp_socks import ProxyConnector
19from music_assistant_models.enums import EventType
20from yarl import URL
21
22from music_assistant.constants import APPLICATION_NAME
23
24from . import ssl as ssl_util
25from .json import json_dumps, json_loads
26
27if TYPE_CHECKING:
28 from aiohttp.typedefs import JSONDecoder
29 from music_assistant_models.event import MassEvent
30
31 from music_assistant.mass import MusicAssistant
32
33
34MAXIMUM_CONNECTIONS = 4096
35MAXIMUM_CONNECTIONS_PER_HOST = 100
36
37
38def encoded_request_url(url: str) -> str | URL:
39 """
40 Return a URL safe to pass to aiohttp without altering existing percent-encoding.
41
42 :param url: The raw request URL.
43 """
44 # already-encoded URLs (e.g. %3F in a query value) must reach the server unchanged or
45 # auth-bearing stream URLs fail with a 401; leave plain URLs for yarl to normalise
46 return URL(url, encoded=True) if "%" in url else url
47
48
49def create_clientsession(
50 mass: MusicAssistant,
51 verify_ssl: bool = True,
52 socks_url: str = "",
53 **kwargs: Any,
54) -> aiohttp.ClientSession:
55 """Create a new ClientSession with kwargs, i.e. for cookies."""
56 clientsession = aiohttp.ClientSession(
57 connector=_get_connector(mass, verify_ssl, socks_url),
58 json_serialize=json_dumps,
59 response_class=MassClientResponse,
60 **kwargs,
61 )
62 # Prevent packages accidentally overriding our default headers
63 # It's important that we identify as Music Assistant
64 # If a package requires a different user agent, override it by passing a headers
65 # dictionary to the request method.
66 user_agent = (
67 f"{APPLICATION_NAME}/{mass.version} "
68 f"aiohttp/{aiohttp.__version__} Python/{sys.version_info[0]}.{sys.version_info[1]}"
69 )
70 clientsession._default_headers = MappingProxyType( # type: ignore[assignment]
71 {USER_AGENT: user_agent},
72 )
73 return clientsession
74
75
76async def async_aiohttp_proxy_stream(
77 mass: MusicAssistant,
78 request: web.BaseRequest,
79 stream: aiohttp.StreamReader,
80 content_type: str | None,
81 buffer_size: int = 102400,
82 timeout: int = 10,
83) -> web.StreamResponse:
84 """Stream a stream to aiohttp web response."""
85 response = web.StreamResponse()
86 if content_type is not None:
87 response.content_type = content_type
88 await response.prepare(request)
89
90 # Suppressing something went wrong fetching data, closed connection
91 with suppress(TimeoutError, aiohttp.ClientError):
92 while not mass.closing:
93 async with asyncio.timeout(timeout):
94 data = await stream.read(buffer_size)
95
96 if not data:
97 break
98 await response.write(data)
99
100 return response
101
102
103class MassAsyncDNSResolver(AsyncDualMDNSResolver):
104 """
105 Music Assistant AsyncDNSResolver.
106
107 This is a wrapper around the AsyncDualMDNSResolver to only
108 close the resolver when the Music Assistant instance is closed.
109 """
110
111 async def real_close(self) -> None:
112 """Close the resolver."""
113 await super().close()
114
115 async def close(self) -> None:
116 """Close the resolver."""
117
118
119class MassClientResponse(aiohttp.ClientResponse):
120 """aiohttp.ClientResponse with a json method that uses json_loads by default."""
121
122 async def json(
123 self,
124 *args: Any,
125 loads: JSONDecoder = json_loads,
126 **kwargs: Any,
127 ) -> Any:
128 """Send a json request and parse the json response."""
129 return await super().json(*args, loads=loads, **kwargs)
130
131
132class ChunkAsyncStreamIterator:
133 """
134 Async iterator for chunked streams.
135
136 Based on aiohttp.streams.ChunkTupleAsyncStreamIterator, but yields
137 bytes instead of tuple[bytes, bool].
138 """
139
140 __slots__ = ("_stream",)
141
142 def __init__(self, stream: aiohttp.StreamReader) -> None:
143 """Initialize."""
144 self._stream = stream
145
146 def __aiter__(self) -> Self:
147 """Iterate."""
148 return self
149
150 async def __anext__(self) -> bytes:
151 """Yield next chunk."""
152 rv = await self._stream.readchunk()
153 if rv == (b"", False):
154 raise StopAsyncIteration
155 return rv[0]
156
157
158class MusicAssistantTCPConnector(aiohttp.TCPConnector):
159 """
160 Music Assistant TCP Connector.
161
162 Same as aiohttp.TCPConnector but with a longer cleanup_closed timeout.
163
164 By default the cleanup_closed timeout is 2 seconds. This is too short
165 for Music Assistant since we churn through a lot of connections. We set
166 it to 60 seconds to reduce the overhead of aborting TLS connections
167 that are likely already closed.
168 """
169
170 # abort transport after 60 seconds (cleanup broken connections)
171 _cleanup_closed_period = 60.0
172
173
174def _get_connector(
175 mass: MusicAssistant,
176 verify_ssl: bool = True,
177 socks_url: str | None = None,
178 family: socket.AddressFamily = socket.AF_UNSPEC,
179 ssl_cipher: ssl_util.SSLCipherList = ssl_util.SSLCipherList.PYTHON_DEFAULT,
180) -> aiohttp.BaseConnector | ProxyConnector:
181 """
182 Return the connector pool for aiohttp.
183
184 This method must be run in the event loop.
185 """
186 if verify_ssl:
187 ssl_context: SSLContext = ssl_util.client_context(ssl_cipher)
188 else:
189 ssl_context = ssl_util.client_context_no_verify(ssl_cipher)
190
191 if socks_url:
192 return ProxyConnector.from_url(
193 socks_url,
194 ssl=ssl_context,
195 limit=MAXIMUM_CONNECTIONS,
196 limit_per_host=MAXIMUM_CONNECTIONS_PER_HOST,
197 resolver=_get_resolver(mass),
198 )
199
200 return MusicAssistantTCPConnector(
201 family=family,
202 # Cleanup closed is no longer needed after https://github.com/python/cpython/pull/118960
203 # which first appeared in Python 3.12.7 and 3.13.1
204 enable_cleanup_closed=False,
205 ssl=ssl_context,
206 limit=MAXIMUM_CONNECTIONS,
207 limit_per_host=MAXIMUM_CONNECTIONS_PER_HOST,
208 resolver=_get_resolver(mass),
209 )
210
211
212@cache
213def _get_resolver(mass: MusicAssistant) -> MassAsyncDNSResolver:
214 """Return the MassAsyncDNSResolver."""
215 resolver = MassAsyncDNSResolver(async_zeroconf=mass.discovery.aiozc)
216
217 async def _close_resolver(event: MassEvent) -> None: # noqa: ARG001
218 await resolver.real_close()
219
220 mass.subscribe(_close_resolver, EventType.SHUTDOWN)
221 return resolver
222
223
224def get_socks5_url(url_string: str) -> str:
225 """Return full socks5 url string from config."""
226 if url_string:
227 if "://" not in url_string:
228 url_string = f"socks5://{url_string}"
229 return str(url_string)
230