Coverage for src/kwai/frontend/etag_file_response.py: 96%

1"""Module that implements an etag response for a file.""" 

2 

3from hashlib import md5 

4from pathlib import Path 

5 

6from fastapi import Request 

7from fastapi.responses import FileResponse, Response 

8 

9 

10class EtagFileResponse(FileResponse): 

11 """A FileResponse that will check the etag when if-none-match header is passed. 

12 

13 _generate_etag uses the same implementation as FileResponse. 

14 FileResponse automatically sets the etag header with this etag value. 

15 """ 

16 

17 def __init__(self, path: Path, **kwargs) -> None: 

18 super().__init__(path, **kwargs) 

19 self._path = path 

20 

21 def _generate_etag(self) -> str: 

22 file_stat = self._path.stat() 

23 etag_base = str(file_stat.st_mtime) + "-" + str(file_stat.st_size) 

24 return f'"{md5(etag_base.encode(), usedforsecurity=False).hexdigest()}"' 

25 

26 async def __call__(self, scope, receive, send) -> None: 

27 """Check the etag, and return 304 when the file is not modified.""" 

28 request = Request(scope, receive) 

29 if_none_match = request.headers.get("if-none-match") 

30 

31 if not if_none_match: 

32 await super().__call__(scope, receive, send) 

33 return 

34 

35 etag = self._generate_etag() 

36 if if_none_match == etag: 

37 response = Response(status_code=304) 

38 await response(scope, receive, send) 

39 return 

40 

41 await super().__call__(scope, receive, send)