Gogs 3 dni temu
rodzic
commit
74ded560a9
2 zmienionych plików z 30 dodań i 18 usunięć
  1. 4 4
      relay_client.py
  2. 26 14
      relay_server.py

+ 4 - 4
relay_client.py

@@ -83,10 +83,10 @@ class RelayConnection:
                     self._dispatch_frame(frame, handler)
                 else:
                     self._record_dropped_frame(frame.kind)
-        except asyncio.IncompleteReadError:
-            print(f"[edge] relay disconnected name={self.node.name} eof=true")
-        except Exception as exc:
-            print(f"[edge] relay pump error name={self.node.name} error={exc!r}")
+        except (asyncio.IncompleteReadError, ConnectionResetError, BrokenPipeError, OSError):
+            pass
+        except Exception:
+            pass
         finally:
             await self.close()
 

+ 26 - 14
relay_server.py

@@ -106,10 +106,12 @@ class RelayChannel:
         if frame.kind == PING:
             await self.safe_send(Frame(PONG, 0, 0, frame.seq, 0, b"pong"))
             return
+        if frame.kind == AUTH:
+            return
         if frame.kind == TCP_OPEN:
-            meta = decode_json(frame.payload)
-            family = int(meta.get("family", 0)) or 0
             try:
+                meta = decode_json(frame.payload) if frame.payload else {}
+                family = int(meta.get("family", 0)) or 0
                 reader, writer = await asyncio.open_connection(meta["host"], int(meta["port"]), family=family or 0)
                 task = asyncio.create_task(self._tcp_pump(frame.session_id, frame.stream_id, reader))
                 self.tcp_sessions[key] = TcpSession(frame.session_id, frame.stream_id, writer, task)
@@ -134,21 +136,31 @@ class RelayChannel:
             meta = None
             payload = frame.payload
             if frame.packet_id > 0:
-                meta = decode_json(frame.payload[: frame.packet_id])
-                payload = frame.payload[frame.packet_id :]
+                try:
+                    meta = decode_json(frame.payload[: frame.packet_id])
+                    payload = frame.payload[frame.packet_id :]
+                except Exception:
+                    if session is None:
+                        return
+                    payload = frame.payload
             if session is None:
                 if meta is None:
                     return
-                family = int(meta.get("family", 0)) or 0
-                transport, protocol = await asyncio.get_running_loop().create_datagram_endpoint(
-                    lambda: RelayUdpProtocol(self, frame.session_id, frame.stream_id),
-                    remote_addr=(meta["host"], int(meta["port"])),
-                    family=family or 0,
-                )
-                session = UdpSession(frame.session_id, frame.stream_id, transport, protocol, meta["host"], int(meta["port"]), family)
-                self.udp_sessions[key] = session
-            with contextlib.suppress(Exception):
-                session.transport.sendto(payload)
+                try:
+                    family = int(meta.get("family", 0)) or 0
+                    transport, protocol = await asyncio.get_running_loop().create_datagram_endpoint(
+                        lambda: RelayUdpProtocol(self, frame.session_id, frame.stream_id),
+                        remote_addr=(meta["host"], int(meta["port"])),
+                        family=family or 0,
+                    )
+                    session = UdpSession(frame.session_id, frame.stream_id, transport, protocol, meta["host"], int(meta["port"]), family)
+                    self.udp_sessions[key] = session
+                except Exception as exc:
+                    await self.safe_send(Frame(TCP_STATUS, frame.session_id, frame.stream_id, 0, STATUS_ERR, str(exc).encode()))
+                    return
+            if session.transport is not None:
+                with contextlib.suppress(Exception):
+                    session.transport.sendto(payload)
             return
 
     async def _tcp_pump(self, session_id: int, stream_id: int, reader: asyncio.StreamReader) -> None: