airsltd commited on
Commit
542be03
·
verified ·
1 Parent(s): 33b16a5

Upload folder using huggingface_hub

Browse files
Files changed (2) hide show
  1. Dockerfile +6 -7
  2. server.py +117 -0
Dockerfile CHANGED
@@ -1,8 +1,6 @@
1
- FROM nginxinc/nginx-unprivileged:1.25-alpine
2
 
3
- USER root
4
-
5
- RUN apk add --no-cache curl
6
 
7
  ARG WSTUNNEL_VERSION=v10.5.5
8
  RUN curl -fsSL \
@@ -12,10 +10,11 @@ RUN curl -fsSL \
12
  && rm /tmp/wstunnel.tar.gz \
13
  && chmod +x /usr/local/bin/wstunnel
14
 
15
- COPY nginx.conf /etc/nginx/nginx.conf
 
16
 
17
- COPY --chmod=755 entrypoint.sh /entrypoint.sh
18
 
19
  EXPOSE 7860
20
 
21
- CMD ["/entrypoint.sh"]
 
1
+ FROM alpine:3.19
2
 
3
+ RUN apk add --no-cache curl python3
 
 
4
 
5
  ARG WSTUNNEL_VERSION=v10.5.5
6
  RUN curl -fsSL \
 
10
  && rm /tmp/wstunnel.tar.gz \
11
  && chmod +x /usr/local/bin/wstunnel
12
 
13
+ RUN adduser -D -u 1000 user
14
+ USER user
15
 
16
+ COPY --chown=user server.py /app/server.py
17
 
18
  EXPOSE 7860
19
 
20
+ CMD ["python3", "/app/server.py"]
server.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import asyncio
2
+ import os
3
+ import subprocess
4
+ import signal
5
+ import sys
6
+
7
+ WSTUNNEL_PORT = 7861
8
+ PROXY_PORT = 7860
9
+ WSTUNNEL_BIN = "/usr/local/bin/wstunnel"
10
+
11
+
12
+ async def health_check(reader, writer):
13
+ writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 3\r\n\r\nOK\n")
14
+ await writer.drain()
15
+ writer.close()
16
+
17
+
18
+ async def tcp_proxy(reader, writer):
19
+ try:
20
+ backend_reader, backend_writer = await asyncio.wait_for(
21
+ asyncio.open_connection("127.0.0.1", WSTUNNEL_PORT), timeout=5
22
+ )
23
+ except (OSError, asyncio.TimeoutError):
24
+ writer.write(b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 15\r\n\r\nBackend unavailable\n")
25
+ await writer.drain()
26
+ writer.close()
27
+ return
28
+
29
+ async def forward(src, dst):
30
+ try:
31
+ while True:
32
+ data = await src.read(65536)
33
+ if not data:
34
+ break
35
+ dst.write(data)
36
+ await dst.drain()
37
+ except (ConnectionResetError, BrokenPipeError, OSError):
38
+ pass
39
+ finally:
40
+ try:
41
+ dst.close()
42
+ except OSError:
43
+ pass
44
+
45
+ await asyncio.gather(
46
+ forward(reader, backend_writer),
47
+ forward(backend_reader, writer),
48
+ )
49
+
50
+
51
+ async def handle_client(reader, writer):
52
+ try:
53
+ data = await asyncio.wait_for(reader.read(4096), timeout=10)
54
+ if not data:
55
+ writer.close()
56
+ return
57
+ is_health_check = (
58
+ data.startswith(b"GET / ")
59
+ and b"Upgrade: websocket" not in data
60
+ )
61
+ if is_health_check:
62
+ await health_check(reader, writer)
63
+ else:
64
+ writer.write(data)
65
+ await writer.drain()
66
+ await tcp_proxy(reader, writer)
67
+ except (ConnectionResetError, BrokenPipeError, OSError, asyncio.TimeoutError):
68
+ pass
69
+ finally:
70
+ try:
71
+ writer.close()
72
+ except OSError:
73
+ pass
74
+
75
+
76
+ async def run_wstunnel():
77
+ token = os.environ.get("PROXY_TOKEN", "")
78
+ args = [WSTUNNEL_BIN, "server"]
79
+ if token:
80
+ args += ["--restrict-http-upgrade-path-prefix", token]
81
+ args += [f"ws://0.0.0.0:{WSTUNNEL_PORT}"]
82
+
83
+ proc = await asyncio.create_subprocess_exec(*args, stdout=sys.stderr, stderr=sys.stderr)
84
+ return proc
85
+
86
+
87
+ async def main():
88
+ print(f"starting wstunnel on port {WSTUNNEL_PORT}...", flush=True)
89
+ wstunnel_proc = await run_wstunnel()
90
+ await asyncio.sleep(0.5)
91
+
92
+ if wstunnel_proc.returncode is not None:
93
+ print(f"ERROR: wstunnel exited with code {wstunnel_proc.returncode}", flush=True)
94
+ sys.exit(1)
95
+
96
+ print(f"starting TCP proxy on port {PROXY_PORT}...", flush=True)
97
+ server = await asyncio.start_server(handle_client, "0.0.0.0", PROXY_PORT)
98
+
99
+ async def shutdown():
100
+ server.close()
101
+ wstunnel_proc.terminate()
102
+ await wstunnel_proc.wait()
103
+
104
+ loop = asyncio.get_event_loop()
105
+ for sig in (signal.SIGTERM, signal.SIGINT):
106
+ try:
107
+ loop.add_signal_handler(sig, lambda: asyncio.create_task(shutdown()))
108
+ except NotImplementedError:
109
+ pass
110
+
111
+ async with server:
112
+ print(f"proxy server ready on port {PROXY_PORT}", flush=True)
113
+ await server.serve_forever()
114
+
115
+
116
+ if __name__ == "__main__":
117
+ asyncio.run(main())