setup unix socket manually, and allow changing the permissions on it

1 files changed, 35 insertions, 10 deletions

diff --git a/pypaste/__main__.py b/pypaste/__main__.py
index 21ac508..342d4b0 100644
--- a/pypaste/__main__.py
+++ b/pypaste/__main__.py
@@ -14,12 +14,14 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 import sys
+import os
 import asyncio
 import secrets
 import aiohttp
 import aiosqlite
 import zstandard
 from . import s3
+from socket import socket, AF_UNIX, SOCK_STREAM
 from hashlib import sha256
 from argparse import ArgumentParser
 from aiohttp import web
@@ -295,11 +297,10 @@ class App:
             await self.database.delete(oldest)
 
 
-def main() -> int:
+async def main() -> int:
     parser = ArgumentParser()
-    parser.add_argument("--host")
-    parser.add_argument("--port", type=int)
-    parser.add_argument("--path")
+    parser.add_argument("--socket", required=True)
+    parser.add_argument("--socket-mode", default="0600")
     parser.add_argument("--site", required=True)
     parser.add_argument("--content-length-max-bytes", type=int, required=True)
     parser.add_argument("--key-length", type=int, required=True)
@@ -368,18 +369,42 @@ def main() -> int:
         secret_key,
     )
 
-    webapp = web.Application()
+    app = web.Application()
 
-    app = App(config, database, bucket)
+    pypaste = App(config, database, bucket)
 
-    webapp.add_routes(
-        [web.get("/paste/{key}", app.download), web.post("/paste", app.upload)]
+    app.add_routes(
+        [web.get("/paste/{key}", pypaste.download), web.post("/paste", pypaste.upload)]
     )
 
-    web.run_app(webapp, host=args.host, port=args.port, path=args.path)
+    runner = web.AppRunner(app)
+
+    await runner.setup()
+
+    sock = socket(AF_UNIX, SOCK_STREAM)
+    sock.bind(args.socket)
+
+    os.chmod(args.socket, int(args.socket_mode, 8))
+
+    site = web.SockSite(runner, sock)
+
+    await site.start()
+
+    try:
+        await asyncio.Event().wait()
+    except asyncio.exceptions.CancelledError:
+        pass
+    finally:
+        os.remove(args.socket)
 
     return 0
 
 
 if __name__ == "__main__":
-    sys.exit(main())
+    try:
+        sys.exit(asyncio.run(main()))
+    except KeyboardInterrupt:
+        sys.exit(0)
+    except Exception as e:
+        print(f"failure: {e}", file=sys.stderr)
+        sys.exit(1)