diff --git a/README.md b/README.md index f2cc410..8aa999d 100644 --- a/README.md +++ b/README.md @@ -148,6 +148,16 @@ await tunnel.forward_to("127.0.0.1", 8000) It keeps accepting rstream streams and relays them to the local TCP service until the tunnel or control channel is closed. +The control channel uses a negotiated heartbeat deadline. If that deadline +expires, or the control transport ends unexpectedly, the SDK immediately stops +accepting new streams but lets already accepted streams and active +`forward_to()` relays drain normally. This avoids interrupting an established +application session because of a transient or asymmetric control-path outage. +An explicit tunnel/control close or a protocol violation remains a hard close +and terminates active local forwarding. Applications should use each stream's +EOF/error as its payload-lifecycle signal; `control.done()` only describes the +control plane. + ## Private dial ```python diff --git a/proto/rstream.proto b/proto/rstream.proto index 12d23db..84aa301 100644 --- a/proto/rstream.proto +++ b/proto/rstream.proto @@ -14,7 +14,7 @@ extend google.protobuf.FieldOptions { string access = 51234; } -option (protocol_version) = "1.4.4"; +option (protocol_version) = "1.4.5"; package rstream.io_rstrm.protobuf; @@ -136,6 +136,7 @@ message TunnelProperties { message OpenControlChannelReq { ClientDetails client_details = 1; + ControlChannelLiveness liveness = 2; } // The server responds to an 'OpenControlChannelReq' message with an @@ -148,6 +149,7 @@ message OpenControlChannelRsp { message Ok { string client_id = 1; ServerDetails server_details = 2; + ControlChannelLiveness liveness = 3; } oneof payload { Ok ok = 1; @@ -263,7 +265,15 @@ message DatagramChannelClose { // Sent by client and/or server to maintain the control channel active -message Heartbeat { } +message ControlChannelLiveness { + uint32 heartbeat_interval_ms = 1; + uint32 heartbeat_timeout_ms = 2; +} + +message Heartbeat { + uint64 sequence = 1; + uint64 acknowledgement = 2; +} // Allows the server to send unsolicited messages to the client message ServerMessage { diff --git a/src/rstream/_proto/rstream_pb2.py b/src/rstream/_proto/rstream_pb2.py index 0296de7..6d682a4 100644 --- a/src/rstream/_proto/rstream_pb2.py +++ b/src/rstream/_proto/rstream_pb2.py @@ -2,7 +2,7 @@ # Generated by the protocol buffer compiler. DO NOT EDIT! # NO CHECKED-IN PROTOBUF GENCODE # source: rstream.proto -# Protobuf Python Version: 6.31.1 +# Protobuf Python Version: 7.35.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool @@ -11,8 +11,8 @@ from google.protobuf.internal import builder as _builder _runtime_version.ValidateProtobufRuntimeVersion( _runtime_version.Domain.PUBLIC, - 6, - 31, + 7, + 35, 1, '', 'rstream.proto' @@ -27,14 +27,14 @@ from google.protobuf import wrappers_pb2 as google_dot_protobuf_dot_wrappers__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\rrstream.proto\x12\x19rstream.io_rstrm.protobuf\x1a google/protobuf/descriptor.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1egoogle/protobuf/wrappers.proto\"/\n\tIpAddress\x12\x0c\n\x02v4\x18\x01 \x01(\x07H\x00\x12\x0c\n\x02v6\x18\x02 \x01(\x0cH\x00\x42\x06\n\x04\x61\x64\x64r\"j\n\x05\x45rror\x12\x32\n\x04\x63ode\x18\x01 \x01(\x0e\x32$.rstream.io_rstrm.protobuf.ErrorCode\x12-\n\x07message\x18\x02 \x01(\x0b\x32\x1c.google.protobuf.StringValue\"\xd5\x02\n\rClientDetails\x12+\n\x05\x61gent\x18\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07\x63hannel\x18\x02 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07version\x18\x03 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12(\n\x02os\x18\x04 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x04\x61rch\x18\x05 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x05token\x18\x06 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x10protocol_version\x18\x07 \x01(\x0b\x32\x1c.google.protobuf.StringValue\"\xd2\x02\n\rServerDetails\x12+\n\x05\x61gent\x18\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07\x63hannel\x18\x02 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07version\x18\x03 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x04plan\x18\x04 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08provider\x18\x05 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06region\x18\x06 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06update\x18\x07 \x01(\x0b\x32\x1c.google.protobuf.StringValue\"\x90\r\n\x10TunnelProperties\x12\x37\n\x02id\x18\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\r\x92\x82\x19\tread-only\x12@\n\rcreation_date\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.TimestampB\r\x92\x82\x19\tread-only\x12:\n\x04name\x18\x03 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12:\n\x04type\x18\x04 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12;\n\x07publish\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12>\n\x08protocol\x18\x06 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12W\n\x06labels\x18\x07 \x03(\x0b\x32\x37.rstream.io_rstrm.protobuf.TunnelProperties.LabelsEntryB\x0e\x92\x82\x19\nread-write\x12\x1d\n\x05geoip\x18\x08 \x03(\tB\x0e\x92\x82\x19\nread-write\x12#\n\x0btrusted_ips\x18\t \x03(\tB\x0e\x92\x82\x19\nread-write\x12;\n\x04host\x18\n \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0f\x18\x01\x92\x82\x19\tread-only\x12>\n\x08tls_mode\x18\x0b \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12!\n\ttls_alpns\x18\x0c \x03(\tB\x0e\x92\x82\x19\nread-write\x12\x45\n\x0ftls_min_version\x18\r \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12#\n\x0btls_ciphers\x18\x0e \x03(\tB\x0e\x92\x82\x19\nread-write\x12=\n\tmtls_auth\x18\x0f \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12G\n\x0fmtls_cacert_pem\x18\x10 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x10\x18\x01\x92\x82\x19\nread-write\x12\x42\n\x0chttp_version\x18\x11 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12\x42\n\x0chttp_use_tls\x18\x12 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x10\x18\x01\x92\x82\x19\nread-write\x12>\n\ntoken_auth\x18\x13 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12@\n\x0crstream_auth\x18\x14 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12\x42\n\x0e\x63hallenge_mode\x18\x15 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12>\n\x08hostname\x18\x16 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12:\n\x04port\x18\x17 \x01(\x0b\x32\x1c.google.protobuf.UInt32ValueB\x0e\x92\x82\x19\nread-write\x12@\n\x0cupstream_tls\x18\x18 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12P\n\x1c\x64\x61tagram_guaranteed_delivery\x18\x19 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12N\n\x1a\x61llow_cross_region_routing\x18\x1a \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x1a-\n\x0bLabelsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"Y\n\x15OpenControlChannelReq\x12@\n\x0e\x63lient_details\x18\x01 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.ClientDetails\"\xf3\x01\n\x15OpenControlChannelRsp\x12\x41\n\x02ok\x18\x01 \x01(\x0b\x32\x33.rstream.io_rstrm.protobuf.OpenControlChannelRsp.OkH\x00\x12\x31\n\x05\x65rror\x18\x02 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.ErrorH\x00\x1aY\n\x02Ok\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12@\n\x0eserver_details\x18\x02 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.ServerDetailsB\t\n\x07payload\"\x18\n\x16\x43loseControlChannelReq\"\x18\n\x16\x43loseControlChannelRsp\"k\n\rOpenTunnelReq\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x46\n\x11tunnel_properties\x18\x02 \x01(\x0b\x32+.rstream.io_rstrm.protobuf.TunnelProperties\"\xab\x01\n\rOpenTunnelRsp\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12H\n\x11tunnel_properties\x18\x02 \x01(\x0b\x32+.rstream.io_rstrm.protobuf.TunnelPropertiesH\x00\x12\x31\n\x05\x65rror\x18\x03 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.ErrorH\x00\x42\t\n\x07payload\"#\n\x0e\x43loseTunnelReq\x12\x11\n\ttunnel_id\x18\x01 \x01(\t\"#\n\x0e\x43loseTunnelRsp\x12\x11\n\ttunnel_id\x18\x01 \x01(\t\"\xd1\x01\n\x0cProxyConnReq\x12\x11\n\ttunnel_id\x18\x01 \x01(\t\x12\x11\n\tstream_id\x18\x02 \x01(\t\x12,\n\x06secret\x18\x03 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\tsource_ip\x18\x04 \x01(\x0b\x32$.rstream.io_rstrm.protobuf.IpAddress\x12\x34\n\x0eproxy_endpoint\x18\x05 \x01(\x0b\x32\x1c.google.protobuf.StringValue\"R\n\x0cProxyConnRsp\x12\x11\n\tstream_id\x18\x01 \x01(\t\x12/\n\x05\x65rror\x18\x02 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.Error\"\x8d\x01\n\x08ProxyReq\x12@\n\x0e\x63lient_details\x18\x01 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.ClientDetails\x12\x11\n\tstream_id\x18\x02 \x01(\t\x12,\n\x08zero_rtt\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\";\n\x08ProxyRsp\x12/\n\x05\x65rror\x18\x01 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.Error\"\xc9\x01\n\tStreamReq\x12@\n\x0e\x63lient_details\x18\x01 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.ClientDetails\x12\x16\n\x0etunnel_id_name\x18\x02 \x01(\t\x12,\n\x08zero_rtt\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x10\x64\x61tagram_channel\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\"^\n\tStreamRsp\x12\x13\n\tstream_id\x18\x01 \x01(\tH\x00\x12\x31\n\x05\x65rror\x18\x02 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.ErrorH\x00\x42\t\n\x07payload\"Z\n\x14\x44\x61tagramChannelClose\x12\x11\n\tstream_id\x18\x01 \x01(\t\x12/\n\x05\x65rror\x18\x02 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.Error\"\x0b\n\tHeartbeat\" \n\rServerMessage\x12\x0f\n\x07message\x18\x01 \x01(\t\"\xcc\t\n\x07Message\x12T\n\x18open_control_channel_req\x18\x01 \x01(\x0b\x32\x30.rstream.io_rstrm.protobuf.OpenControlChannelReqH\x00\x12T\n\x18open_control_channel_rsp\x18\x02 \x01(\x0b\x32\x30.rstream.io_rstrm.protobuf.OpenControlChannelRspH\x00\x12V\n\x19\x63lose_control_channel_req\x18\x03 \x01(\x0b\x32\x31.rstream.io_rstrm.protobuf.CloseControlChannelReqH\x00\x12V\n\x19\x63lose_control_channel_rsp\x18\x04 \x01(\x0b\x32\x31.rstream.io_rstrm.protobuf.CloseControlChannelRspH\x00\x12\x43\n\x0fopen_tunnel_req\x18\x05 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.OpenTunnelReqH\x00\x12\x43\n\x0fopen_tunnel_rsp\x18\x06 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.OpenTunnelRspH\x00\x12\x45\n\x10\x63lose_tunnel_req\x18\x07 \x01(\x0b\x32).rstream.io_rstrm.protobuf.CloseTunnelReqH\x00\x12\x45\n\x10\x63lose_tunnel_rsp\x18\x08 \x01(\x0b\x32).rstream.io_rstrm.protobuf.CloseTunnelRspH\x00\x12\x41\n\x0eproxy_conn_req\x18\t \x01(\x0b\x32\'.rstream.io_rstrm.protobuf.ProxyConnReqH\x00\x12\x41\n\x0eproxy_conn_rsp\x18\n \x01(\x0b\x32\'.rstream.io_rstrm.protobuf.ProxyConnRspH\x00\x12\x38\n\tproxy_req\x18\x0b \x01(\x0b\x32#.rstream.io_rstrm.protobuf.ProxyReqH\x00\x12\x38\n\tproxy_rsp\x18\x0c \x01(\x0b\x32#.rstream.io_rstrm.protobuf.ProxyRspH\x00\x12:\n\nstream_req\x18\r \x01(\x0b\x32$.rstream.io_rstrm.protobuf.StreamReqH\x00\x12:\n\nstream_rsp\x18\x0e \x01(\x0b\x32$.rstream.io_rstrm.protobuf.StreamRspH\x00\x12\x39\n\theartbeat\x18\x0f \x01(\x0b\x32$.rstream.io_rstrm.protobuf.HeartbeatH\x00\x12\x42\n\x0eserver_message\x18\x10 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.ServerMessageH\x00\x12Q\n\x16\x64\x61tagram_channel_close\x18\x11 \x01(\x0b\x32/.rstream.io_rstrm.protobuf.DatagramChannelCloseH\x00\x42\t\n\x07payload*\xb5\x03\n\tErrorCode\x12\x1a\n\x16\x45RROR_CODE_UNSPECIFIED\x10\x00\x12\x1c\n\x17\x45RROR_CODE_UNAUTHORIZED\x10\xe8\x07\x12\x1f\n\x1a\x45RROR_CODE_INVALID_REQUEST\x10\xd0\x0f\x12(\n#ERROR_CODE_PROTOCOL_VERSION_MISSING\x10\xda\x0f\x12(\n#ERROR_CODE_PROTOCOL_VERSION_INVALID\x10\xe4\x0f\x12-\n(ERROR_CODE_PROTOCOL_VERSION_INCOMPATIBLE\x10\xee\x0f\x12 \n\x1b\x45RROR_CODE_TUNNEL_NOT_FOUND\x10\xb8\x17\x12\x1e\n\x19\x45RROR_CODE_INVALID_STREAM\x10\xa0\x1f\x12%\n ERROR_CODE_FEATURE_NOT_AVAILABLE\x10\x88\'\x12#\n\x1e\x45RROR_CODE_SERVICE_UNAVAILABLE\x10\xf0.\x12\"\n\x1d\x45RROR_CODE_CAPACITY_EXHAUSTED\x10\xd8\x36\x12\x18\n\x13\x45RROR_CODE_INTERNAL\x10\xa8\x46:8\n\x10protocol_version\x12\x1c.google.protobuf.FileOptions\x18\xb8\x8e\x03 \x01(\t:/\n\x06\x61\x63\x63\x65ss\x12\x1d.google.protobuf.FieldOptions\x18\xa2\x90\x03 \x01(\tB\x0fZ\x04./pb\xc2\xf3\x18\x05\x31.4.4b\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\rrstream.proto\x12\x19rstream.io_rstrm.protobuf\x1a google/protobuf/descriptor.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1egoogle/protobuf/wrappers.proto\"/\n\tIpAddress\x12\x0c\n\x02v4\x18\x01 \x01(\x07H\x00\x12\x0c\n\x02v6\x18\x02 \x01(\x0cH\x00\x42\x06\n\x04\x61\x64\x64r\"j\n\x05\x45rror\x12\x32\n\x04\x63ode\x18\x01 \x01(\x0e\x32$.rstream.io_rstrm.protobuf.ErrorCode\x12-\n\x07message\x18\x02 \x01(\x0b\x32\x1c.google.protobuf.StringValue\"\xd5\x02\n\rClientDetails\x12+\n\x05\x61gent\x18\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07\x63hannel\x18\x02 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07version\x18\x03 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12(\n\x02os\x18\x04 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x04\x61rch\x18\x05 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x05token\x18\x06 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x10protocol_version\x18\x07 \x01(\x0b\x32\x1c.google.protobuf.StringValue\"\xd2\x02\n\rServerDetails\x12+\n\x05\x61gent\x18\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07\x63hannel\x18\x02 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07version\x18\x03 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x04plan\x18\x04 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08provider\x18\x05 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06region\x18\x06 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06update\x18\x07 \x01(\x0b\x32\x1c.google.protobuf.StringValue\"\x90\r\n\x10TunnelProperties\x12\x37\n\x02id\x18\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\r\x92\x82\x19\tread-only\x12@\n\rcreation_date\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.TimestampB\r\x92\x82\x19\tread-only\x12:\n\x04name\x18\x03 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12:\n\x04type\x18\x04 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12;\n\x07publish\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12>\n\x08protocol\x18\x06 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12W\n\x06labels\x18\x07 \x03(\x0b\x32\x37.rstream.io_rstrm.protobuf.TunnelProperties.LabelsEntryB\x0e\x92\x82\x19\nread-write\x12\x1d\n\x05geoip\x18\x08 \x03(\tB\x0e\x92\x82\x19\nread-write\x12#\n\x0btrusted_ips\x18\t \x03(\tB\x0e\x92\x82\x19\nread-write\x12;\n\x04host\x18\n \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0f\x18\x01\x92\x82\x19\tread-only\x12>\n\x08tls_mode\x18\x0b \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12!\n\ttls_alpns\x18\x0c \x03(\tB\x0e\x92\x82\x19\nread-write\x12\x45\n\x0ftls_min_version\x18\r \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12#\n\x0btls_ciphers\x18\x0e \x03(\tB\x0e\x92\x82\x19\nread-write\x12=\n\tmtls_auth\x18\x0f \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12G\n\x0fmtls_cacert_pem\x18\x10 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x10\x18\x01\x92\x82\x19\nread-write\x12\x42\n\x0chttp_version\x18\x11 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12\x42\n\x0chttp_use_tls\x18\x12 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x10\x18\x01\x92\x82\x19\nread-write\x12>\n\ntoken_auth\x18\x13 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12@\n\x0crstream_auth\x18\x14 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12\x42\n\x0e\x63hallenge_mode\x18\x15 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12>\n\x08hostname\x18\x16 \x01(\x0b\x32\x1c.google.protobuf.StringValueB\x0e\x92\x82\x19\nread-write\x12:\n\x04port\x18\x17 \x01(\x0b\x32\x1c.google.protobuf.UInt32ValueB\x0e\x92\x82\x19\nread-write\x12@\n\x0cupstream_tls\x18\x18 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12P\n\x1c\x64\x61tagram_guaranteed_delivery\x18\x19 \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x12N\n\x1a\x61llow_cross_region_routing\x18\x1a \x01(\x0b\x32\x1a.google.protobuf.BoolValueB\x0e\x92\x82\x19\nread-write\x1a-\n\x0bLabelsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x9e\x01\n\x15OpenControlChannelReq\x12@\n\x0e\x63lient_details\x18\x01 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.ClientDetails\x12\x43\n\x08liveness\x18\x02 \x01(\x0b\x32\x31.rstream.io_rstrm.protobuf.ControlChannelLiveness\"\xb9\x02\n\x15OpenControlChannelRsp\x12\x41\n\x02ok\x18\x01 \x01(\x0b\x32\x33.rstream.io_rstrm.protobuf.OpenControlChannelRsp.OkH\x00\x12\x31\n\x05\x65rror\x18\x02 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.ErrorH\x00\x1a\x9e\x01\n\x02Ok\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12@\n\x0eserver_details\x18\x02 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.ServerDetails\x12\x43\n\x08liveness\x18\x03 \x01(\x0b\x32\x31.rstream.io_rstrm.protobuf.ControlChannelLivenessB\t\n\x07payload\"\x18\n\x16\x43loseControlChannelReq\"\x18\n\x16\x43loseControlChannelRsp\"k\n\rOpenTunnelReq\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x46\n\x11tunnel_properties\x18\x02 \x01(\x0b\x32+.rstream.io_rstrm.protobuf.TunnelProperties\"\xab\x01\n\rOpenTunnelRsp\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12H\n\x11tunnel_properties\x18\x02 \x01(\x0b\x32+.rstream.io_rstrm.protobuf.TunnelPropertiesH\x00\x12\x31\n\x05\x65rror\x18\x03 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.ErrorH\x00\x42\t\n\x07payload\"#\n\x0e\x43loseTunnelReq\x12\x11\n\ttunnel_id\x18\x01 \x01(\t\"#\n\x0e\x43loseTunnelRsp\x12\x11\n\ttunnel_id\x18\x01 \x01(\t\"\xd1\x01\n\x0cProxyConnReq\x12\x11\n\ttunnel_id\x18\x01 \x01(\t\x12\x11\n\tstream_id\x18\x02 \x01(\t\x12,\n\x06secret\x18\x03 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\tsource_ip\x18\x04 \x01(\x0b\x32$.rstream.io_rstrm.protobuf.IpAddress\x12\x34\n\x0eproxy_endpoint\x18\x05 \x01(\x0b\x32\x1c.google.protobuf.StringValue\"R\n\x0cProxyConnRsp\x12\x11\n\tstream_id\x18\x01 \x01(\t\x12/\n\x05\x65rror\x18\x02 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.Error\"\x8d\x01\n\x08ProxyReq\x12@\n\x0e\x63lient_details\x18\x01 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.ClientDetails\x12\x11\n\tstream_id\x18\x02 \x01(\t\x12,\n\x08zero_rtt\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\";\n\x08ProxyRsp\x12/\n\x05\x65rror\x18\x01 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.Error\"\xc9\x01\n\tStreamReq\x12@\n\x0e\x63lient_details\x18\x01 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.ClientDetails\x12\x16\n\x0etunnel_id_name\x18\x02 \x01(\t\x12,\n\x08zero_rtt\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x10\x64\x61tagram_channel\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\"^\n\tStreamRsp\x12\x13\n\tstream_id\x18\x01 \x01(\tH\x00\x12\x31\n\x05\x65rror\x18\x02 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.ErrorH\x00\x42\t\n\x07payload\"Z\n\x14\x44\x61tagramChannelClose\x12\x11\n\tstream_id\x18\x01 \x01(\t\x12/\n\x05\x65rror\x18\x02 \x01(\x0b\x32 .rstream.io_rstrm.protobuf.Error\"U\n\x16\x43ontrolChannelLiveness\x12\x1d\n\x15heartbeat_interval_ms\x18\x01 \x01(\r\x12\x1c\n\x14heartbeat_timeout_ms\x18\x02 \x01(\r\"6\n\tHeartbeat\x12\x10\n\x08sequence\x18\x01 \x01(\x04\x12\x17\n\x0f\x61\x63knowledgement\x18\x02 \x01(\x04\" \n\rServerMessage\x12\x0f\n\x07message\x18\x01 \x01(\t\"\xcc\t\n\x07Message\x12T\n\x18open_control_channel_req\x18\x01 \x01(\x0b\x32\x30.rstream.io_rstrm.protobuf.OpenControlChannelReqH\x00\x12T\n\x18open_control_channel_rsp\x18\x02 \x01(\x0b\x32\x30.rstream.io_rstrm.protobuf.OpenControlChannelRspH\x00\x12V\n\x19\x63lose_control_channel_req\x18\x03 \x01(\x0b\x32\x31.rstream.io_rstrm.protobuf.CloseControlChannelReqH\x00\x12V\n\x19\x63lose_control_channel_rsp\x18\x04 \x01(\x0b\x32\x31.rstream.io_rstrm.protobuf.CloseControlChannelRspH\x00\x12\x43\n\x0fopen_tunnel_req\x18\x05 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.OpenTunnelReqH\x00\x12\x43\n\x0fopen_tunnel_rsp\x18\x06 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.OpenTunnelRspH\x00\x12\x45\n\x10\x63lose_tunnel_req\x18\x07 \x01(\x0b\x32).rstream.io_rstrm.protobuf.CloseTunnelReqH\x00\x12\x45\n\x10\x63lose_tunnel_rsp\x18\x08 \x01(\x0b\x32).rstream.io_rstrm.protobuf.CloseTunnelRspH\x00\x12\x41\n\x0eproxy_conn_req\x18\t \x01(\x0b\x32\'.rstream.io_rstrm.protobuf.ProxyConnReqH\x00\x12\x41\n\x0eproxy_conn_rsp\x18\n \x01(\x0b\x32\'.rstream.io_rstrm.protobuf.ProxyConnRspH\x00\x12\x38\n\tproxy_req\x18\x0b \x01(\x0b\x32#.rstream.io_rstrm.protobuf.ProxyReqH\x00\x12\x38\n\tproxy_rsp\x18\x0c \x01(\x0b\x32#.rstream.io_rstrm.protobuf.ProxyRspH\x00\x12:\n\nstream_req\x18\r \x01(\x0b\x32$.rstream.io_rstrm.protobuf.StreamReqH\x00\x12:\n\nstream_rsp\x18\x0e \x01(\x0b\x32$.rstream.io_rstrm.protobuf.StreamRspH\x00\x12\x39\n\theartbeat\x18\x0f \x01(\x0b\x32$.rstream.io_rstrm.protobuf.HeartbeatH\x00\x12\x42\n\x0eserver_message\x18\x10 \x01(\x0b\x32(.rstream.io_rstrm.protobuf.ServerMessageH\x00\x12Q\n\x16\x64\x61tagram_channel_close\x18\x11 \x01(\x0b\x32/.rstream.io_rstrm.protobuf.DatagramChannelCloseH\x00\x42\t\n\x07payload*\xb5\x03\n\tErrorCode\x12\x1a\n\x16\x45RROR_CODE_UNSPECIFIED\x10\x00\x12\x1c\n\x17\x45RROR_CODE_UNAUTHORIZED\x10\xe8\x07\x12\x1f\n\x1a\x45RROR_CODE_INVALID_REQUEST\x10\xd0\x0f\x12(\n#ERROR_CODE_PROTOCOL_VERSION_MISSING\x10\xda\x0f\x12(\n#ERROR_CODE_PROTOCOL_VERSION_INVALID\x10\xe4\x0f\x12-\n(ERROR_CODE_PROTOCOL_VERSION_INCOMPATIBLE\x10\xee\x0f\x12 \n\x1b\x45RROR_CODE_TUNNEL_NOT_FOUND\x10\xb8\x17\x12\x1e\n\x19\x45RROR_CODE_INVALID_STREAM\x10\xa0\x1f\x12%\n ERROR_CODE_FEATURE_NOT_AVAILABLE\x10\x88\'\x12#\n\x1e\x45RROR_CODE_SERVICE_UNAVAILABLE\x10\xf0.\x12\"\n\x1d\x45RROR_CODE_CAPACITY_EXHAUSTED\x10\xd8\x36\x12\x18\n\x13\x45RROR_CODE_INTERNAL\x10\xa8\x46:8\n\x10protocol_version\x12\x1c.google.protobuf.FileOptions\x18\xb8\x8e\x03 \x01(\t:/\n\x06\x61\x63\x63\x65ss\x12\x1d.google.protobuf.FieldOptions\x18\xa2\x90\x03 \x01(\tB\x0fZ\x04./pb\xc2\xf3\x18\x05\x31.4.5b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'rstream_pb2', _globals) if not _descriptor._USE_C_DESCRIPTORS: _globals['DESCRIPTOR']._loaded_options = None - _globals['DESCRIPTOR']._serialized_options = b'Z\004./pb\302\363\030\0051.4.4' + _globals['DESCRIPTOR']._serialized_options = b'Z\004./pb\302\363\030\0051.4.5' _globals['_TUNNELPROPERTIES_LABELSENTRY']._loaded_options = None _globals['_TUNNELPROPERTIES_LABELSENTRY']._serialized_options = b'8\001' _globals['_TUNNELPROPERTIES'].fields_by_name['id']._loaded_options = None @@ -89,8 +89,8 @@ _globals['_TUNNELPROPERTIES'].fields_by_name['datagram_guaranteed_delivery']._serialized_options = b'\222\202\031\nread-write' _globals['_TUNNELPROPERTIES'].fields_by_name['allow_cross_region_routing']._loaded_options = None _globals['_TUNNELPROPERTIES'].fields_by_name['allow_cross_region_routing']._serialized_options = b'\222\202\031\nread-write' - _globals['_ERRORCODE']._serialized_start=5586 - _globals['_ERRORCODE']._serialized_end=6023 + _globals['_ERRORCODE']._serialized_start=5856 + _globals['_ERRORCODE']._serialized_end=6293 _globals['_IPADDRESS']._serialized_start=143 _globals['_IPADDRESS']._serialized_end=190 _globals['_ERROR']._serialized_start=192 @@ -103,42 +103,44 @@ _globals['_TUNNELPROPERTIES']._serialized_end=2666 _globals['_TUNNELPROPERTIES_LABELSENTRY']._serialized_start=2621 _globals['_TUNNELPROPERTIES_LABELSENTRY']._serialized_end=2666 - _globals['_OPENCONTROLCHANNELREQ']._serialized_start=2668 - _globals['_OPENCONTROLCHANNELREQ']._serialized_end=2757 - _globals['_OPENCONTROLCHANNELRSP']._serialized_start=2760 - _globals['_OPENCONTROLCHANNELRSP']._serialized_end=3003 - _globals['_OPENCONTROLCHANNELRSP_OK']._serialized_start=2903 - _globals['_OPENCONTROLCHANNELRSP_OK']._serialized_end=2992 - _globals['_CLOSECONTROLCHANNELREQ']._serialized_start=3005 - _globals['_CLOSECONTROLCHANNELREQ']._serialized_end=3029 - _globals['_CLOSECONTROLCHANNELRSP']._serialized_start=3031 - _globals['_CLOSECONTROLCHANNELRSP']._serialized_end=3055 - _globals['_OPENTUNNELREQ']._serialized_start=3057 - _globals['_OPENTUNNELREQ']._serialized_end=3164 - _globals['_OPENTUNNELRSP']._serialized_start=3167 - _globals['_OPENTUNNELRSP']._serialized_end=3338 - _globals['_CLOSETUNNELREQ']._serialized_start=3340 - _globals['_CLOSETUNNELREQ']._serialized_end=3375 - _globals['_CLOSETUNNELRSP']._serialized_start=3377 - _globals['_CLOSETUNNELRSP']._serialized_end=3412 - _globals['_PROXYCONNREQ']._serialized_start=3415 - _globals['_PROXYCONNREQ']._serialized_end=3624 - _globals['_PROXYCONNRSP']._serialized_start=3626 - _globals['_PROXYCONNRSP']._serialized_end=3708 - _globals['_PROXYREQ']._serialized_start=3711 - _globals['_PROXYREQ']._serialized_end=3852 - _globals['_PROXYRSP']._serialized_start=3854 - _globals['_PROXYRSP']._serialized_end=3913 - _globals['_STREAMREQ']._serialized_start=3916 - _globals['_STREAMREQ']._serialized_end=4117 - _globals['_STREAMRSP']._serialized_start=4119 - _globals['_STREAMRSP']._serialized_end=4213 - _globals['_DATAGRAMCHANNELCLOSE']._serialized_start=4215 - _globals['_DATAGRAMCHANNELCLOSE']._serialized_end=4305 - _globals['_HEARTBEAT']._serialized_start=4307 - _globals['_HEARTBEAT']._serialized_end=4318 - _globals['_SERVERMESSAGE']._serialized_start=4320 - _globals['_SERVERMESSAGE']._serialized_end=4352 - _globals['_MESSAGE']._serialized_start=4355 - _globals['_MESSAGE']._serialized_end=5583 + _globals['_OPENCONTROLCHANNELREQ']._serialized_start=2669 + _globals['_OPENCONTROLCHANNELREQ']._serialized_end=2827 + _globals['_OPENCONTROLCHANNELRSP']._serialized_start=2830 + _globals['_OPENCONTROLCHANNELRSP']._serialized_end=3143 + _globals['_OPENCONTROLCHANNELRSP_OK']._serialized_start=2974 + _globals['_OPENCONTROLCHANNELRSP_OK']._serialized_end=3132 + _globals['_CLOSECONTROLCHANNELREQ']._serialized_start=3145 + _globals['_CLOSECONTROLCHANNELREQ']._serialized_end=3169 + _globals['_CLOSECONTROLCHANNELRSP']._serialized_start=3171 + _globals['_CLOSECONTROLCHANNELRSP']._serialized_end=3195 + _globals['_OPENTUNNELREQ']._serialized_start=3197 + _globals['_OPENTUNNELREQ']._serialized_end=3304 + _globals['_OPENTUNNELRSP']._serialized_start=3307 + _globals['_OPENTUNNELRSP']._serialized_end=3478 + _globals['_CLOSETUNNELREQ']._serialized_start=3480 + _globals['_CLOSETUNNELREQ']._serialized_end=3515 + _globals['_CLOSETUNNELRSP']._serialized_start=3517 + _globals['_CLOSETUNNELRSP']._serialized_end=3552 + _globals['_PROXYCONNREQ']._serialized_start=3555 + _globals['_PROXYCONNREQ']._serialized_end=3764 + _globals['_PROXYCONNRSP']._serialized_start=3766 + _globals['_PROXYCONNRSP']._serialized_end=3848 + _globals['_PROXYREQ']._serialized_start=3851 + _globals['_PROXYREQ']._serialized_end=3992 + _globals['_PROXYRSP']._serialized_start=3994 + _globals['_PROXYRSP']._serialized_end=4053 + _globals['_STREAMREQ']._serialized_start=4056 + _globals['_STREAMREQ']._serialized_end=4257 + _globals['_STREAMRSP']._serialized_start=4259 + _globals['_STREAMRSP']._serialized_end=4353 + _globals['_DATAGRAMCHANNELCLOSE']._serialized_start=4355 + _globals['_DATAGRAMCHANNELCLOSE']._serialized_end=4445 + _globals['_CONTROLCHANNELLIVENESS']._serialized_start=4447 + _globals['_CONTROLCHANNELLIVENESS']._serialized_end=4532 + _globals['_HEARTBEAT']._serialized_start=4534 + _globals['_HEARTBEAT']._serialized_end=4588 + _globals['_SERVERMESSAGE']._serialized_start=4590 + _globals['_SERVERMESSAGE']._serialized_end=4622 + _globals['_MESSAGE']._serialized_start=4625 + _globals['_MESSAGE']._serialized_end=5853 # @@protoc_insertion_point(module_scope) diff --git a/src/rstream/_proto/rstream_pb2.pyi b/src/rstream/_proto/rstream_pb2.pyi index 6d4241e..f7c2bfc 100644 --- a/src/rstream/_proto/rstream_pb2.pyi +++ b/src/rstream/_proto/rstream_pb2.pyi @@ -159,20 +159,24 @@ class TunnelProperties(_message.Message): def __init__(self, id: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ..., creation_date: _Optional[_Union[datetime.datetime, _timestamp_pb2.Timestamp, _Mapping]] = ..., name: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ..., type: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ..., publish: _Optional[_Union[_wrappers_pb2.BoolValue, _Mapping]] = ..., protocol: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ..., labels: _Optional[_Mapping[str, str]] = ..., geoip: _Optional[_Iterable[str]] = ..., trusted_ips: _Optional[_Iterable[str]] = ..., host: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ..., tls_mode: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ..., tls_alpns: _Optional[_Iterable[str]] = ..., tls_min_version: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ..., tls_ciphers: _Optional[_Iterable[str]] = ..., mtls_auth: _Optional[_Union[_wrappers_pb2.BoolValue, _Mapping]] = ..., mtls_cacert_pem: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ..., http_version: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ..., http_use_tls: _Optional[_Union[_wrappers_pb2.BoolValue, _Mapping]] = ..., token_auth: _Optional[_Union[_wrappers_pb2.BoolValue, _Mapping]] = ..., rstream_auth: _Optional[_Union[_wrappers_pb2.BoolValue, _Mapping]] = ..., challenge_mode: _Optional[_Union[_wrappers_pb2.BoolValue, _Mapping]] = ..., hostname: _Optional[_Union[_wrappers_pb2.StringValue, _Mapping]] = ..., port: _Optional[_Union[_wrappers_pb2.UInt32Value, _Mapping]] = ..., upstream_tls: _Optional[_Union[_wrappers_pb2.BoolValue, _Mapping]] = ..., datagram_guaranteed_delivery: _Optional[_Union[_wrappers_pb2.BoolValue, _Mapping]] = ..., allow_cross_region_routing: _Optional[_Union[_wrappers_pb2.BoolValue, _Mapping]] = ...) -> None: ... class OpenControlChannelReq(_message.Message): - __slots__ = ("client_details",) + __slots__ = ("client_details", "liveness") CLIENT_DETAILS_FIELD_NUMBER: _ClassVar[int] + LIVENESS_FIELD_NUMBER: _ClassVar[int] client_details: ClientDetails - def __init__(self, client_details: _Optional[_Union[ClientDetails, _Mapping]] = ...) -> None: ... + liveness: ControlChannelLiveness + def __init__(self, client_details: _Optional[_Union[ClientDetails, _Mapping]] = ..., liveness: _Optional[_Union[ControlChannelLiveness, _Mapping]] = ...) -> None: ... class OpenControlChannelRsp(_message.Message): __slots__ = ("ok", "error") class Ok(_message.Message): - __slots__ = ("client_id", "server_details") + __slots__ = ("client_id", "server_details", "liveness") CLIENT_ID_FIELD_NUMBER: _ClassVar[int] SERVER_DETAILS_FIELD_NUMBER: _ClassVar[int] + LIVENESS_FIELD_NUMBER: _ClassVar[int] client_id: str server_details: ServerDetails - def __init__(self, client_id: _Optional[str] = ..., server_details: _Optional[_Union[ServerDetails, _Mapping]] = ...) -> None: ... + liveness: ControlChannelLiveness + def __init__(self, client_id: _Optional[str] = ..., server_details: _Optional[_Union[ServerDetails, _Mapping]] = ..., liveness: _Optional[_Union[ControlChannelLiveness, _Mapping]] = ...) -> None: ... OK_FIELD_NUMBER: _ClassVar[int] ERROR_FIELD_NUMBER: _ClassVar[int] ok: OpenControlChannelRsp.Ok @@ -283,9 +287,21 @@ class DatagramChannelClose(_message.Message): error: Error def __init__(self, stream_id: _Optional[str] = ..., error: _Optional[_Union[Error, _Mapping]] = ...) -> None: ... +class ControlChannelLiveness(_message.Message): + __slots__ = ("heartbeat_interval_ms", "heartbeat_timeout_ms") + HEARTBEAT_INTERVAL_MS_FIELD_NUMBER: _ClassVar[int] + HEARTBEAT_TIMEOUT_MS_FIELD_NUMBER: _ClassVar[int] + heartbeat_interval_ms: int + heartbeat_timeout_ms: int + def __init__(self, heartbeat_interval_ms: _Optional[int] = ..., heartbeat_timeout_ms: _Optional[int] = ...) -> None: ... + class Heartbeat(_message.Message): - __slots__ = () - def __init__(self) -> None: ... + __slots__ = ("sequence", "acknowledgement") + SEQUENCE_FIELD_NUMBER: _ClassVar[int] + ACKNOWLEDGEMENT_FIELD_NUMBER: _ClassVar[int] + sequence: int + acknowledgement: int + def __init__(self, sequence: _Optional[int] = ..., acknowledgement: _Optional[int] = ...) -> None: ... class ServerMessage(_message.Message): __slots__ = ("message",) diff --git a/src/rstream/client.py b/src/rstream/client.py index 2339da3..9346e4a 100644 --- a/src/rstream/client.py +++ b/src/rstream/client.py @@ -44,6 +44,28 @@ from rstream.stream import RstreamStream _T = TypeVar("_T") +_MAX_HEARTBEAT_TIMEOUT_MS = 900_000 + + +def _negotiated_heartbeat_timeout( + heartbeat: bool, + heartbeat_interval_ms: int, + response: pb.OpenControlChannelRsp.Ok, +) -> float: + if not response.HasField("liveness"): + return 0.0 + liveness = response.liveness + if ( + not heartbeat + or liveness.heartbeat_interval_ms != heartbeat_interval_ms + or liveness.heartbeat_timeout_ms < liveness.heartbeat_interval_ms + or liveness.heartbeat_timeout_ms > _MAX_HEARTBEAT_TIMEOUT_MS + ): + raise ProtocolError( + "Engine returned an invalid liveness policy.", + code="ERR_RSTREAM_PROTOCOL", + ) + return liveness.heartbeat_timeout_ms / 1_000 class Client: @@ -146,8 +168,17 @@ async def connect(self) -> ControlChannel: engine = await self._resolve_engine(resolved) token = await self._resolve_token(resolved, engine) reader, writer = await self._dial_engine(engine, resolved) + heartbeat_interval_ms = ( + round(resolved.heartbeat_interval * 1_000) if resolved.heartbeat else 0 + ) try: - await write_message(writer, message_with_open_control_channel_req(token)) + await write_message( + writer, + message_with_open_control_channel_req( + token, + heartbeat_interval_ms if resolved.heartbeat else None, + ), + ) response = await _wait_for_operation( read_message(reader), resolved.operation_timeout, @@ -172,6 +203,11 @@ async def connect(self) -> ControlChannel: writer, heartbeat=resolved.heartbeat, heartbeat_interval=resolved.heartbeat_interval, + heartbeat_timeout=_negotiated_heartbeat_timeout( + resolved.heartbeat, + heartbeat_interval_ms, + payload.ok, + ), operation_timeout=resolved.operation_timeout, open_proxy_connection=lambda request: self._open_proxy_connection( engine, diff --git a/src/rstream/config.py b/src/rstream/config.py index 7b6bd04..da4772e 100644 --- a/src/rstream/config.py +++ b/src/rstream/config.py @@ -4,6 +4,7 @@ import base64 import json +import math import os import re import ssl @@ -17,6 +18,8 @@ from rstream.errors import ConfigurationError, UnsupportedFeatureError DEFAULT_API_URL = "https://rstream.io" +MIN_HEARTBEAT_INTERVAL_MS = 1_000 +MAX_HEARTBEAT_INTERVAL_MS = 300_000 @dataclass(frozen=True) @@ -229,6 +232,30 @@ async def resolve_client_options(options: ClientOptions) -> ResolvedClientOption "operation_timeout must be positive.", code="ERR_RSTREAM_INVALID_TIMEOUT", ) + if options.heartbeat and not math.isfinite(options.heartbeat_interval): + raise ConfigurationError( + "heartbeat_interval must be between 1 and 300 seconds with " + "millisecond precision.", + code="ERR_RSTREAM_INVALID_CONFIG", + ) + heartbeat_interval_ms = ( + round(options.heartbeat_interval * 1_000) if options.heartbeat else 0 + ) + if options.heartbeat and ( + not math.isclose( + options.heartbeat_interval * 1_000, + heartbeat_interval_ms, + abs_tol=1e-9, + ) + or not MIN_HEARTBEAT_INTERVAL_MS + <= heartbeat_interval_ms + <= MAX_HEARTBEAT_INTERVAL_MS + ): + raise ConfigurationError( + "heartbeat_interval must be between 1 and 300 seconds with " + "millisecond precision.", + code="ERR_RSTREAM_INVALID_CONFIG", + ) region = _normalize_region( _first_defined(options.region, env.region, config.region) ) diff --git a/src/rstream/control.py b/src/rstream/control.py index ce16792..4f69d4e 100644 --- a/src/rstream/control.py +++ b/src/rstream/control.py @@ -4,6 +4,7 @@ import asyncio import uuid +from collections import deque from collections.abc import Awaitable, Callable, Mapping from contextlib import suppress from typing import TypeVar @@ -37,6 +38,11 @@ OpenProxyConnection = Callable[[pb.ProxyConnReq], Awaitable[RstreamStream]] _T = TypeVar("_T") +# Each active request owns a socket and a TLS handshake. Keep the worker budget +# aligned with the Go SDK and queue excess bursts instead of multiplying memory, +# file descriptors, and handshakes without a bound. +_MAX_ACTIVE_PROXY_CONNECTIONS = 64 +_MAX_QUEUED_PROXY_CONNECTIONS = 256 class ControlChannel: @@ -49,6 +55,7 @@ def __init__( *, heartbeat: bool, heartbeat_interval: float, + heartbeat_timeout: float, operation_timeout: float, open_proxy_connection: OpenProxyConnection, server_details: ServerDetails | None = None, @@ -57,6 +64,7 @@ def __init__( self._writer = writer self._heartbeat = heartbeat self._heartbeat_interval = heartbeat_interval + self._heartbeat_timeout = heartbeat_timeout self._operation_timeout = operation_timeout self._open_proxy_connection = open_proxy_connection self._server_details = server_details @@ -64,6 +72,11 @@ def __init__( self._pending_closes: dict[str, asyncio.Future[None]] = {} self._tunnels: dict[str, BytestreamTunnel] = {} self._write_lock = asyncio.Lock() + self._proxy_queue: deque[pb.ProxyConnReq] = deque() + self._proxy_tasks: set[asyncio.Task[None]] = set() + self._heartbeat_sequence = 0 + self._heartbeat_acknowledgement = 0 + self._liveness_handle: asyncio.TimerHandle | None = None self._closed = False self._closing = False self._close_error: BaseException | None = None @@ -73,6 +86,8 @@ def __init__( self._heartbeat_task: asyncio.Task[None] | None = None if heartbeat and heartbeat_interval > 0: self._heartbeat_task = asyncio.create_task(self._heartbeat_loop()) + if heartbeat_timeout > 0: + self._arm_liveness_timer() @property def closed(self) -> bool: @@ -231,9 +246,15 @@ async def __aexit__( async def _heartbeat_loop(self) -> None: try: while not self._closed: + self._heartbeat_sequence += 1 + await self._write( + message_with_heartbeat( + self._heartbeat_sequence + if self._heartbeat_timeout > 0 + else None + ) + ) await asyncio.sleep(self._heartbeat_interval) - if not self._closed: - await self._write(message_with_heartbeat()) except BaseException as error: if not self._closed: self._fail(error) @@ -242,6 +263,8 @@ async def _read_loop(self) -> None: try: while not self._closed: await self._handle_message(await read_message(self._reader)) + if not self._closed and self._heartbeat_timeout > 0: + self._arm_liveness_timer() except BaseException as error: if not self._closed: self._fail(error) @@ -255,7 +278,10 @@ async def _handle_message(self, message: pb.Message) -> None: self._handle_close_tunnel_rsp(message.close_tunnel_rsp.tunnel_id) return if payload == "proxy_conn_req": - await self._handle_proxy_conn_req(message.proxy_conn_req) + self._dispatch_proxy_conn_req(message.proxy_conn_req) + return + if payload == "heartbeat": + self._handle_heartbeat(message.heartbeat) return if payload == "close_control_channel_rsp": self._finish() @@ -290,6 +316,51 @@ def _handle_close_tunnel_rsp(self, tunnel_id: str) -> None: if pending is not None and not pending.done(): pending.set_result(None) + def _handle_heartbeat(self, heartbeat: pb.Heartbeat) -> None: + if not self._heartbeat: + raise self._heartbeat_protocol_error() + if self._heartbeat_timeout == 0: + if heartbeat.sequence != 0 or heartbeat.acknowledgement != 0: + raise self._heartbeat_protocol_error() + return + if ( + heartbeat.sequence != 0 + or heartbeat.acknowledgement == 0 + or heartbeat.acknowledgement <= self._heartbeat_acknowledgement + or heartbeat.acknowledgement > self._heartbeat_sequence + ): + raise self._heartbeat_protocol_error() + self._heartbeat_acknowledgement = heartbeat.acknowledgement + + def _dispatch_proxy_conn_req(self, request: pb.ProxyConnReq) -> None: + if len(self._proxy_tasks) >= _MAX_ACTIVE_PROXY_CONNECTIONS: + if len(self._proxy_queue) >= _MAX_QUEUED_PROXY_CONNECTIONS: + self._fail( + RuntimeError( + "Control channel proxy queue is full.", + code="ERR_RSTREAM_CONTROL_OVERLOAD", + ) + ) + return + queued = pb.ProxyConnReq() + queued.CopyFrom(request) + self._proxy_queue.append(queued) + return + task = asyncio.create_task(self._handle_proxy_conn_req(request)) + self._proxy_tasks.add(task) + task.add_done_callback(self._proxy_task_done) + + def _proxy_task_done(self, task: asyncio.Task[None]) -> None: + self._proxy_tasks.discard(task) + if not task.cancelled(): + try: + task.result() + except BaseException as error: + if not self._closed: + self._fail(error) + if not self._closed and self._proxy_queue: + self._dispatch_proxy_conn_req(self._proxy_queue.popleft()) + async def _handle_proxy_conn_req(self, request: pb.ProxyConnReq) -> None: tunnel = self._tunnels.get(request.tunnel_id) if tunnel is None: @@ -302,7 +373,13 @@ async def _handle_proxy_conn_req(self, request: pb.ProxyConnReq) -> None: return try: stream = await self._open_proxy_connection(request) + if self._closed: + stream.close() + await stream.wait_closed() + return if not tunnel.deliver(stream): + stream.close() + await stream.wait_closed() await self._write( message_with_proxy_conn_rsp( request.stream_id, @@ -311,15 +388,49 @@ async def _handle_proxy_conn_req(self, request: pb.ProxyConnReq) -> None: ) return await self._write(message_with_proxy_conn_rsp(request.stream_id)) - except BaseException as error: + except asyncio.CancelledError: + raise + except Exception as error: + if self._closed: + return await self._write( message_with_proxy_conn_rsp(request.stream_id, error_to_pb(str(error))) ) async def _write(self, message: pb.Message) -> None: + if self._closed: + raise RuntimeError( + "Control channel is closed.", + code="ERR_RSTREAM_CONTROL_CLOSED", + ) async with self._write_lock: await write_message(self._writer, message) + def _arm_liveness_timer(self) -> None: + if self._closed: + return + if self._liveness_handle is not None: + self._liveness_handle.cancel() + self._liveness_handle = asyncio.get_running_loop().call_later( + self._heartbeat_timeout, + self._liveness_expired, + ) + + def _liveness_expired(self) -> None: + self._fail( + RuntimeError( + "Control channel liveness timeout expired.", + code="ERR_RSTREAM_CONTROL_LIVENESS", + ) + ) + + @staticmethod + def _heartbeat_protocol_error() -> ProtocolError: + return ProtocolError( + "Engine returned an invalid heartbeat.", + code="ERR_RSTREAM_PROTOCOL", + ) + async def _await_pending( self, future: asyncio.Future[_T], @@ -336,14 +447,21 @@ async def _await_pending( def _fail(self, error: BaseException) -> None: self._close_error = error - self._finish() + self._finish(preserve_forwarders=_preserve_payloads_after(error)) - def _finish(self) -> None: + def _finish(self, *, preserve_forwarders: bool = False) -> None: if self._closed: return self._closed = True + if self._liveness_handle is not None: + self._liveness_handle.cancel() if self._heartbeat_task is not None: self._heartbeat_task.cancel() + current_task = asyncio.current_task() + for proxy_task in self._proxy_tasks: + if proxy_task is not current_task: + proxy_task.cancel() + self._proxy_queue.clear() if self._read_task is not asyncio.current_task(): self._read_task.cancel() self._writer.close() @@ -358,7 +476,7 @@ def _finish(self) -> None: if not close_pending.done(): close_pending.set_exception(error) for tunnel in self._tunnels.values(): - tunnel.on_close(error) + tunnel.on_close(error, preserve_forwarders=preserve_forwarders) self._pending_tunnels.clear() self._pending_closes.clear() self._tunnels.clear() @@ -369,6 +487,12 @@ def _finish(self) -> None: self._done.set_exception(self._close_error) +def _preserve_payloads_after(error: BaseException) -> bool: + if isinstance(error, RuntimeError): + return error.code == "ERR_RSTREAM_CONTROL_LIVENESS" + return isinstance(error, (asyncio.IncompleteReadError, ConnectionError, OSError)) + + def _normalize_bytestream_options(options: CreateTunnelOptions) -> TunnelProperties: if options.port is not None and options.protocol != "tcp": raise RuntimeError( diff --git a/src/rstream/protocol.py b/src/rstream/protocol.py index ab6881a..b05c00c 100644 --- a/src/rstream/protocol.py +++ b/src/rstream/protocol.py @@ -243,11 +243,18 @@ def create_client_details(token: str | None) -> pb.ClientDetails: return details -def message_with_open_control_channel_req(token: str | None) -> pb.Message: +def message_with_open_control_channel_req( + token: str | None, + heartbeat_interval_ms: int | None = None, +) -> pb.Message: message = pb.Message() message.open_control_channel_req.client_details.CopyFrom( create_client_details(token) ) + if heartbeat_interval_ms is not None: + message.open_control_channel_req.liveness.heartbeat_interval_ms = ( + heartbeat_interval_ms + ) return message @@ -310,9 +317,12 @@ def message_with_stream_req( return message -def message_with_heartbeat() -> pb.Message: +def message_with_heartbeat(sequence: int | None = None) -> pb.Message: message = pb.Message() - message.heartbeat.CopyFrom(pb.Heartbeat()) + if sequence is None: + message.heartbeat.CopyFrom(pb.Heartbeat()) + else: + message.heartbeat.sequence = sequence return message diff --git a/src/rstream/tunnel.py b/src/rstream/tunnel.py index 81c3c02..24be173 100644 --- a/src/rstream/tunnel.py +++ b/src/rstream/tunnel.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from collections import deque from contextlib import suppress from typing import Protocol @@ -13,6 +14,52 @@ DEFAULT_PUBLISHED_PORT = 443 +class _StreamQueue: + def __init__(self) -> None: + self._streams: deque[RstreamStream] = deque() + self._waiters: deque[asyncio.Future[RstreamStream]] = deque() + self._close_error: BaseException | None = None + + async def get(self) -> RstreamStream: + if self._streams: + return self._streams.popleft() + if self._close_error is not None: + raise self._close_error + future: asyncio.Future[RstreamStream] = ( + asyncio.get_running_loop().create_future() + ) + self._waiters.append(future) + try: + return await future + finally: + with suppress(ValueError): + self._waiters.remove(future) + + def put(self, stream: RstreamStream) -> bool: + if self._close_error is not None: + stream.close() + return False + while self._waiters: + waiter = self._waiters.popleft() + if waiter.done(): + continue + waiter.set_result(stream) + return True + self._streams.append(stream) + return True + + def close(self, error: BaseException) -> None: + if self._close_error is not None: + return + self._close_error = error + while self._streams: + self._streams.popleft().close() + while self._waiters: + waiter = self._waiters.popleft() + if not waiter.done(): + waiter.set_exception(error) + + class _TunnelControl(Protocol): async def close_tunnel(self, tunnel_id: str) -> None: ... @@ -32,11 +79,10 @@ def __init__( ) self._control = control self._properties = properties - self._queue: asyncio.Queue[RstreamStream | BaseException | None] = ( - asyncio.Queue() - ) + self._queue = _StreamQueue() self._closed = False - self._forward_tasks: set[asyncio.Task[None]] = set() + self._hard_closed = False + self._forward_tasks: dict[asyncio.Task[None], RstreamStream] = {} @property def id(self) -> str: @@ -55,31 +101,29 @@ def properties(self) -> TunnelProperties: return self._properties async def accept(self) -> RstreamStream: - item = await self._queue.get() - if item is None: + stream = await self._queue.get() + if self._hard_closed: + stream.close() raise RuntimeError( "Tunnel closed.", code="ERR_RSTREAM_TUNNEL_CLOSED", ) - if isinstance(item, BaseException): - raise item - return item + return stream async def close(self) -> None: if self._closed: + self.on_close() await self.wait_forwarders_closed() return - try: - await self._control.close_tunnel(self.id) - finally: - await self.wait_forwarders_closed() + await self._control.close_tunnel(self.id) + await self.wait_forwarders_closed() async def forward_to(self, host: str, port: int) -> None: _validate_local_endpoint(host, port) async for stream in self: task = asyncio.create_task(pipe_stream_to_local(stream, host, port)) - self._forward_tasks.add(task) - task.add_done_callback(self._forward_tasks.discard) + self._forward_tasks[task] = stream + task.add_done_callback(self._forward_task_done) def __aiter__(self) -> BytestreamTunnel: return self @@ -96,18 +140,29 @@ def deliver(self, stream: RstreamStream) -> bool: if self._closed: stream.close() return False - self._queue.put_nowait(stream) - return True + return self._queue.put(stream) - def on_close(self, error: BaseException | None = None) -> None: - if self._closed: - return - self._closed = True - if error is not None: - self._queue.put_nowait(error) - self._queue.put_nowait(None) - for task in self._forward_tasks: - task.cancel() + def on_close( + self, + error: BaseException | None = None, + *, + preserve_forwarders: bool = False, + ) -> None: + if not preserve_forwarders: + self._hard_closed = True + if not self._closed: + self._closed = True + self._queue.close( + error + or RuntimeError( + "Tunnel closed.", + code="ERR_RSTREAM_TUNNEL_CLOSED", + ) + ) + if not preserve_forwarders: + for task, stream in tuple(self._forward_tasks.items()): + stream.close() + task.cancel() async def wait_forwarders_closed(self) -> None: if not self._forward_tasks: @@ -115,6 +170,11 @@ async def wait_forwarders_closed(self) -> None: with suppress(asyncio.CancelledError): await asyncio.gather(*self._forward_tasks, return_exceptions=True) + def _forward_task_done(self, task: asyncio.Task[None]) -> None: + self._forward_tasks.pop(task, None) + if not task.cancelled(): + task.exception() + def format_forwarding_address(properties: TunnelProperties) -> str: published = _published_host(properties) diff --git a/tests/e2e/test_real_engine.py b/tests/e2e/test_real_engine.py index 11e2ad1..71783a0 100644 --- a/tests/e2e/test_real_engine.py +++ b/tests/e2e/test_real_engine.py @@ -70,6 +70,22 @@ async def test_real_engine_private_bytestream_matrix() -> None: finally: responder.cancel() await asyncio.gather(responder, return_exceptions=True) + + await asyncio.sleep(2.2) + responder = asyncio.create_task( + echo_one_stream(tunnel, b"heartbeat", b"HEARTBEAT") + ) + try: + await assert_private_round_trip( + client, + tunnel_name, + zero_rtt=False, + request=b"heartbeat", + response=b"HEARTBEAT", + ) + finally: + responder.cancel() + await asyncio.gather(responder, return_exceptions=True) finally: await tunnel.close() @@ -372,6 +388,7 @@ def real_engine_client(*, zero_rtt: bool) -> rstream.Client: ) return rstream.Client( engine=engine, + heartbeat_interval=1.0, no_token=token is None, read_config_file=False, tls=tls, diff --git a/tests/integration/test_runtime_fake_engine.py b/tests/integration/test_runtime_fake_engine.py index 9bf521e..f16f1cc 100644 --- a/tests/integration/test_runtime_fake_engine.py +++ b/tests/integration/test_runtime_fake_engine.py @@ -2,6 +2,7 @@ import asyncio import ssl +from contextlib import suppress from pathlib import Path import pytest @@ -13,6 +14,27 @@ from rstream.protocol import read_message, write_message +async def start_uppercase_server() -> tuple[asyncio.AbstractServer, str, int]: + async def echo( + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + try: + while data := await reader.read(64 * 1024): + writer.write(data.upper()) + await writer.drain() + finally: + writer.close() + await writer.wait_closed() + + server = await asyncio.start_server(echo, "127.0.0.1", 0) + sockets = server.sockets + assert sockets + address = sockets[0].getsockname() + assert isinstance(address, tuple) + return server, str(address[0]), int(address[1]) + + @pytest.mark.asyncio async def test_connect_create_and_close_tunnel(tmp_path: Path) -> None: async with await FakeEngine.start(tmp_path) as engine: @@ -32,6 +54,168 @@ async def test_connect_create_and_close_tunnel(tmp_path: Path) -> None: assert engine.close_tunnel_requests == 1 +@pytest.mark.asyncio +async def test_negotiated_liveness_tolerates_delayed_acknowledgement( + tmp_path: Path, +) -> None: + async with await FakeEngine.start( + tmp_path, + liveness=(1_000, 1_000), + acknowledge_heartbeats=True, + heartbeat_acknowledgement_delay=0.8, + ) as engine: + control = await client_for(engine, heartbeat_interval=1.0).connect() + + await asyncio.wait_for(engine.first_heartbeat.wait(), timeout=0.5) + assert engine.open_control_request is not None + assert engine.open_control_request.liveness.heartbeat_interval_ms == 1_000 + assert engine.heartbeats[0].sequence == 1 + await asyncio.sleep(1.1) + assert not control.closed + + await control.close() + + +@pytest.mark.asyncio +async def test_negotiated_liveness_tolerates_a_dropped_heartbeat( + tmp_path: Path, +) -> None: + async with await FakeEngine.start( + tmp_path, + liveness=(1_000, 2_500), + acknowledge_heartbeats=True, + heartbeat_acknowledgement_every=2, + ) as engine: + control = await client_for(engine, heartbeat_interval=1.0).connect() + + for sequence in range(1, 5): + heartbeat = await asyncio.wait_for( + engine.observed_heartbeats.get(), timeout=2.0 + ) + assert heartbeat.sequence == sequence + assert not control.closed + + await control.close() + + +@pytest.mark.asyncio +async def test_liveness_is_not_starved_by_a_stalled_proxy_tls_handshake( + tmp_path: Path, +) -> None: + accepted = asyncio.Event() + release = asyncio.Event() + blackhole_writers: set[asyncio.StreamWriter] = set() + + async def blackhole( + _reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + blackhole_writers.add(writer) + accepted.set() + await release.wait() + writer.close() + await writer.wait_closed() + + blackhole_server = await asyncio.start_server(blackhole, "127.0.0.1", 0) + socket = blackhole_server.sockets[0] + host, port = socket.getsockname()[:2] + try: + async with await FakeEngine.start( + tmp_path, + liveness=(1_000, 1_500), + acknowledge_heartbeats=True, + ) as engine: + control = await client_for(engine, heartbeat_interval=1.0).connect() + tunnel = await control.create_tunnel() + proxy_response = asyncio.create_task( + engine.request_proxy_connection( + tunnel.id, + "blocked-stream", + proxy_endpoint=f"{host}:{port}", + ) + ) + + await asyncio.wait_for(accepted.wait(), timeout=0.5) + await asyncio.sleep(1.9) + assert not control.closed + + await control.close() + proxy_response.cancel() + await asyncio.gather(proxy_response, return_exceptions=True) + finally: + release.set() + for writer in blackhole_writers: + writer.close() + await asyncio.gather( + *(writer.wait_closed() for writer in blackhole_writers), + return_exceptions=True, + ) + blackhole_server.close() + await blackhole_server.wait_closed() + + +@pytest.mark.asyncio +async def test_negotiated_liveness_expires_when_acknowledgements_stop( + tmp_path: Path, +) -> None: + async with await FakeEngine.start( + tmp_path, + liveness=(1_000, 1_000), + ) as engine: + control = await client_for(engine, heartbeat_interval=1.0).connect() + + with pytest.raises(rstream.RstreamRuntimeError) as exc: + await asyncio.wait_for(control.done(), timeout=1.5) + + assert exc.value.code == "ERR_RSTREAM_CONTROL_LIVENESS" + + +@pytest.mark.asyncio +async def test_negotiated_liveness_rejects_future_acknowledgement( + tmp_path: Path, +) -> None: + async with await FakeEngine.start( + tmp_path, + liveness=(1_000, 60_000), + acknowledge_heartbeats=True, + heartbeat_acknowledgement_offset=1, + ) as engine: + control = await client_for(engine, heartbeat_interval=1.0).connect() + + with pytest.raises(rstream.ProtocolError): + await asyncio.wait_for(control.done(), timeout=0.5) + + +@pytest.mark.asyncio +async def test_negotiated_liveness_rejects_replayed_acknowledgement( + tmp_path: Path, +) -> None: + async with await FakeEngine.start( + tmp_path, + liveness=(1_000, 60_000), + acknowledge_heartbeats=True, + duplicate_heartbeat_acknowledgement=True, + ) as engine: + control = await client_for(engine, heartbeat_interval=1.0).connect() + + with pytest.raises(rstream.ProtocolError): + await asyncio.wait_for(control.done(), timeout=0.5) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "liveness", + [(2_000, 60_000), (1_000, 999), (1_000, 900_001)], +) +async def test_connect_rejects_invalid_server_liveness_policy( + tmp_path: Path, + liveness: tuple[int, int], +) -> None: + async with await FakeEngine.start(tmp_path, liveness=liveness) as engine: + with pytest.raises(rstream.ProtocolError): + await client_for(engine, heartbeat_interval=1.0).connect() + + @pytest.mark.asyncio async def test_published_tcp_options_and_local_validation(tmp_path: Path) -> None: async with await FakeEngine.start(tmp_path) as engine: @@ -222,6 +406,268 @@ async def test_proxy_connection_delivery_round_trip(tmp_path: Path) -> None: ) +@pytest.mark.asyncio +async def test_proxy_handshakes_progress_concurrently_without_blocking_control( + tmp_path: Path, +) -> None: + stream_ids = [f"stream_concurrent_{index}" for index in range(16)] + async with await FakeEngine.start(tmp_path) as engine: + engine.proxy_handshake_gate = asyncio.Event() + client = client_for(engine, zero_rtt=False) + + async with await client.connect() as control: + tunnel = await control.create_tunnel() + response_tasks = [ + asyncio.create_task( + engine.request_proxy_connection( + tunnel.id, + stream_id, + timeout=3, + ) + ) + for stream_id in stream_ids + ] + observed = await asyncio.wait_for( + asyncio.gather( + *(engine.observed_proxy_requests.get() for _ in stream_ids) + ), + timeout=2, + ) + + assert {stream_id for stream_id, _ in observed} == set(stream_ids) + assert all(not zero_rtt for _, zero_rtt in observed) + + engine.proxy_handshake_gate.set() + streams = await asyncio.gather( + *(engine.next_proxy_stream() for _ in stream_ids) + ) + responses = await asyncio.wait_for( + asyncio.gather(*response_tasks), + timeout=2, + ) + + assert {response.proxy_conn_rsp.stream_id for response in responses} == set( + stream_ids + ) + assert all( + not response.proxy_conn_rsp.HasField("error") for response in responses + ) + for stream in streams: + stream.close() + await asyncio.gather( + *(stream.wait_closed() for stream in streams), + return_exceptions=True, + ) + + +@pytest.mark.asyncio +async def test_forwarded_stream_survives_liveness_timeout(tmp_path: Path) -> None: + server, host, port = await start_uppercase_server() + async with ( + server, + await FakeEngine.start(tmp_path, liveness=(1_000, 1_000)) as engine, + ): + control = await client_for(engine, heartbeat_interval=1.0).connect() + tunnel = await control.create_tunnel() + forwarding = asyncio.create_task(tunnel.forward_to(host, port)) + response_waiter = asyncio.create_task( + engine.request_proxy_connection(tunnel.id, "draining-forward") + ) + engine_stream = await engine.next_proxy_stream() + response = await response_waiter + assert not response.proxy_conn_rsp.HasField("error") + engine_stream.write(b"before") + await engine_stream.drain() + assert await engine_stream.readexactly(6) == b"BEFORE" + with pytest.raises(rstream.RstreamRuntimeError) as failure: + await asyncio.wait_for(control.done(), timeout=1.5) + assert failure.value.code == "ERR_RSTREAM_CONTROL_LIVENESS" + engine_stream.write(b"after") + await engine_stream.drain() + assert ( + await asyncio.wait_for( + engine_stream.readexactly(5), + timeout=0.5, + ) + == b"AFTER" + ) + engine_stream.close() + await engine_stream.wait_closed() + await tunnel.close() + await asyncio.gather(forwarding, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_accepted_stream_survives_liveness_timeout(tmp_path: Path) -> None: + async with await FakeEngine.start(tmp_path, liveness=(1_000, 1_000)) as engine: + control = await client_for(engine, heartbeat_interval=1.0).connect() + tunnel = await control.create_tunnel() + response_waiter = asyncio.create_task( + engine.request_proxy_connection(tunnel.id, "draining-accepted") + ) + app_stream = await tunnel.accept() + engine_stream = await engine.next_proxy_stream() + response = await response_waiter + assert not response.proxy_conn_rsp.HasField("error") + engine_stream.write(b"before") + await engine_stream.drain() + assert await app_stream.readexactly(6) == b"before" + with pytest.raises(rstream.RstreamRuntimeError) as failure: + await asyncio.wait_for(control.done(), timeout=1.5) + assert failure.value.code == "ERR_RSTREAM_CONTROL_LIVENESS" + with pytest.raises(rstream.RstreamRuntimeError): + await tunnel.accept() + engine_stream.write(b"after") + await engine_stream.drain() + assert await app_stream.readexactly(5) == b"after" + app_stream.write(b"return") + await app_stream.drain() + assert await engine_stream.readexactly(6) == b"return" + app_stream.close() + engine_stream.close() + await asyncio.gather( + app_stream.wait_closed(), + engine_stream.wait_closed(), + return_exceptions=True, + ) + + +@pytest.mark.asyncio +async def test_forwarded_stream_survives_control_transport_eof(tmp_path: Path) -> None: + server, host, port = await start_uppercase_server() + async with server, await FakeEngine.start(tmp_path) as engine: + control = await client_for(engine).connect() + tunnel = await control.create_tunnel() + forwarding = asyncio.create_task(tunnel.forward_to(host, port)) + response_waiter = asyncio.create_task( + engine.request_proxy_connection(tunnel.id, "eof-forward") + ) + engine_stream = await engine.next_proxy_stream() + response = await response_waiter + assert not response.proxy_conn_rsp.HasField("error") + await engine.close_control_transport() + with pytest.raises(asyncio.IncompleteReadError): + await asyncio.wait_for(control.done(), timeout=1) + engine_stream.write(b"survives") + await engine_stream.drain() + assert ( + await asyncio.wait_for( + engine_stream.readexactly(8), + timeout=0.5, + ) + == b"SURVIVES" + ) + engine_stream.close() + await engine_stream.wait_closed() + await tunnel.close() + await asyncio.gather(forwarding, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_explicit_control_close_stops_forwarded_stream(tmp_path: Path) -> None: + server, host, port = await start_uppercase_server() + async with server, await FakeEngine.start(tmp_path) as engine: + control = await client_for(engine).connect() + tunnel = await control.create_tunnel() + forwarding = asyncio.create_task(tunnel.forward_to(host, port)) + response_waiter = asyncio.create_task( + engine.request_proxy_connection(tunnel.id, "hard-close-forward") + ) + engine_stream = await engine.next_proxy_stream() + response = await response_waiter + assert not response.proxy_conn_rsp.HasField("error") + await control.close() + assert await asyncio.wait_for(engine_stream.read(), timeout=0.5) == b"" + engine_stream.close() + await engine_stream.wait_closed() + await asyncio.gather(forwarding, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_protocol_failure_stops_forwarded_stream(tmp_path: Path) -> None: + server, host, port = await start_uppercase_server() + async with server, await FakeEngine.start(tmp_path) as engine: + control = await client_for(engine).connect() + tunnel = await control.create_tunnel() + forwarding = asyncio.create_task(tunnel.forward_to(host, port)) + response_waiter = asyncio.create_task( + engine.request_proxy_connection(tunnel.id, "protocol-close-forward") + ) + engine_stream = await engine.next_proxy_stream() + response = await response_waiter + assert not response.proxy_conn_rsp.HasField("error") + await engine.send_heartbeat(acknowledgement=1) + with pytest.raises(rstream.ProtocolError): + await asyncio.wait_for(control.done(), timeout=0.5) + assert await asyncio.wait_for(engine_stream.read(), timeout=0.5) == b"" + engine_stream.close() + await engine_stream.wait_closed() + await asyncio.gather(forwarding, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_control_close_closes_an_unaccepted_proxy_stream(tmp_path: Path) -> None: + async with await FakeEngine.start(tmp_path) as engine: + control = await client_for(engine).connect() + tunnel = await control.create_tunnel() + response_waiter = asyncio.create_task( + engine.request_proxy_connection(tunnel.id, "unaccepted-stream") + ) + engine_stream = await engine.next_proxy_stream() + response = await response_waiter + assert not response.proxy_conn_rsp.HasField("error") + + await control.close() + + assert await asyncio.wait_for(engine_stream.read(), timeout=0.5) == b"" + engine_stream.close() + await engine_stream.wait_closed() + + +@pytest.mark.asyncio +async def test_proxy_stream_is_closed_when_tunnel_closes_during_dial( + tmp_path: Path, +) -> None: + async with await FakeEngine.start(tmp_path) as engine: + engine.proxy_handshake_gate = asyncio.Event() + async with await client_for(engine).connect() as control: + tunnel = await control.create_tunnel() + response_waiter = asyncio.create_task( + engine.request_proxy_connection(tunnel.id, "closing-stream") + ) + await asyncio.wait_for(engine.first_proxy_request.wait(), timeout=1) + await tunnel.close() + engine.proxy_handshake_gate.set() + engine_stream = await engine.next_proxy_stream() + response = await response_waiter + + assert response.proxy_conn_rsp.HasField("error") + assert await asyncio.wait_for(engine_stream.read(), timeout=0.5) == b"" + engine_stream.close() + await engine_stream.wait_closed() + + +@pytest.mark.asyncio +async def test_control_close_releases_all_concurrent_accept_waiters( + tmp_path: Path, +) -> None: + async with await FakeEngine.start(tmp_path) as engine: + control = await client_for(engine).connect() + tunnel = await control.create_tunnel() + accept_tasks = [asyncio.create_task(tunnel.accept()) for _ in range(8)] + await asyncio.sleep(0) + + await control.close() + + results = await asyncio.wait_for( + asyncio.gather(*accept_tasks, return_exceptions=True), + timeout=0.5, + ) + assert all( + isinstance(result, rstream.RstreamRuntimeError) for result in results + ) + + @pytest.mark.asyncio async def test_proxy_connection_can_dial_ingress_engine(tmp_path: Path) -> None: ca = trustme.CA() @@ -353,12 +799,14 @@ async def test_client_close_closes_open_control_channels(tmp_path: Path) -> None def client_for( engine: FakeEngine, *, + heartbeat_interval: float = 5.0, operation_timeout: float = 1, token: str | None = None, zero_rtt: bool = False, ) -> rstream.Client: return rstream.Client( engine=engine.address, + heartbeat_interval=heartbeat_interval, no_token=token is None, operation_timeout=operation_timeout, read_config_file=False, @@ -392,7 +840,10 @@ def __init__( self.open_tunnel_requests = 0 self.stream_requests: list[tuple[str, bool]] = [] self.proxy_requests: list[tuple[str, bool]] = [] + self.observed_proxy_requests: asyncio.Queue[tuple[str, bool]] = asyncio.Queue() self.proxy_tokens: list[str | None] = [] + self.proxy_handshake_gate: asyncio.Event | None = None + self.first_proxy_request = asyncio.Event() self._control_writer: asyncio.StreamWriter | None = None self._pending_proxy_responses: asyncio.Queue[pb.Message] = asyncio.Queue() self._proxy_streams: asyncio.Queue[rstream.RstreamStream] = asyncio.Queue() @@ -400,6 +851,16 @@ def __init__( self._writers: set[asyncio.StreamWriter] = set() self._hold_tasks: set[asyncio.Task[None]] = set() self._closed_event = asyncio.Event() + self.liveness: tuple[int, int] | None = None + self.acknowledge_heartbeats = False + self.heartbeat_acknowledgement_delay = 0.0 + self.heartbeat_acknowledgement_offset = 0 + self.heartbeat_acknowledgement_every = 1 + self.duplicate_heartbeat_acknowledgement = False + self.open_control_request: pb.OpenControlChannelReq | None = None + self.heartbeats: list[pb.Heartbeat] = [] + self.observed_heartbeats: asyncio.Queue[pb.Heartbeat] = asyncio.Queue() + self.first_heartbeat = asyncio.Event() @classmethod async def start( @@ -408,6 +869,12 @@ async def start( *, name: str = "engine", ca: trustme.CA | None = None, + liveness: tuple[int, int] | None = None, + acknowledge_heartbeats: bool = False, + heartbeat_acknowledgement_delay: float = 0.0, + heartbeat_acknowledgement_offset: int = 0, + heartbeat_acknowledgement_every: int = 1, + duplicate_heartbeat_acknowledgement: bool = False, ) -> FakeEngine: certificate_authority = ca or trustme.CA() cert = certificate_authority.issue_cert("localhost", "127.0.0.1") @@ -430,6 +897,12 @@ async def start( socket = server.sockets[0] host, port = socket.getsockname()[:2] cls.__init__(engine, server, f"{host}:{port}", ca_file) + engine.liveness = liveness + engine.acknowledge_heartbeats = acknowledge_heartbeats + engine.heartbeat_acknowledgement_delay = heartbeat_acknowledgement_delay + engine.heartbeat_acknowledgement_offset = heartbeat_acknowledgement_offset + engine.heartbeat_acknowledgement_every = heartbeat_acknowledgement_every + engine.duplicate_heartbeat_acknowledgement = duplicate_heartbeat_acknowledgement return engine async def __aenter__(self) -> FakeEngine: @@ -465,6 +938,7 @@ async def request_proxy_connection( proxy_endpoint: str | None = None, include_secret: bool = True, secret: str = "stream-secret", + timeout: float = 1, ) -> pb.Message: writer = self._control_writer assert writer is not None @@ -478,7 +952,10 @@ async def request_proxy_connection( if include_secret: message.proxy_conn_req.secret.CopyFrom(StringValue(value=secret)) await write_message(writer, message) - return await asyncio.wait_for(self._pending_proxy_responses.get(), timeout=1) + return await asyncio.wait_for( + self._pending_proxy_responses.get(), + timeout=timeout, + ) async def next_proxy_stream(self) -> rstream.RstreamStream: return await asyncio.wait_for(self._proxy_streams.get(), timeout=1) @@ -487,7 +964,15 @@ async def close_control_transport(self) -> None: writer = self._control_writer assert writer is not None writer.close() - await writer.wait_closed() + with suppress(ssl.SSLError, ConnectionError): + await writer.wait_closed() + + async def send_heartbeat(self, *, acknowledgement: int) -> None: + writer = self._control_writer + assert writer is not None + message = pb.Message() + message.heartbeat.acknowledgement = acknowledgement + await write_message(writer, message) async def _handle_connection( self, @@ -499,6 +984,7 @@ async def _handle_connection( message = await read_message(reader) payload = message.WhichOneof("payload") if payload == "open_control_channel_req": + self.open_control_request = message.open_control_channel_req await self._handle_control_channel(reader, writer) return if payload == "stream_req": @@ -510,7 +996,8 @@ async def _handle_connection( finally: self._writers.discard(writer) writer.close() - await writer.wait_closed() + with suppress(ssl.SSLError, ConnectionError): + await writer.wait_closed() async def _handle_control_channel( self, @@ -525,6 +1012,14 @@ async def _handle_control_channel( response.open_control_channel_rsp.ok.server_details.agent.CopyFrom( StringValue(value="fake-engine") ) + if self.liveness is not None: + interval_ms, timeout_ms = self.liveness + response.open_control_channel_rsp.ok.liveness.heartbeat_interval_ms = ( + interval_ms + ) + response.open_control_channel_rsp.ok.liveness.heartbeat_timeout_ms = ( + timeout_ms + ) await write_message(writer, response) while True: message = await read_message(reader) @@ -542,6 +1037,23 @@ async def _handle_control_channel( await self._pending_proxy_responses.put(message) continue if payload == "heartbeat": + heartbeat = pb.Heartbeat() + heartbeat.CopyFrom(message.heartbeat) + self.heartbeats.append(heartbeat) + await self.observed_heartbeats.put(heartbeat) + self.first_heartbeat.set() + if self.acknowledge_heartbeats: + if len(self.heartbeats) % self.heartbeat_acknowledgement_every != 0: + continue + if len(self.heartbeats) == 1: + await asyncio.sleep(self.heartbeat_acknowledgement_delay) + response = pb.Message() + response.heartbeat.acknowledgement = ( + heartbeat.sequence + self.heartbeat_acknowledgement_offset + ) + await write_message(writer, response) + if self.duplicate_heartbeat_acknowledgement: + await write_message(writer, response) continue async def _handle_open_tunnel( @@ -634,6 +1146,8 @@ async def _handle_proxy( ) -> None: zero_rtt = request.HasField("zero_rtt") and request.zero_rtt.value self.proxy_requests.append((request.stream_id, zero_rtt)) + await self.observed_proxy_requests.put((request.stream_id, zero_rtt)) + self.first_proxy_request.set() self.proxy_tokens.append( request.client_details.token.value if request.client_details.HasField("token") @@ -641,6 +1155,8 @@ async def _handle_proxy( ) if self._consume("next_proxy_hang"): await self._hold() + if self.proxy_handshake_gate is not None: + await self.proxy_handshake_gate.wait() if not zero_rtt: response = pb.Message() response.proxy_rsp.CopyFrom(pb.ProxyRsp()) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 3d1e767..6f5120c 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -44,6 +44,22 @@ def test_default_runtime_timeouts_are_resolved() -> None: ClientOptions(operation_timeout=-1, read_config_file=False), "operation_timeout", ), + ( + ClientOptions(heartbeat_interval=0.999, read_config_file=False), + "heartbeat_interval", + ), + ( + ClientOptions(heartbeat_interval=300.001, read_config_file=False), + "heartbeat_interval", + ), + ( + ClientOptions(heartbeat_interval=1.0005, read_config_file=False), + "heartbeat_interval", + ), + ( + ClientOptions(heartbeat_interval=float("inf"), read_config_file=False), + "heartbeat_interval", + ), ], ) def test_resolve_client_options_rejects_invalid_timeouts( diff --git a/tests/unit/test_protocol.py b/tests/unit/test_protocol.py index 70db37c..ae3143a 100644 --- a/tests/unit/test_protocol.py +++ b/tests/unit/test_protocol.py @@ -10,6 +10,8 @@ create_client_details, decode_message, encode_message, + message_with_heartbeat, + message_with_open_control_channel_req, protocol_version, read_message, tunnel_properties_from_pb, @@ -89,6 +91,16 @@ def test_client_details_use_protocol_version_from_proto_descriptor() -> None: assert __version__ != "unknown" +def test_control_liveness_messages_preserve_sequence_and_interval() -> None: + request = message_with_open_control_channel_req(None, 1_250) + heartbeat = message_with_heartbeat(42) + + assert request.open_control_channel_req.liveness.heartbeat_interval_ms == 1_250 + assert request.open_control_channel_req.liveness.heartbeat_timeout_ms == 0 + assert heartbeat.heartbeat.sequence == 42 + assert heartbeat.heartbeat.acknowledgement == 0 + + @pytest.mark.asyncio async def test_read_message_rejects_oversized_frame() -> None: reader = asyncio.StreamReader() diff --git a/tests/unit/test_tunnel.py b/tests/unit/test_tunnel.py index 923523f..b8e9952 100644 --- a/tests/unit/test_tunnel.py +++ b/tests/unit/test_tunnel.py @@ -99,8 +99,8 @@ async def forwarder() -> None: stopped.set() task = asyncio.create_task(forwarder()) - tunnel._forward_tasks.add(task) - task.add_done_callback(tunnel._forward_tasks.discard) + tunnel._forward_tasks[task] = _stream_double() + task.add_done_callback(tunnel._forward_task_done) await started.wait() await tunnel.close() @@ -109,6 +109,41 @@ async def forwarder() -> None: assert task.done() +@pytest.mark.asyncio +async def test_hard_close_cancels_a_preserved_forwarder() -> None: + control = _Control() + tunnel = BytestreamTunnel(control, TunnelProperties(id="tun_123")) + started = asyncio.Event() + stopped = asyncio.Event() + + async def forwarder() -> None: + started.set() + try: + await asyncio.Future() + finally: + stopped.set() + + task = asyncio.create_task(forwarder()) + tunnel._forward_tasks[task] = _stream_double() + task.add_done_callback(tunnel._forward_task_done) + await started.wait() + tunnel.on_close( + RstreamRuntimeError( + "Control channel liveness timeout expired.", + code="ERR_RSTREAM_CONTROL_LIVENESS", + ), + preserve_forwarders=True, + ) + await asyncio.sleep(0) + assert not task.done() + + await tunnel.close() + + assert control.closed_tunnels == [] + assert stopped.is_set() + assert task.done() + + @pytest.mark.asyncio async def test_tunnel_accepts_delivered_stream() -> None: stream = _stream_double() @@ -120,6 +155,42 @@ async def test_tunnel_accepts_delivered_stream() -> None: assert accepted is stream +@pytest.mark.asyncio +async def test_hard_close_wins_delivery_acceptance_race() -> None: + stream = _stream_double() + tunnel = BytestreamTunnel(_Control(), TunnelProperties(id="tun_123")) + acceptance = asyncio.create_task(tunnel.accept()) + await asyncio.sleep(0) + + assert tunnel.deliver(stream) is True + tunnel.on_close() + + with pytest.raises(RstreamRuntimeError) as failure: + await acceptance + assert failure.value.code == "ERR_RSTREAM_TUNNEL_CLOSED" + assert stream.writer.is_closing() + + +@pytest.mark.asyncio +async def test_soft_close_preserves_delivery_acceptance_race() -> None: + stream = _stream_double() + tunnel = BytestreamTunnel(_Control(), TunnelProperties(id="tun_123")) + acceptance = asyncio.create_task(tunnel.accept()) + await asyncio.sleep(0) + + assert tunnel.deliver(stream) is True + tunnel.on_close( + RstreamRuntimeError( + "Control channel liveness timeout expired.", + code="ERR_RSTREAM_CONTROL_LIVENESS", + ), + preserve_forwarders=True, + ) + + assert await acceptance is stream + assert not stream.writer.is_closing() + + @pytest.mark.asyncio async def test_tunnel_rejects_delivery_after_close() -> None: stream = _stream_double()