File size: 5,324 Bytes
e12a676
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
"""媒体缓存文件的有界读取、Range交付与真实线程关闭。"""
import asyncio
import os
from secrets import token_hex
import stat

from fastapi.responses import FileResponse
from starlette.datastructures import MutableHeaders
from owned_thread import run_owned_thread
from workflow_stream_response import OwnedDelivery


class OwnedMediaFileResponse(OwnedDelivery, FileResponse):
    total_timeout = 300
    operation_timeout = 15
    eof_timeout = 5

    def __init__(self, path, **kwargs):
        snapshot = os.stat(path)
        if not stat.S_ISREG(snapshot.st_mode):
            raise ValueError('invalid media cache file')
        super().__init__(path, stat_result=snapshot, **kwargs)
        self._file = None
        self._init_delivery(self._close_file)

    async def send_owned(self, scope, receive, send):
        async def bounded_send(message):
            final = message['type'] == 'http.response.body' and not message.get('more_body', False)
            async with asyncio.timeout(self.eof_timeout if final else self.operation_timeout):
                await send(message)
        async with asyncio.timeout(self.total_timeout):
            return await super().send_owned(scope, receive, bounded_send)

    async def _file_work(self, function, *args):
        async with asyncio.timeout(self.operation_timeout):
            return await run_owned_thread(function, *args)

    async def _open_file(self):
        def open_file():
            self._file = open(self.path, 'rb')
            self._check_snapshot(os.fstat(self._file.fileno()))
        await self._file_work(open_file)

    def _check_snapshot(self, current):
        fields = ('st_dev', 'st_ino', 'st_size', 'st_mtime_ns')
        if os.name != 'nt':
            fields += ('st_ctime_ns',)
        if (not stat.S_ISREG(current.st_mode) or not current.st_ino or not self.stat_result.st_ino
                or any(getattr(current, key) != getattr(self.stat_result, key) for key in fields)):
            raise RuntimeError('media cache file changed')

    async def _prepare_file(self, header_only):
        if header_only:
            current = await self._file_work(os.stat, self.path)
            self._check_snapshot(current)
        else:
            await self._open_file()

    async def _close_file(self):
        if self._file is None:
            return
        def close_file():
            try:
                self._file.close()
            finally:
                self._file = None
        await run_owned_thread(close_file)

    async def _send_part(self, send, start, end, *, final):
        await self._file_work(self._file.seek, start)
        remaining = end - start
        if remaining == 0:
            await send({'type': 'http.response.body', 'body': b'', 'more_body': not final})
        while remaining:
            chunk = await self._file_work(self._file.read, min(self.chunk_size, remaining))
            if not chunk:
                raise RuntimeError('media cache file truncated')
            remaining -= len(chunk)
            await send({'type': 'http.response.body', 'body': chunk, 'more_body': remaining > 0 or not final})

    async def _handle_simple(self, send, send_header_only, send_pathsend):
        await self._prepare_file(send_header_only)
        await send({'type': 'http.response.start', 'status': self.status_code, 'headers': self.raw_headers})
        if send_header_only:
            await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
            return
        await self._send_part(send, 0, self.stat_result.st_size, final=True)

    async def _handle_single_range(self, send, start, end, file_size, send_header_only):
        await self._prepare_file(send_header_only)
        headers = MutableHeaders(raw=list(self.raw_headers))
        headers['content-range'] = f'bytes {start}-{end - 1}/{file_size}'
        headers['content-length'] = str(end - start)
        await send({'type': 'http.response.start', 'status': 206, 'headers': headers.raw})
        if send_header_only:
            await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
            return
        await self._send_part(send, start, end, final=True)

    async def _handle_multiple_ranges(self, send, ranges, file_size, send_header_only):
        await self._prepare_file(send_header_only)
        boundary = token_hex(13)
        length, header = self.generate_multipart(ranges, boundary, file_size, self.headers['content-type'])
        headers = MutableHeaders(raw=list(self.raw_headers))
        headers['content-type'] = f'multipart/byteranges; boundary={boundary}'
        headers['content-length'] = str(length)
        await send({'type': 'http.response.start', 'status': 206, 'headers': headers.raw})
        if send_header_only:
            await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
            return
        for start, end in ranges:
            await send({'type': 'http.response.body', 'body': header(start, end), 'more_body': True})
            await self._send_part(send, start, end, final=False)
            await send({'type': 'http.response.body', 'body': b'\r\n', 'more_body': True})
        await send({'type': 'http.response.body', 'body': f'--{boundary}--'.encode('latin-1'), 'more_body': False})