452 lines
14 KiB
Plaintext
452 lines
14 KiB
Plaintext
from stdint import *
|
|
import string
|
|
import platmacro
|
|
import stdio
|
|
import win32.winsock2
|
|
|
|
|
|
# ============================================================
|
|
# socket.py — Cross-platform wrapper for Socket
|
|
# Python-style OOP API + Compatibility with legacy procedural functions
|
|
#
|
|
# Example:
|
|
# sock = socket.Socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
# sock.connect("example.com", 80)
|
|
# sock.send(data, length)
|
|
# n = sock.recv(buf, bufsize)
|
|
# sock.close()
|
|
#
|
|
# Only 3 platform differences (via platmacro.IS_WINDOWS branch):
|
|
# 1. WSAStartup() / WSACleanup()
|
|
# 2. closesocket() vs close()
|
|
# 3. WSAGetLastError() vs errno
|
|
# ============================================================
|
|
|
|
|
|
# ============================================================
|
|
# Constants
|
|
# ============================================================
|
|
define AF_INET = 2
|
|
define AF_INET6 = 10
|
|
|
|
define SOCK_STREAM = 1
|
|
define SOCK_DGRAM = 2
|
|
define SOCK_RAW = 3
|
|
|
|
define IPPROTO_TCP = 6
|
|
define IPPROTO_UDP = 17
|
|
|
|
define SOL_SOCKET = 1
|
|
define SO_RCVTIMEO = 20
|
|
define SO_SNDTIMEO = 21
|
|
define SO_REUSEADDR = 2
|
|
|
|
define INADDR_ANY = 0
|
|
|
|
define SOCKET_ERROR = -1
|
|
define INVALID_SOCKET = -1
|
|
|
|
define SOCKET_DEFAULT_TIMEOUT = 0
|
|
|
|
|
|
# ============================================================
|
|
# Structures
|
|
# ============================================================
|
|
|
|
class SocketAddr:
|
|
"""sockaddr_in — IPv4 address (16 bytes layout)"""
|
|
family: u16
|
|
port: u16
|
|
addr: u32
|
|
zero: u64
|
|
|
|
|
|
class HostEnt:
|
|
"""struct hostent — DNS resolution result"""
|
|
h_name: str
|
|
h_aliases: bytes
|
|
h_addrtype: int
|
|
h_length: int
|
|
h_addr_list: bytes
|
|
|
|
|
|
class timeval:
|
|
"""struct timeval — Timeout setting"""
|
|
tv_sec: int
|
|
tv_usec: int
|
|
|
|
|
|
# ============================================================
|
|
# External C Functions — Using u64 Handle, Cross-platform Compatibility
|
|
# (Windows SOCKET=u64, POSIX fd=int → u64 Zero-Extension Compatible)
|
|
# ============================================================
|
|
|
|
extern def socket(family: INT, type: INT, protocol: INT) -> u64
|
|
extern def connect(fd: u64, addr: SocketAddr*, addrlen: INT) -> INT
|
|
extern def send(fd: u64, buf: t.CPtr, len: INT, flags: INT) -> INT
|
|
extern def recv(fd: u64, buf: t.CPtr, len: INT, flags: INT) -> INT
|
|
extern def bind(fd: u64, addr: SocketAddr*, addrlen: INT) -> INT
|
|
extern def listen(fd: u64, backlog: INT) -> INT
|
|
extern def accept(fd: u64, addr: SocketAddr*, addrlen: INT*) -> u64
|
|
extern def setsockopt(fd: u64, level: INT, optname: INT, optval: t.CPtr, optlen: INT) -> INT
|
|
extern def gethostbyname(name: const str) -> HostEnt*
|
|
extern def ntohs(netshort: u16) -> u16
|
|
extern def htons(hostshort: u16) -> u16
|
|
extern def close(fd: u64) -> INT
|
|
|
|
# ============================================================
|
|
# Internal Auxiliary Functions
|
|
# ============================================================
|
|
|
|
def _CloseFd(fd: INT) -> INT:
|
|
"""Platform differences: Windows uses closesocket, POSIX uses close。"""
|
|
if platmacro.IS_WINDOWS:
|
|
return win32.winsock2.closesocket(u64(fd))
|
|
return close(fd)
|
|
|
|
|
|
def _MakeAddr(host: str, port: INT) -> SocketAddr:
|
|
"""Build SocketAddr. First try decimal IP, then DNS resolution.
|
|
If decimal IP fails, try DNS resolution.
|
|
If DNS resolution fails, return INADDR_ANY."""
|
|
addr: SocketAddr
|
|
string.memset(&addr, 0, 16)
|
|
addr.family = u16(AF_INET)
|
|
addr.port = htons(u16(port))
|
|
|
|
ip: u32 = 0
|
|
if _ParseIPv4(host, &ip):
|
|
addr.addr = ip
|
|
return addr
|
|
|
|
# DNS resolution — gethostbyname returns struct hostent*
|
|
# h_addr_list is char** → h_addr_list[0] is char* → *(uint32_t*)h_addr_list[0] is IP
|
|
if platmacro.IS_WINDOWS:
|
|
he_raw: win32.winsock2.WSAHostEnt* = win32.winsock2.gethostbyname(host)
|
|
if he_raw is not None:
|
|
# he_raw.h_addr_list → he_raw->h_addr_list (char**)
|
|
# c.Deref(uint64(he_raw.h_addr_list, t.CPtr)) → *(uint64_t*)(he_raw->h_addr_list) → h_addr_list[0]
|
|
ip_ptr: u64 = u64(*(uint64*(he_raw.h_addr_list)))
|
|
if ip_ptr != 0:
|
|
addr.addr = u32(*(uint32*(ip_ptr)))
|
|
else:
|
|
he: HostEnt* = gethostbyname(host)
|
|
if he is not None:
|
|
ip_ptr: u64 = u64(*(uint64*(he.h_addr_list)))
|
|
if ip_ptr != 0:
|
|
addr.addr = u32(*(uint32*(ip_ptr)))
|
|
|
|
return addr
|
|
|
|
|
|
def _ParseIPv4(ip_str: str, out_addr: u32*) -> bool:
|
|
"""Parse IPv4 address."""
|
|
if ip_str is None or out_addr is None: return False
|
|
|
|
parts: t.CArray[INT, 4]
|
|
part_count: INT = 0
|
|
current: INT = 0
|
|
started: bool = False
|
|
|
|
for ch in ip_str:
|
|
if ch >= '0' and ch <= '9':
|
|
if not started: started = True
|
|
current = current * 10 + INT(ch - '0')
|
|
if current > 255: return False
|
|
elif ch == '.':
|
|
if not started: return False
|
|
if part_count >= 4: return False
|
|
parts[part_count] = current
|
|
part_count += 1
|
|
current = 0
|
|
started = False
|
|
else:
|
|
break
|
|
|
|
if started and part_count < 4:
|
|
parts[part_count] = current
|
|
part_count += 1
|
|
|
|
if part_count != 4: return False
|
|
|
|
result: u32 = 0
|
|
result = result | (u32(parts[0]) << 24)
|
|
result = result | (u32(parts[1]) << 16)
|
|
result = result | (u32(parts[2]) << 8)
|
|
result = result | u32(parts[3])
|
|
*out_addr = result
|
|
return True
|
|
|
|
|
|
# ============================================================
|
|
# Socket class — Python OOP API
|
|
# ============================================================
|
|
|
|
class Socket:
|
|
"""Cross-platform Socket class, similar to Python socket.socket.
|
|
|
|
Usage:
|
|
sock = Socket(AF_INET, SOCK_STREAM)
|
|
sock.connect("example.com", 80)
|
|
sock.send(data, length)
|
|
n = sock.recv(buf, bufsize)
|
|
sock.close()
|
|
"""
|
|
fd: INT
|
|
family: INT
|
|
type: INT
|
|
protocol: INT
|
|
|
|
def __init__(self, family: INT, type: INT, protocol: INT = 0):
|
|
"""Create Socket.
|
|
|
|
Args:
|
|
family: Address family (AF_INET)
|
|
type: Socket type (SOCK_STREAM)
|
|
protocol: Protocol (default 0, auto-match)
|
|
"""
|
|
self.family = family
|
|
self.type = type
|
|
self.protocol = protocol
|
|
|
|
if platmacro.IS_WINDOWS:
|
|
self.fd = INT(win32.winsock2.socket(family, type, protocol))
|
|
else:
|
|
self.fd = INT(socket(family, type, protocol))
|
|
|
|
def connect(self, host: str, port: INT) -> INT:
|
|
"""Connect to remote host.
|
|
|
|
Args:
|
|
host: Hostname or IP
|
|
port: Port number
|
|
Returns:
|
|
0 success, non 0 otherwise
|
|
"""
|
|
if self.fd == INVALID_SOCKET: return -1
|
|
addr: SocketAddr = _MakeAddr(host, port)
|
|
|
|
result: INT = 0
|
|
if platmacro.IS_WINDOWS:
|
|
result = win32.winsock2.connect(u64(self.fd), &addr, 16)
|
|
if result != 0:
|
|
err: INT = win32.winsock2.WSAGetLastError()
|
|
else:
|
|
result = connect(self.fd, &addr, 16)
|
|
return result
|
|
|
|
def send(self, data: t.CPtr, length: INT) -> INT:
|
|
"""Send data.
|
|
|
|
Returns:
|
|
Sent bytes, SOCKET_ERROR on error
|
|
"""
|
|
if self.fd == INVALID_SOCKET: return SOCKET_ERROR
|
|
if data is None or length <= 0: return 0
|
|
|
|
if platmacro.IS_WINDOWS:
|
|
return win32.winsock2.send(u64(self.fd), data, length, 0)
|
|
return send(self.fd, data, length, 0)
|
|
|
|
def recv(self, buf: t.CPtr, bufsize: INT) -> INT:
|
|
"""Receive data.
|
|
|
|
Returns:
|
|
Received bytes, SOCKET_ERROR on error, 0 on connection close
|
|
"""
|
|
if self.fd == INVALID_SOCKET: return SOCKET_ERROR
|
|
if buf is None or bufsize <= 0: return 0
|
|
|
|
if platmacro.IS_WINDOWS:
|
|
return win32.winsock2.recv(u64(self.fd), buf, bufsize, 0)
|
|
return recv(self.fd, buf, bufsize, 0)
|
|
|
|
def sendall(self, data: t.CPtr, length: INT) -> INT:
|
|
"""Send all data.
|
|
|
|
Returns:
|
|
0 success, SOCKET_ERROR on error
|
|
"""
|
|
if self.fd == INVALID_SOCKET: return SOCKET_ERROR
|
|
if data is None or length <= 0: return 0
|
|
|
|
total: INT = 0
|
|
while total < length:
|
|
n: INT = self.send(t.CPtr(uint64(data) + total), length - total)
|
|
if n <= 0: return SOCKET_ERROR
|
|
total += n
|
|
return 0
|
|
|
|
def close(self) -> INT:
|
|
"""Close Socket."""
|
|
if self.fd == INVALID_SOCKET: return 0
|
|
result: INT = _CloseFd(self.fd)
|
|
self.fd = INVALID_SOCKET
|
|
return result
|
|
|
|
def bind(self, host: str, port: INT) -> INT:
|
|
"""Bind address and port (server-side).
|
|
|
|
Returns:
|
|
0 success, non 0 otherwise
|
|
"""
|
|
if self.fd == INVALID_SOCKET: return -1
|
|
addr: SocketAddr = _MakeAddr(host, port)
|
|
|
|
if platmacro.IS_WINDOWS:
|
|
return win32.winsock2.bind(u64(self.fd), &addr, 16)
|
|
return bind(self.fd, &addr, 16)
|
|
|
|
def listen(self, backlog: INT) -> INT:
|
|
"""Start for incoming connections (server-side).
|
|
|
|
Returns:
|
|
0 success, non 0 otherwise
|
|
"""
|
|
if self.fd == INVALID_SOCKET: return -1
|
|
|
|
if platmacro.IS_WINDOWS:
|
|
return win32.winsock2.listen(u64(self.fd), backlog)
|
|
return listen(self.fd, backlog)
|
|
|
|
def accept(self, out_client: 'Socket' | t.CPtr) -> INT:
|
|
"""Accept a client connection.
|
|
|
|
Args:
|
|
out_client: Output parameter, receives client Socket
|
|
Returns:
|
|
Client fd (>= 0 on success), SOCKET_ERROR on error
|
|
"""
|
|
if self.fd == INVALID_SOCKET: return SOCKET_ERROR
|
|
if out_client is None: return SOCKET_ERROR
|
|
|
|
addr: SocketAddr
|
|
string.memset(&addr, 0, 16)
|
|
addrlen: INT = 16
|
|
|
|
if platmacro.IS_WINDOWS:
|
|
client_fd: u64 = win32.winsock2.accept(u64(self.fd), &addr, &addrlen)
|
|
if client_fd == win32.winsock2.INVALID_SOCKET:
|
|
return SOCKET_ERROR
|
|
out_client.fd = INT(client_fd)
|
|
out_client.family = self.family
|
|
out_client.type = self.type
|
|
out_client.protocol = self.protocol
|
|
return INT(client_fd)
|
|
|
|
client_fd: INT = INT(accept(self.fd, &addr, &addrlen))
|
|
if client_fd == SOCKET_ERROR:
|
|
return SOCKET_ERROR
|
|
out_client.fd = INT(client_fd)
|
|
out_client.family = self.family
|
|
out_client.type = self.type
|
|
out_client.protocol = self.protocol
|
|
return INT(client_fd)
|
|
|
|
def settimeout(self, timeout_ms: INT) -> INT:
|
|
"""Set timeout for receive and send operations (milliseconds).
|
|
|
|
Returns:
|
|
0 success, non 0 otherwise
|
|
"""
|
|
if self.fd == INVALID_SOCKET: return -1
|
|
|
|
tv: timeval
|
|
tv.tv_sec = timeout_ms / 1000
|
|
tv.tv_usec = (timeout_ms % 1000) * 1000
|
|
optlen: INT = 8
|
|
|
|
if platmacro.IS_WINDOWS:
|
|
r1: INT = win32.winsock2.setsockopt(u64(self.fd), SOL_SOCKET, SO_RCVTIMEO, &tv, optlen)
|
|
r2: INT = win32.winsock2.setsockopt(u64(self.fd), SOL_SOCKET, SO_SNDTIMEO, &tv, optlen)
|
|
else:
|
|
r1 = setsockopt(self.fd, SOL_SOCKET, SO_RCVTIMEO, &tv, optlen)
|
|
r2 = setsockopt(self.fd, SOL_SOCKET, SO_SNDTIMEO, &tv, optlen)
|
|
|
|
if r1 != 0: return r1
|
|
return r2
|
|
|
|
def fileno(self) -> INT:
|
|
"""Return underlying file descriptor / SOCKET handle."""
|
|
return self.fd
|
|
|
|
|
|
# ============================================================
|
|
# Platform initialization
|
|
# ============================================================
|
|
|
|
def SocketInit() -> INT:
|
|
"""Initialize the Socket library. On Windows, call WSAStartup; on POSIX, do nothing."""
|
|
if platmacro.IS_WINDOWS:
|
|
wsa_data: t.CArray[BYTE, 400]
|
|
string.memset(&wsa_data, 0, 400)
|
|
return win32.winsock2.WSAStartup(win32.winsock2.WINSOCK_VERSION, &wsa_data)
|
|
return 0
|
|
|
|
|
|
def SocketCleanup() -> INT:
|
|
"""Cleanup the Socket library."""
|
|
if platmacro.IS_WINDOWS:
|
|
return win32.winsock2.WSACleanup()
|
|
return 0
|
|
|
|
|
|
# ============================================================
|
|
# Compatibility with old procedural API (delegates to Socket methods)
|
|
# ============================================================
|
|
|
|
def SocketCreate(family: INT, type: INT, protocol: INT, out_sock: Socket*) -> INT:
|
|
"""Create a Socket handle (compatible with old API). It is recommended to use Socket(...) constructor."""
|
|
if out_sock is None: return -1
|
|
out_sock.family = family
|
|
out_sock.type = type
|
|
out_sock.protocol = protocol
|
|
|
|
if platmacro.IS_WINDOWS:
|
|
fd: u64 = win32.winsock2.socket(family, type, protocol)
|
|
out_sock.fd = INT(fd)
|
|
if fd == win32.winsock2.INVALID_SOCKET: return -1
|
|
else:
|
|
fd: INT = INT(socket(family, type, protocol))
|
|
out_sock.fd = INT(fd)
|
|
if fd == INVALID_SOCKET: return -1
|
|
return 0
|
|
|
|
# But maybe using OOP would be better, though that can be wrapped up later.
|
|
|
|
def SocketConnect(sock: Socket*, host: str, port: INT) -> INT:
|
|
return sock.connect(host, port)
|
|
|
|
|
|
def SocketSend(sock: Socket*, data: t.CPtr, length: INT) -> INT:
|
|
return sock.send(data, length)
|
|
|
|
|
|
def SocketRecv(sock: Socket*, buf: t.CPtr, bufsize: INT) -> INT:
|
|
return sock.recv(buf, bufsize)
|
|
|
|
|
|
def SocketClose(sock: Socket*) -> INT:
|
|
return sock.close()
|
|
|
|
|
|
def SocketBind(sock: Socket*, host: str, port: INT) -> INT:
|
|
return sock.bind(host, port)
|
|
|
|
|
|
def SocketListen(sock: Socket*, backlog: INT) -> INT:
|
|
return sock.listen(backlog)
|
|
|
|
|
|
def SocketAccept(sock: Socket*, out_client: Socket*) -> INT:
|
|
return sock.accept(out_client)
|
|
|
|
|
|
def SocketSetTimeout(sock: Socket*, timeout_ms: INT) -> INT:
|
|
return sock.settimeout(timeout_ms)
|
|
|
|
|
|
def SocketIsValid(sock: Socket*) -> bool:
|
|
"""Check if a Socket is valid."""
|
|
if sock is None: return False
|
|
return sock.fd != INVALID_SOCKET |