From 14e4cd56d0f7ae6259c10282505044d48b1c2b49 Mon Sep 17 00:00:00 2001 From: toki Date: Sat, 11 Apr 2026 08:33:46 +0900 Subject: [PATCH] feat: add WebSocket/WSS support and Go implementation - Add WebSocket binary frame support to protocol specification - Update README and PROTOCOL.md to reflect dual TCP/WebSocket transport - Add Go implementation with TCP, WebSocket, and heartbeat support - Include .claude settings configuration --- .claude/settings.json | 10 + PROTOCOL.md | 30 +-- README.md | 92 +++++++- go/communicator.go | 358 +++++++++++++++++++++++++++++++ go/communicator_test.go | 99 +++++++++ go/example/tcp_echo/main.go | 60 ++++++ go/example/ws_echo/main.go | 60 ++++++ go/go.mod | 8 + go/go.sum | 8 + go/heartbeat_timer.go | 38 ++++ go/packets/message_common.pb.go | 244 +++++++++++++++++++++ go/packets/message_common.proto | 17 ++ go/socket_test.go | 368 ++++++++++++++++++++++++++++++++ go/tcp_client.go | 216 +++++++++++++++++++ go/tcp_server.go | 146 +++++++++++++ go/ws_client.go | 206 ++++++++++++++++++ go/ws_server.go | 156 ++++++++++++++ 17 files changed, 2100 insertions(+), 16 deletions(-) create mode 100644 .claude/settings.json create mode 100644 go/communicator.go create mode 100644 go/communicator_test.go create mode 100644 go/example/tcp_echo/main.go create mode 100644 go/example/ws_echo/main.go create mode 100644 go/go.mod create mode 100644 go/go.sum create mode 100644 go/heartbeat_timer.go create mode 100644 go/packets/message_common.pb.go create mode 100644 go/packets/message_common.proto create mode 100644 go/socket_test.go create mode 100644 go/tcp_client.go create mode 100644 go/tcp_server.go create mode 100644 go/ws_client.go create mode 100644 go/ws_server.go diff --git a/.claude/settings.json b/.claude/settings.json new file mode 100644 index 0000000..446f86f --- /dev/null +++ b/.claude/settings.json @@ -0,0 +1,10 @@ +{ + "permissions": { + "allow": [ + "Bash(apt-get update:*)", + "Bash(apt-get install:*)", + "Bash(sudo apt-get:*)", + "Bash(protoc --version)" + ] + } +} diff --git a/PROTOCOL.md b/PROTOCOL.md index 5af3627..8f5e674 100644 --- a/PROTOCOL.md +++ b/PROTOCOL.md @@ -1,6 +1,6 @@ # Toki Socket Protocol -Binary TCP protocol using Protocol Buffers for message framing and type-based routing. +Binary TCP and WebSocket protocol using Protocol Buffers for message framing and type-based routing. Designed for bidirectional, heterogeneous communication across languages and platforms. --- @@ -9,7 +9,7 @@ Designed for bidirectional, heterogeneous communication across languages and pla This protocol intentionally defines only the transport-level contract shared across implementations. -- In scope: framing, protobuf payload transport, `typeName` routing, `nonce`, `responseNonce`, and heartbeat +- In scope: TCP/WebSocket framing, protobuf payload transport, `typeName` routing, `nonce`, `responseNonce`, TLS/WSS transport variants, and heartbeat - Out of scope: authentication, session lifecycle, agent semantics, chat semantics, game rules, room logic, and other application-specific behavior - Higher-level concerns should be implemented as messages and conventions on top of this protocol, not inside the protocol core @@ -21,16 +21,17 @@ This protocol intentionally defines only the transport-level contract shared acr ``` ┌──────────────────────────────────────────┐ -│ Header (4 bytes, big-endian int32) │ ← PacketBase payload length +│ Header (4 bytes, big-endian integer) │ ← PacketBase payload length ├──────────────────────────────────────────┤ │ PacketBase (protobuf, N bytes) │ ← typeName + nonce + data └──────────────────────────────────────────┘ ``` #### Header -- 4 bytes, big-endian signed int32 +- 4 bytes, big-endian positive integer. Dart reads/writes it as `int32`; Go reads/writes the same 4 bytes as `uint32`. - Value: byte length of the following `PacketBase` protobuf payload - Value of `0`: reserved / no-op, receiver clears buffer +- Implementations may apply safety limits to payload size. The current Go implementation rejects TCP packets larger than 64 MiB. ### WebSocket / WSS @@ -47,7 +48,7 @@ This protocol intentionally defines only the transport-level contract shared acr ### PacketBase ```protobuf message PacketBase { - string typeName = 1; // fully-qualified protobuf message name + string typeName = 1; // protobuf message name used for routing int32 nonce = 2; // monotonically increasing per-connection counter bytes data = 3; // serialized inner message bytes int32 responseNonce = 4; // > 0: this packet is a response to request with that nonce @@ -56,7 +57,7 @@ message PacketBase { | Field | Description | |-------|-------------| -| `typeName` | `GeneratedMessage.info_.qualifiedMessageName` — language-agnostic routing key | +| `typeName` | Protobuf message name used as the language-agnostic routing key | | `nonce` | Starts at 1, increments by 1 per send call per connection | | `data` | Protobuf-serialized bytes of the inner message | | `responseNonce` | `0` for regular messages. Set to the request's `nonce` when sending a response | @@ -92,24 +93,27 @@ Side A Side B - After any received message → heartbeat interval timer resets - If no data received within `heartbeatIntervalTime` seconds → send `HeartBeat` - If no response within `heartbeatWaitTime` seconds → `onDisconnected()` → `dispose()` -- Only the initiating side responds to HeartBeat echo (ping-pong prevention via `_waitingHeartbeatResponse` flag) +- A side that is already waiting for a heartbeat response consumes the echo and does not send another heartbeat. A side that is not waiting replies with `HeartBeat`. --- ## typeName Convention -Uses `GeneratedMessage.info_.qualifiedMessageName` on the send side. +Uses the protobuf message name on the send side. On the receive side, use the same value as the registration key. | Language | Registration key | |----------|-----------------| | Dart | `T.toString()` (equals qualified name for top-level proto messages) | +| Go | `string(proto.MessageName(m))` via `TypeNameOf(m)` | | C# | `typeof(T).Name` — verify matches proto qualified name | | Kotlin | `T::class.simpleName` — verify matches | | Swift | `String(describing: T.self)` — verify matches | | Python | `descriptor.name` from `MessageClass.DESCRIPTOR` — verify matches | | Rust | `M::default().descriptor_dyn().name().to_string()` (protobuf crate) — verify matches | +The current packet proto has no `package` declaration, so Dart and Go both use simple names such as `TestData` and `HeartBeat`. If a future proto adds a `package`, `proto.MessageName` may become fully qualified, and all implementations must use the same value. + **Important**: Verify typeName consistency across languages before connecting heterogeneous clients. --- @@ -145,7 +149,7 @@ Requester Responder - `responseNonce > 0`: response → matched to the pending `sendRequest` by `responseNonce` - A response must also have the expected `typeName` for the waiting `sendRequest`; mismatches are protocol errors - Both sides can be requester and responder simultaneously on the same connection -- For a given `typeName` on one connection, register it with either `addListener` or `addRequestListener`, not both. Mixed registration is ambiguous and rejected by the Dart implementation. +- For a given `typeName` on one connection, register it with either `addListener` or `addRequestListener`, not both. Mixed registration is ambiguous and rejected by the Dart and Go implementations. --- @@ -153,7 +157,7 @@ Requester Responder Sending `HeartBeat {}`: ``` -00 00 00 05 ← header: PacketBase length = 5 bytes +00 00 00 0B ← header: PacketBase length = 11 bytes 0A 09 48 65 61 72 74 42 65 61 74 ← PacketBase { typeName: "HeartBeat" } ``` @@ -167,7 +171,7 @@ Sending `HeartBeat {}`: | C# (Unity) | Planned | `csharp/` | | Kotlin | Planned | `kotlin/` | | Swift | Planned | `swift/` | -| Go | Planned | `go/` | +| Go | Available | `go/` | | Python | Planned | `python/` | | Rust | Planned | `rust/` | @@ -176,4 +180,6 @@ Sending `HeartBeat {}`: ## Proto Source `dart/lib/src/packets/message_common.proto` is the canonical packet definition. -All language implementations must generate bindings from this file. +All language implementations must keep the same message schema and generate bindings from it. + +The Go implementation keeps a copy at `go/packets/message_common.proto` because Go generation needs `option go_package`. Keep the message fields in sync with the canonical Dart proto before regenerating `go/packets/message_common.pb.go`. diff --git a/README.md b/README.md index 75a02b6..3ccf8ca 100644 --- a/README.md +++ b/README.md @@ -1,15 +1,15 @@ # Toki Socket -Binary TCP socket protocol library for bidirectional, heterogeneous communication across languages and platforms. +Binary socket protocol library for bidirectional, heterogeneous communication across languages and platforms. -Built on Protocol Buffers with a fixed 4-byte length-prefixed framing, type-based message routing, and built-in heartbeat. +Built on Protocol Buffers with TCP length-prefixed framing, WebSocket binary frames, type-based message routing, request-response correlation, and built-in heartbeat. --- ## Design Principles - Keep the core transport layer thin and stable -- Provide only the minimum common foundation for cross-language communication5 +- Provide only the minimum common foundation for cross-language communication - Standardize framing, serialization, routing, request-response correlation, and heartbeat - Do not embed application semantics into the protocol core - Domain-specific concerns such as auth, session, agent workflow, chat features, and game logic belong in upper-layer implementations @@ -24,6 +24,8 @@ See [PROTOCOL.md](PROTOCOL.md) for the full wire format specification. [4-byte big-endian length] [PacketBase protobuf bytes] ``` +For WebSocket/WSS transports, each binary frame contains one `PacketBase` protobuf payload without the TCP length header. + --- ## Implementations @@ -34,7 +36,7 @@ See [PROTOCOL.md](PROTOCOL.md) for the full wire format specification. | C# | Planned | `csharp/` | Unity, .NET | | Kotlin | Planned | `kotlin/` | Android | | Swift | Planned | `swift/` | iOS, macOS | -| Go | Planned | `go/` | Server, tooling, scripting | +| Go | Available | [go/](go/) | Server, tooling, scripting | | Python | Planned | `python/` | Server, tooling, scripting | | Rust | Planned | `rust/` | High-performance server, embedded | @@ -86,6 +88,83 @@ dart pub global activate protoc_plugin protoc --dart_out=lib/src/packets lib/src/packets/message_common.proto ``` +For Go, keep `go/packets/message_common.proto` in sync with the canonical Dart proto. The Go copy includes `option go_package`, then regenerate from the Go module root: + +```bash +cd go +protoc --go_out=. --go_opt=paths=source_relative packets/message_common.proto +``` + +--- + +## Quick Start (Go) + +```go +package main + +import ( + "context" + "net" + "time" + + "google.golang.org/protobuf/proto" + + toki "toki-labs.com/toki_socket/go" + "toki-labs.com/toki_socket/go/packets" +) + +func parserMap() toki.ParserMap { + return toki.ParserMap{ + toki.TypeNameOf(&packets.TestData{}): func(b []byte) (proto.Message, error) { + m := &packets.TestData{} + return m, proto.Unmarshal(b, m) + }, + } +} + +func main() { + ctx := context.Background() + + server := toki.NewTcpServer("127.0.0.1", 9090, func(conn net.Conn) *toki.TcpClient { + return toki.NewTcpClient(conn, 30, 10, parserMap()) + }) + server.OnClientConnected = func(client *toki.TcpClient) { + toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData]( + &client.Communicator, + func(req *packets.TestData) (*packets.TestData, error) { + return &packets.TestData{Index: req.GetIndex(), Message: "echo: " + req.GetMessage()}, nil + }, + ) + } + if err := server.Start(ctx); err != nil { + panic(err) + } + defer server.Stop() + + client, err := toki.DialTcp(ctx, "127.0.0.1", 9090, 30, 10, parserMap()) + if err != nil { + panic(err) + } + defer client.Close() + + res, err := toki.SendRequestTyped[*packets.TestData, *packets.TestData]( + &client.Communicator, + &packets.TestData{Index: 1, Message: "hello"}, + 2*time.Second, + ) + if err != nil { + panic(err) + } + println(res.GetMessage()) +} +``` + +Go also provides: + +- TCP: `NewTcpServer`, `DialTcp`, `NewTcpServerTLS`, `DialTcpTLS` +- WebSocket: `NewWsServer`, `DialWs`, `NewWsServerTLS`, `DialWss` +- Shared helpers: `Send`, `SendRequest`, `AddListenerTyped`, `AddRequestListenerTyped`, `Broadcast` + --- ## Running Tests @@ -95,3 +174,8 @@ cd dart dart pub get dart test ``` + +```bash +cd go +go test ./... +``` diff --git a/go/communicator.go b/go/communicator.go new file mode 100644 index 0000000..e54faa7 --- /dev/null +++ b/go/communicator.go @@ -0,0 +1,358 @@ +package toki_socket + +import ( + "errors" + "fmt" + "reflect" + "sync" + "sync/atomic" + "time" + + "google.golang.org/protobuf/proto" + + "toki-labs.com/toki_socket/go/packets" +) + +var ErrNotConnected = errors.New("not connected") + +type ParserMap map[string]func([]byte) (proto.Message, error) + +type Transport interface { + WritePacket(base *packets.PacketBase) error + Close() error +} + +type pendingRequest struct { + expectedTypeName string + ch chan proto.Message + errCh chan error +} + +type queuedPacket struct { + base *packets.PacketBase + done chan error +} + +type Communicator struct { + mu sync.RWMutex + nonce atomic.Int32 + isAlive atomic.Bool + parserMap ParserMap + handlers map[string][]func(proto.Message) + reqHandlers map[string]func(proto.Message, int32) + pendingRequests map[int32]*pendingRequest + writeQueue chan queuedPacket + closed chan struct{} + closeOnce sync.Once + transport Transport + writeErrHandler func(error) +} + +func TypeNameOf(m proto.Message) string { + return string(proto.MessageName(m)) +} + +func NewCommunicator(transport Transport, parserMap ParserMap) *Communicator { + c := &Communicator{} + c.Initialize(transport, parserMap) + return c +} + +func (c *Communicator) Initialize(transport Transport, parserMap ParserMap) { + c.transport = transport + c.parserMap = make(ParserMap, len(parserMap)+1) + for k, v := range parserMap { + c.parserMap[k] = v + } + c.parserMap[TypeNameOf(&packets.HeartBeat{})] = func(b []byte) (proto.Message, error) { + m := &packets.HeartBeat{} + return m, proto.Unmarshal(b, m) + } + c.handlers = make(map[string][]func(proto.Message)) + c.reqHandlers = make(map[string]func(proto.Message, int32)) + c.pendingRequests = make(map[int32]*pendingRequest) + c.writeQueue = make(chan queuedPacket, 64) + c.closed = make(chan struct{}) + c.isAlive.Store(true) + go c.writeLoop() +} + +func (c *Communicator) IsAlive() bool { + return c.isAlive.Load() +} + +func (c *Communicator) SetWriteErrorHandler(fn func(error)) { + c.mu.Lock() + defer c.mu.Unlock() + c.writeErrHandler = fn +} + +func (c *Communicator) nextNonce() int32 { + return c.nonce.Add(1) +} + +func (c *Communicator) shutdown() { + c.isAlive.Store(false) + c.closeOnce.Do(func() { + close(c.closed) + }) +} + +func (c *Communicator) writeLoop() { + for { + select { + case item := <-c.writeQueue: + err := c.transport.WritePacket(item.base) + item.done <- err + if err != nil { + c.mu.RLock() + handler := c.writeErrHandler + c.mu.RUnlock() + if handler != nil { + handler(err) + } + } + case <-c.closed: + return + } + } +} + +func (c *Communicator) QueuePacket(base *packets.PacketBase) error { + if !c.IsAlive() { + return ErrNotConnected + } + done := make(chan error, 1) + select { + case c.writeQueue <- queuedPacket{base: base, done: done}: + case <-c.closed: + return ErrNotConnected + } + select { + case err := <-done: + return err + case <-c.closed: + return ErrNotConnected + } +} + +func (c *Communicator) Send(m proto.Message) error { + if !c.IsAlive() { + return ErrNotConnected + } + data, err := proto.Marshal(m) + if err != nil { + return err + } + return c.QueuePacket(&packets.PacketBase{ + TypeName: TypeNameOf(m), + Nonce: c.nextNonce(), + Data: data, + }) +} + +func (c *Communicator) SendRequest(req proto.Message, resType proto.Message, timeout time.Duration) (proto.Message, error) { + if !c.IsAlive() { + return nil, ErrNotConnected + } + if timeout <= 0 { + timeout = 30 * time.Second + } + + requestNonce := c.nextNonce() + pending := &pendingRequest{ + expectedTypeName: TypeNameOf(resType), + ch: make(chan proto.Message, 1), + errCh: make(chan error, 1), + } + + c.mu.Lock() + c.pendingRequests[requestNonce] = pending + c.mu.Unlock() + + data, err := proto.Marshal(req) + if err != nil { + c.removePending(requestNonce) + return nil, err + } + err = c.QueuePacket(&packets.PacketBase{ + TypeName: TypeNameOf(req), + Nonce: requestNonce, + Data: data, + }) + if err != nil { + c.removePending(requestNonce) + return nil, err + } + + timer := time.NewTimer(timeout) + defer timer.Stop() + + select { + case res := <-pending.ch: + return res, nil + case err := <-pending.errCh: + return nil, err + case <-timer.C: + c.removePending(requestNonce) + return nil, fmt.Errorf("request timeout for nonce %d", requestNonce) + case <-c.closed: + c.removePending(requestNonce) + return nil, ErrNotConnected + } +} + +func (c *Communicator) AddListener(typeName string, fn func(proto.Message)) { + c.mu.Lock() + defer c.mu.Unlock() + + if _, ok := c.reqHandlers[typeName]; ok { + panic(fmt.Sprintf("type %s is already registered with AddRequestListener", typeName)) + } + c.handlers[typeName] = append(c.handlers[typeName], fn) +} + +func (c *Communicator) RemoveListeners(typeName string) { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.handlers, typeName) +} + +func (c *Communicator) AddRequestListener(typeName string, fn func(proto.Message, int32)) { + c.mu.Lock() + defer c.mu.Unlock() + + if len(c.handlers[typeName]) > 0 { + panic(fmt.Sprintf("type %s is already registered with AddListener", typeName)) + } + if _, ok := c.reqHandlers[typeName]; ok { + panic(fmt.Sprintf("type %s is already registered with AddRequestListener", typeName)) + } + c.reqHandlers[typeName] = fn +} + +func (c *Communicator) OnReceivedData(typeName string, data []byte, incomingNonce, responseNonce int32) { + if responseNonce > 0 { + c.handleResponse(typeName, data, responseNonce) + return + } + + c.mu.RLock() + reqHandler := c.reqHandlers[typeName] + listeners := append([]func(proto.Message){}, c.handlers[typeName]...) + c.mu.RUnlock() + + if reqHandler != nil { + msg, err := c.parse(typeName, data) + if err != nil { + return + } + go reqHandler(msg, incomingNonce) + return + } + + if len(listeners) == 0 { + return + } + msg, err := c.parse(typeName, data) + if err != nil { + return + } + for _, listener := range listeners { + listener(msg) + } +} + +func (c *Communicator) handleResponse(typeName string, data []byte, responseNonce int32) { + pending := c.removePending(responseNonce) + if pending == nil { + return + } + if typeName != pending.expectedTypeName { + pending.errCh <- fmt.Errorf("response type mismatch for nonce %d: expected %s, got %s", responseNonce, pending.expectedTypeName, typeName) + return + } + msg, err := c.parse(typeName, data) + if err != nil { + pending.errCh <- err + return + } + pending.ch <- msg +} + +func (c *Communicator) parse(typeName string, data []byte) (proto.Message, error) { + c.mu.RLock() + parser := c.parserMap[typeName] + c.mu.RUnlock() + if parser == nil { + return nil, fmt.Errorf("protobuf parser is not registered for type %s", typeName) + } + return parser(data) +} + +func (c *Communicator) removePending(nonce int32) *pendingRequest { + c.mu.Lock() + defer c.mu.Unlock() + pending := c.pendingRequests[nonce] + delete(c.pendingRequests, nonce) + return pending +} + +func AddListenerTyped[T proto.Message](c *Communicator, fn func(T)) { + typeName := TypeNameOf(newMessageOf[T]()) + c.AddListener(typeName, func(m proto.Message) { + typed, ok := m.(T) + if !ok { + panic(fmt.Sprintf("received %T for listener %s", m, typeName)) + } + fn(typed) + }) +} + +func AddRequestListenerTyped[Req proto.Message, Res proto.Message](c *Communicator, fn func(Req) (Res, error)) { + reqTypeName := TypeNameOf(newMessageOf[Req]()) + c.AddRequestListener(reqTypeName, func(m proto.Message, requestNonce int32) { + req, ok := m.(Req) + if !ok { + return + } + res, err := fn(req) + if err != nil || !c.IsAlive() { + return + } + data, err := proto.Marshal(res) + if err != nil { + return + } + _ = c.QueuePacket(&packets.PacketBase{ + TypeName: TypeNameOf(res), + Nonce: c.nextNonce(), + ResponseNonce: requestNonce, + Data: data, + }) + }) +} + +func SendRequestTyped[Req proto.Message, Res proto.Message](c *Communicator, req Req, timeout time.Duration) (Res, error) { + resType := newMessageOf[Res]() + msg, err := c.SendRequest(req, resType, timeout) + if err != nil { + var zero Res + return zero, err + } + res, ok := msg.(Res) + if !ok { + var zero Res + return zero, fmt.Errorf("received %T, expected %T", msg, resType) + } + return res, nil +} + +func newMessageOf[T proto.Message]() T { + var zero T + t := reflect.TypeOf(zero) + if t == nil || t.Kind() != reflect.Ptr { + panic("protobuf type parameter must be a pointer message type") + } + return reflect.New(t.Elem()).Interface().(T) +} diff --git a/go/communicator_test.go b/go/communicator_test.go new file mode 100644 index 0000000..1e4785d --- /dev/null +++ b/go/communicator_test.go @@ -0,0 +1,99 @@ +package toki_socket + +import ( + "sync" + "testing" + "time" + + "google.golang.org/protobuf/proto" + + "toki-labs.com/toki_socket/go/packets" +) + +type fakeTransport struct { + mu sync.Mutex + packets []*packets.PacketBase +} + +func (f *fakeTransport) WritePacket(base *packets.PacketBase) error { + f.mu.Lock() + defer f.mu.Unlock() + f.packets = append(f.packets, base) + return nil +} + +func (f *fakeTransport) Close() error { + return nil +} + +func (f *fakeTransport) sent() []*packets.PacketBase { + f.mu.Lock() + defer f.mu.Unlock() + return append([]*packets.PacketBase{}, f.packets...) +} + +func testParserMap() ParserMap { + return ParserMap{ + TypeNameOf(&packets.TestData{}): func(b []byte) (proto.Message, error) { + m := &packets.TestData{} + return m, proto.Unmarshal(b, m) + }, + } +} + +func TestSendRequestTypeMismatch(t *testing.T) { + transport := &fakeTransport{} + communicator := NewCommunicator(transport, testParserMap()) + defer communicator.shutdown() + + errCh := make(chan error, 1) + go func() { + _, err := SendRequestTyped[*packets.TestData, *packets.TestData]( + communicator, + &packets.TestData{Index: 1, Message: "hello"}, + time.Second, + ) + errCh <- err + }() + + var requestNonce int32 + for deadline := time.Now().Add(time.Second); time.Now().Before(deadline); { + sent := transport.sent() + if len(sent) == 1 { + requestNonce = sent[0].GetNonce() + break + } + time.Sleep(time.Millisecond) + } + if requestNonce == 0 { + t.Fatal("request packet was not sent") + } + + hb, err := proto.Marshal(&packets.HeartBeat{}) + if err != nil { + t.Fatal(err) + } + communicator.OnReceivedData(TypeNameOf(&packets.HeartBeat{}), hb, 0, requestNonce) + + err = <-errCh + if err == nil { + t.Fatal("expected response type mismatch error") + } +} + +func TestListenerAndRequestListenerConflict(t *testing.T) { + transport := &fakeTransport{} + communicator := NewCommunicator(transport, testParserMap()) + defer communicator.shutdown() + + AddListenerTyped[*packets.TestData](communicator, func(*packets.TestData) {}) + + defer func() { + if recover() == nil { + t.Fatal("expected panic") + } + }() + AddRequestListenerTyped[*packets.TestData, *packets.TestData](communicator, func(req *packets.TestData) (*packets.TestData, error) { + return req, nil + }) +} diff --git a/go/example/tcp_echo/main.go b/go/example/tcp_echo/main.go new file mode 100644 index 0000000..bf3ef6d --- /dev/null +++ b/go/example/tcp_echo/main.go @@ -0,0 +1,60 @@ +package main + +import ( + "context" + "fmt" + "log" + "net" + "time" + + toki "toki-labs.com/toki_socket/go" + "toki-labs.com/toki_socket/go/packets" + + "google.golang.org/protobuf/proto" +) + +func parserMap() toki.ParserMap { + return toki.ParserMap{ + toki.TypeNameOf(&packets.TestData{}): func(b []byte) (proto.Message, error) { + m := &packets.TestData{} + return m, proto.Unmarshal(b, m) + }, + } +} + +func main() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + server := toki.NewTcpServer("127.0.0.1", 19090, func(conn net.Conn) *toki.TcpClient { + return toki.NewTcpClient(conn, 30, 10, parserMap()) + }) + server.OnClientConnected = func(client *toki.TcpClient) { + toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](&client.Communicator, func(req *packets.TestData) (*packets.TestData, error) { + return &packets.TestData{ + Index: req.GetIndex(), + Message: "echo: " + req.GetMessage(), + }, nil + }) + } + if err := server.Start(ctx); err != nil { + log.Fatal(err) + } + defer server.Stop() + + client, err := toki.DialTcp(ctx, "127.0.0.1", 19090, 30, 10, parserMap()) + if err != nil { + log.Fatal(err) + } + defer client.Close() + + res, err := toki.SendRequestTyped[*packets.TestData, *packets.TestData]( + &client.Communicator, + &packets.TestData{Index: 1, Message: "hello tcp"}, + 2*time.Second, + ) + if err != nil { + log.Fatal(err) + } + fmt.Println(res.GetMessage()) +} diff --git a/go/example/ws_echo/main.go b/go/example/ws_echo/main.go new file mode 100644 index 0000000..562dd47 --- /dev/null +++ b/go/example/ws_echo/main.go @@ -0,0 +1,60 @@ +package main + +import ( + "context" + "fmt" + "log" + "time" + + toki "toki-labs.com/toki_socket/go" + "toki-labs.com/toki_socket/go/packets" + + "google.golang.org/protobuf/proto" + "nhooyr.io/websocket" +) + +func parserMap() toki.ParserMap { + return toki.ParserMap{ + toki.TypeNameOf(&packets.TestData{}): func(b []byte) (proto.Message, error) { + m := &packets.TestData{} + return m, proto.Unmarshal(b, m) + }, + } +} + +func main() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + server := toki.NewWsServer("127.0.0.1", 19092, "/", func(conn *websocket.Conn) *toki.WsClient { + return toki.NewWsClient(conn, 30, 10, parserMap()) + }) + server.OnClientConnected = func(client *toki.WsClient) { + toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](&client.Communicator, func(req *packets.TestData) (*packets.TestData, error) { + return &packets.TestData{ + Index: req.GetIndex(), + Message: "echo: " + req.GetMessage(), + }, nil + }) + } + if err := server.Start(ctx); err != nil { + log.Fatal(err) + } + defer server.Stop() + + client, err := toki.DialWs(ctx, "127.0.0.1", 19092, "/", parserMap()) + if err != nil { + log.Fatal(err) + } + defer client.Close() + + res, err := toki.SendRequestTyped[*packets.TestData, *packets.TestData]( + &client.Communicator, + &packets.TestData{Index: 1, Message: "hello ws"}, + 2*time.Second, + ) + if err != nil { + log.Fatal(err) + } + fmt.Println(res.GetMessage()) +} diff --git a/go/go.mod b/go/go.mod new file mode 100644 index 0000000..645d80e --- /dev/null +++ b/go/go.mod @@ -0,0 +1,8 @@ +module toki-labs.com/toki_socket/go + +go 1.22 + +require ( + google.golang.org/protobuf v1.36.5 + nhooyr.io/websocket v1.8.17 +) diff --git a/go/go.sum b/go/go.sum new file mode 100644 index 0000000..f714725 --- /dev/null +++ b/go/go.sum @@ -0,0 +1,8 @@ +github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU= +github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM= +google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE= +nhooyr.io/websocket v1.8.17 h1:KEVeLJkUywCKVsnLIDlD/5gtayKp8VoCkksHCGGfT9Y= +nhooyr.io/websocket v1.8.17/go.mod h1:rN9OFWIUwuxg4fR5tELlYC04bXYowCP9GX47ivo2l+c= diff --git a/go/heartbeat_timer.go b/go/heartbeat_timer.go new file mode 100644 index 0000000..7168e00 --- /dev/null +++ b/go/heartbeat_timer.go @@ -0,0 +1,38 @@ +package toki_socket + +import ( + "sync" + "time" +) + +type HeartbeatTimer struct { + mu sync.Mutex + timer *time.Timer + callback func() +} + +func NewHeartbeatTimer(d time.Duration, cb func()) *HeartbeatTimer { + h := &HeartbeatTimer{callback: cb} + h.Reset(d) + return h +} + +func (h *HeartbeatTimer) Reset(d time.Duration) { + h.mu.Lock() + defer h.mu.Unlock() + + if h.timer != nil { + h.timer.Stop() + } + h.timer = time.AfterFunc(d, h.callback) +} + +func (h *HeartbeatTimer) Stop() { + h.mu.Lock() + defer h.mu.Unlock() + + if h.timer != nil { + h.timer.Stop() + h.timer = nil + } +} diff --git a/go/packets/message_common.pb.go b/go/packets/message_common.pb.go new file mode 100644 index 0000000..7054c5a --- /dev/null +++ b/go/packets/message_common.pb.go @@ -0,0 +1,244 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.36.11 +// protoc v5.29.3 +// source: packets/message_common.proto + +package packets + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" + unsafe "unsafe" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type PacketBase struct { + state protoimpl.MessageState `protogen:"open.v1"` + TypeName string `protobuf:"bytes,1,opt,name=typeName,proto3" json:"typeName,omitempty"` + Nonce int32 `protobuf:"varint,2,opt,name=nonce,proto3" json:"nonce,omitempty"` + Data []byte `protobuf:"bytes,3,opt,name=data,proto3" json:"data,omitempty"` + ResponseNonce int32 `protobuf:"varint,4,opt,name=responseNonce,proto3" json:"responseNonce,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *PacketBase) Reset() { + *x = PacketBase{} + mi := &file_packets_message_common_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *PacketBase) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PacketBase) ProtoMessage() {} + +func (x *PacketBase) ProtoReflect() protoreflect.Message { + mi := &file_packets_message_common_proto_msgTypes[0] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PacketBase.ProtoReflect.Descriptor instead. +func (*PacketBase) Descriptor() ([]byte, []int) { + return file_packets_message_common_proto_rawDescGZIP(), []int{0} +} + +func (x *PacketBase) GetTypeName() string { + if x != nil { + return x.TypeName + } + return "" +} + +func (x *PacketBase) GetNonce() int32 { + if x != nil { + return x.Nonce + } + return 0 +} + +func (x *PacketBase) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +func (x *PacketBase) GetResponseNonce() int32 { + if x != nil { + return x.ResponseNonce + } + return 0 +} + +type HeartBeat struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *HeartBeat) Reset() { + *x = HeartBeat{} + mi := &file_packets_message_common_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HeartBeat) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HeartBeat) ProtoMessage() {} + +func (x *HeartBeat) ProtoReflect() protoreflect.Message { + mi := &file_packets_message_common_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HeartBeat.ProtoReflect.Descriptor instead. +func (*HeartBeat) Descriptor() ([]byte, []int) { + return file_packets_message_common_proto_rawDescGZIP(), []int{1} +} + +type TestData struct { + state protoimpl.MessageState `protogen:"open.v1"` + Index int32 `protobuf:"varint,1,opt,name=index,proto3" json:"index,omitempty"` + Message string `protobuf:"bytes,2,opt,name=message,proto3" json:"message,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *TestData) Reset() { + *x = TestData{} + mi := &file_packets_message_common_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *TestData) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TestData) ProtoMessage() {} + +func (x *TestData) ProtoReflect() protoreflect.Message { + mi := &file_packets_message_common_proto_msgTypes[2] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use TestData.ProtoReflect.Descriptor instead. +func (*TestData) Descriptor() ([]byte, []int) { + return file_packets_message_common_proto_rawDescGZIP(), []int{2} +} + +func (x *TestData) GetIndex() int32 { + if x != nil { + return x.Index + } + return 0 +} + +func (x *TestData) GetMessage() string { + if x != nil { + return x.Message + } + return "" +} + +var File_packets_message_common_proto protoreflect.FileDescriptor + +const file_packets_message_common_proto_rawDesc = "" + + "\n" + + "\x1cpackets/message_common.proto\"x\n" + + "\n" + + "PacketBase\x12\x1a\n" + + "\btypeName\x18\x01 \x01(\tR\btypeName\x12\x14\n" + + "\x05nonce\x18\x02 \x01(\x05R\x05nonce\x12\x12\n" + + "\x04data\x18\x03 \x01(\fR\x04data\x12$\n" + + "\rresponseNonce\x18\x04 \x01(\x05R\rresponseNonce\"\v\n" + + "\tHeartBeat\":\n" + + "\bTestData\x12\x14\n" + + "\x05index\x18\x01 \x01(\x05R\x05index\x12\x18\n" + + "\amessage\x18\x02 \x01(\tR\amessageB&Z$toki-labs.com/toki_socket/go/packetsb\x06proto3" + +var ( + file_packets_message_common_proto_rawDescOnce sync.Once + file_packets_message_common_proto_rawDescData []byte +) + +func file_packets_message_common_proto_rawDescGZIP() []byte { + file_packets_message_common_proto_rawDescOnce.Do(func() { + file_packets_message_common_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_packets_message_common_proto_rawDesc), len(file_packets_message_common_proto_rawDesc))) + }) + return file_packets_message_common_proto_rawDescData +} + +var file_packets_message_common_proto_msgTypes = make([]protoimpl.MessageInfo, 3) +var file_packets_message_common_proto_goTypes = []any{ + (*PacketBase)(nil), // 0: PacketBase + (*HeartBeat)(nil), // 1: HeartBeat + (*TestData)(nil), // 2: TestData +} +var file_packets_message_common_proto_depIdxs = []int32{ + 0, // [0:0] is the sub-list for method output_type + 0, // [0:0] is the sub-list for method input_type + 0, // [0:0] is the sub-list for extension type_name + 0, // [0:0] is the sub-list for extension extendee + 0, // [0:0] is the sub-list for field type_name +} + +func init() { file_packets_message_common_proto_init() } +func file_packets_message_common_proto_init() { + if File_packets_message_common_proto != nil { + return + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_packets_message_common_proto_rawDesc), len(file_packets_message_common_proto_rawDesc)), + NumEnums: 0, + NumMessages: 3, + NumExtensions: 0, + NumServices: 0, + }, + GoTypes: file_packets_message_common_proto_goTypes, + DependencyIndexes: file_packets_message_common_proto_depIdxs, + MessageInfos: file_packets_message_common_proto_msgTypes, + }.Build() + File_packets_message_common_proto = out.File + file_packets_message_common_proto_goTypes = nil + file_packets_message_common_proto_depIdxs = nil +} diff --git a/go/packets/message_common.proto b/go/packets/message_common.proto new file mode 100644 index 0000000..b7490f5 --- /dev/null +++ b/go/packets/message_common.proto @@ -0,0 +1,17 @@ +syntax = "proto3"; + +option go_package = "toki-labs.com/toki_socket/go/packets"; + +message PacketBase { + string typeName = 1; + int32 nonce = 2; + bytes data = 3; + int32 responseNonce = 4; +} + +message HeartBeat {} + +message TestData { + int32 index = 1; + string message = 2; +} diff --git a/go/socket_test.go b/go/socket_test.go new file mode 100644 index 0000000..e3063f1 --- /dev/null +++ b/go/socket_test.go @@ -0,0 +1,368 @@ +package toki_socket + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/binary" + "encoding/pem" + "io" + "math/big" + "net" + "testing" + "time" + + "google.golang.org/protobuf/proto" + "nhooyr.io/websocket" + + "toki-labs.com/toki_socket/go/packets" +) + +func TestTcpRequestResponse(t *testing.T) { + port := freePort(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + server := NewTcpServer("127.0.0.1", port, func(conn net.Conn) *TcpClient { + return NewTcpClient(conn, 0, 0, testParserMap()) + }) + server.OnClientConnected = func(client *TcpClient) { + AddRequestListenerTyped[*packets.TestData, *packets.TestData](&client.Communicator, func(req *packets.TestData) (*packets.TestData, error) { + return &packets.TestData{ + Index: req.GetIndex() * 2, + Message: "echo: " + req.GetMessage(), + }, nil + }) + } + if err := server.Start(ctx); err != nil { + t.Fatal(err) + } + defer server.Stop() + + client, err := DialTcp(ctx, "127.0.0.1", port, 0, 0, testParserMap()) + if err != nil { + t.Fatal(err) + } + defer client.Close() + + res, err := SendRequestTyped[*packets.TestData, *packets.TestData]( + &client.Communicator, + &packets.TestData{Index: 21, Message: "hello"}, + 2*time.Second, + ) + if err != nil { + t.Fatal(err) + } + if res.GetIndex() != 42 || res.GetMessage() != "echo: hello" { + t.Fatalf("unexpected response: %v", res) + } +} + +func TestTcpBroadcast(t *testing.T) { + port := freePort(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + server := NewTcpServer("127.0.0.1", port, func(conn net.Conn) *TcpClient { + return NewTcpClient(conn, 0, 0, testParserMap()) + }) + if err := server.Start(ctx); err != nil { + t.Fatal(err) + } + defer server.Stop() + + client, err := DialTcp(ctx, "127.0.0.1", port, 0, 0, testParserMap()) + if err != nil { + t.Fatal(err) + } + defer client.Close() + + received := make(chan *packets.TestData, 1) + AddListenerTyped[*packets.TestData](&client.Communicator, func(m *packets.TestData) { + received <- m + }) + + for deadline := time.Now().Add(time.Second); time.Now().Before(deadline) && len(server.Clients()) == 0; { + time.Sleep(time.Millisecond) + } + if err := server.Broadcast(&packets.TestData{Index: 9, Message: "broadcast"}); err != nil { + t.Fatal(err) + } + + select { + case msg := <-received: + if msg.GetMessage() != "broadcast" { + t.Fatalf("unexpected broadcast: %v", msg) + } + case <-time.After(2 * time.Second): + t.Fatal("broadcast timed out") + } +} + +func TestTcpBroadcastAttemptsAllClients(t *testing.T) { + badConn, badPeer := net.Pipe() + badClient := NewTcpClient(badConn, 0, 0, testParserMap()) + _ = badClient.Close() + _ = badPeer.Close() + + goodConn, goodPeer := net.Pipe() + defer goodPeer.Close() + goodClient := NewTcpClient(goodConn, 0, 0, testParserMap()) + defer goodClient.Close() + + received := make(chan *packets.PacketBase, 1) + go func() { + base, err := readTCPPacket(goodPeer) + if err == nil { + received <- base + } + }() + + server := &TcpServer{clients: []*TcpClient{badClient, goodClient}} + err := server.Broadcast(&packets.TestData{Index: 9, Message: "best effort"}) + if err == nil { + t.Fatal("expected broadcast to report the failed client") + } + + select { + case base := <-received: + if base.GetTypeName() != TypeNameOf(&packets.TestData{}) { + t.Fatalf("unexpected packet type: %s", base.GetTypeName()) + } + case <-time.After(2 * time.Second): + t.Fatal("broadcast did not attempt the healthy client") + } +} + +func TestWsRequestResponse(t *testing.T) { + port := freePort(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + server := NewWsServer("127.0.0.1", port, "/", func(conn *websocket.Conn) *WsClient { + return NewWsClient(conn, 0, 0, testParserMap()) + }) + server.OnClientConnected = func(client *WsClient) { + AddRequestListenerTyped[*packets.TestData, *packets.TestData](&client.Communicator, func(req *packets.TestData) (*packets.TestData, error) { + return &packets.TestData{ + Index: req.GetIndex() * 2, + Message: "echo: " + req.GetMessage(), + }, nil + }) + } + if err := server.Start(ctx); err != nil { + t.Fatal(err) + } + defer server.Stop() + + client, err := DialWsWithHeartbeat(ctx, "127.0.0.1", port, "/", 0, 0, testParserMap()) + if err != nil { + t.Fatal(err) + } + defer client.Close() + + res, err := SendRequestTyped[*packets.TestData, *packets.TestData]( + &client.Communicator, + &packets.TestData{Index: 21, Message: "hello ws"}, + 2*time.Second, + ) + if err != nil { + t.Fatal(err) + } + if res.GetIndex() != 42 || res.GetMessage() != "echo: hello ws" { + t.Fatalf("unexpected response: %v", res) + } +} + +func TestTcpSendReceive(t *testing.T) { + port := freePort(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + received := make(chan *packets.TestData, 1) + server := NewTcpServer("127.0.0.1", port, func(conn net.Conn) *TcpClient { + return NewTcpClient(conn, 0, 0, testParserMap()) + }) + server.OnClientConnected = func(client *TcpClient) { + AddListenerTyped[*packets.TestData](&client.Communicator, func(m *packets.TestData) { + received <- m + }) + } + if err := server.Start(ctx); err != nil { + t.Fatal(err) + } + defer server.Stop() + + client, err := DialTcp(ctx, "127.0.0.1", port, 0, 0, testParserMap()) + if err != nil { + t.Fatal(err) + } + defer client.Close() + + if err := client.Send(&packets.TestData{Index: 42, Message: "hello toki-socket"}); err != nil { + t.Fatal(err) + } + + select { + case msg := <-received: + if msg.GetIndex() != 42 { + t.Fatalf("unexpected message: %v", msg) + } + case <-time.After(2 * time.Second): + t.Fatal("receive timed out") + } +} + +func TestHeartbeatDisconnectsWithoutResponse(t *testing.T) { + clientConn, peerConn := net.Pipe() + defer peerConn.Close() + + go func() { + _, _ = io.Copy(io.Discard, peerConn) + }() + + client := NewTcpClient(clientConn, 1, 1, testParserMap()) + defer client.Close() + + disconnected := make(chan struct{}, 1) + client.AddDisconnectListener(func(*TcpClient) { + disconnected <- struct{}{} + }) + + select { + case <-disconnected: + case <-time.After(3500 * time.Millisecond): + t.Fatal("heartbeat timeout did not disconnect client") + } + if client.IsAlive() { + t.Fatal("client is still marked alive after heartbeat timeout") + } +} + +func TestTypeNameMatchesDartConvention(t *testing.T) { + if got := TypeNameOf(&packets.TestData{}); got != "TestData" { + t.Fatalf("type name = %q, want TestData", got) + } +} + +func freePort(t *testing.T) int { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + return ln.Addr().(*net.TCPAddr).Port +} + +func readTCPPacket(conn net.Conn) (*packets.PacketBase, error) { + header := make([]byte, 4) + if _, err := io.ReadFull(conn, header); err != nil { + return nil, err + } + length := binary.BigEndian.Uint32(header) + packetBytes := make([]byte, int(length)) + if _, err := io.ReadFull(conn, packetBytes); err != nil { + return nil, err + } + base := &packets.PacketBase{} + if err := proto.Unmarshal(packetBytes, base); err != nil { + return nil, err + } + return base, nil +} + +// generateSelfSignedCert creates an in-memory self-signed certificate valid +// for 127.0.0.1. Returns the server tls.Config and a client tls.Config that +// trusts the generated certificate. +func generateSelfSignedCert(t *testing.T) (serverCfg, clientCfg *tls.Config) { + t.Helper() + + priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatal(err) + } + + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "toki-socket-test"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, + } + + certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &priv.PublicKey, priv) + if err != nil { + t.Fatal(err) + } + privDER, err := x509.MarshalECPrivateKey(priv) + if err != nil { + t.Fatal(err) + } + + certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER}) + keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: privDER}) + + cert, err := tls.X509KeyPair(certPEM, keyPEM) + if err != nil { + t.Fatal(err) + } + + serverCfg = &tls.Config{Certificates: []tls.Certificate{cert}} + + pool := x509.NewCertPool() + parsed, err := x509.ParseCertificate(certDER) + if err != nil { + t.Fatal(err) + } + pool.AddCert(parsed) + clientCfg = &tls.Config{RootCAs: pool, ServerName: "127.0.0.1"} + + return serverCfg, clientCfg +} + +func TestTLSTcp(t *testing.T) { + port := freePort(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + serverTLS, clientTLS := generateSelfSignedCert(t) + + server := NewTcpServerTLS("127.0.0.1", port, serverTLS, func(conn net.Conn) *TcpClient { + return NewTcpClient(conn, 0, 0, testParserMap()) + }) + server.OnClientConnected = func(client *TcpClient) { + AddRequestListenerTyped[*packets.TestData, *packets.TestData](&client.Communicator, func(req *packets.TestData) (*packets.TestData, error) { + return &packets.TestData{ + Index: req.GetIndex() * 2, + Message: "tls-echo: " + req.GetMessage(), + }, nil + }) + } + if err := server.Start(ctx); err != nil { + t.Fatal(err) + } + defer server.Stop() + + client, err := DialTcpTLS(ctx, "127.0.0.1", port, clientTLS, 0, 0, testParserMap()) + if err != nil { + t.Fatal(err) + } + defer client.Close() + + res, err := SendRequestTyped[*packets.TestData, *packets.TestData]( + &client.Communicator, + &packets.TestData{Index: 7, Message: "hello tls"}, + 2*time.Second, + ) + if err != nil { + t.Fatal(err) + } + if res.GetIndex() != 14 || res.GetMessage() != "tls-echo: hello tls" { + t.Fatalf("unexpected response: %v", res) + } +} diff --git a/go/tcp_client.go b/go/tcp_client.go new file mode 100644 index 0000000..01fc2d9 --- /dev/null +++ b/go/tcp_client.go @@ -0,0 +1,216 @@ +package toki_socket + +import ( + "context" + "crypto/tls" + "encoding/binary" + "fmt" + "io" + "net" + "sync" + "time" + + "google.golang.org/protobuf/proto" + + "toki-labs.com/toki_socket/go/packets" +) + +const MaxPacketSize = 64 << 20 + +type TcpClient struct { + Communicator + conn net.Conn + writeMu sync.Mutex // Defensive if WritePacket is called outside the queued writer. + heartbeatInterval time.Duration + heartbeatWait time.Duration + waitingHBResponse bool + hbTimer *HeartbeatTimer + hbMu sync.Mutex + disconnectListeners []func(*TcpClient) + disconnectMu sync.Mutex + tcpCloseOnce sync.Once +} + +func NewTcpClient(conn net.Conn, intervalSec, waitSec int, parserMap ParserMap) *TcpClient { + c := &TcpClient{ + conn: conn, + heartbeatInterval: time.Duration(intervalSec) * time.Second, + heartbeatWait: time.Duration(waitSec) * time.Second, + } + c.Communicator.Initialize(c, parserMap) + c.SetWriteErrorHandler(func(error) { + c.onDisconnected() + }) + c.AddListener(TypeNameOf(&packets.HeartBeat{}), func(m proto.Message) { + c.onHeartBeat() + }) + go c.readLoop() + c.sendHeartBeat() + return c +} + +func DialTcp(ctx context.Context, host string, port int, intervalSec, waitSec int, parserMap ParserMap) (*TcpClient, error) { + var d net.Dialer + conn, err := d.DialContext(ctx, "tcp", fmt.Sprintf("%s:%d", host, port)) + if err != nil { + return nil, err + } + return NewTcpClient(conn, intervalSec, waitSec, parserMap), nil +} + +func DialTcpTLS(ctx context.Context, host string, port int, tlsCfg *tls.Config, intervalSec, waitSec int, parserMap ParserMap) (*TcpClient, error) { + d := tls.Dialer{Config: tlsCfg} + conn, err := d.DialContext(ctx, "tcp", fmt.Sprintf("%s:%d", host, port)) + if err != nil { + return nil, err + } + return NewTcpClient(conn, intervalSec, waitSec, parserMap), nil +} + +func (c *TcpClient) readLoop() { + header := make([]byte, 4) + for c.IsAlive() { + if _, err := io.ReadFull(c.conn, header); err != nil { + c.onDisconnected() + return + } + length := binary.BigEndian.Uint32(header) + if length == 0 { + continue + } + if length > MaxPacketSize { + c.onDisconnected() + return + } + packetBytes := make([]byte, int(length)) + if _, err := io.ReadFull(c.conn, packetBytes); err != nil { + c.onDisconnected() + return + } + base := &packets.PacketBase{} + if err := proto.Unmarshal(packetBytes, base); err != nil { + c.onDisconnected() + return + } + c.OnReceivedData(base.GetTypeName(), base.GetData(), base.GetNonce(), base.GetResponseNonce()) + c.sendHeartBeat() + } +} + +func (c *TcpClient) WritePacket(base *packets.PacketBase) error { + b, err := proto.Marshal(base) + if err != nil { + return err + } + header := make([]byte, 4) + binary.BigEndian.PutUint32(header, uint32(len(b))) + + c.writeMu.Lock() + defer c.writeMu.Unlock() + if err := writeFull(c.conn, header); err != nil { + return err + } + return writeFull(c.conn, b) +} + +func (c *TcpClient) AddDisconnectListener(handler func(*TcpClient)) { + c.disconnectMu.Lock() + defer c.disconnectMu.Unlock() + c.disconnectListeners = append(c.disconnectListeners, handler) +} + +func (c *TcpClient) RemoveDisconnectListeners() { + c.disconnectMu.Lock() + defer c.disconnectMu.Unlock() + c.disconnectListeners = nil +} + +func (c *TcpClient) Close() error { + var err error + c.tcpCloseOnce.Do(func() { + c.Communicator.shutdown() + c.stopHeartbeat() + err = c.conn.Close() + c.notifyDisconnected() + }) + return err +} + +func (c *TcpClient) sendHeartBeat() { + if !c.IsAlive() || c.heartbeatInterval <= 0 { + return + } + c.hbMu.Lock() + if c.hbTimer != nil { + c.hbTimer.Stop() + } + c.hbTimer = NewHeartbeatTimer(c.heartbeatInterval, func() { + if !c.IsAlive() { + return + } + c.hbMu.Lock() + c.waitingHBResponse = true + c.hbMu.Unlock() + _ = c.Send(&packets.HeartBeat{}) + c.hbMu.Lock() + if c.hbTimer != nil { + c.hbTimer.Stop() + } + c.hbTimer = NewHeartbeatTimer(c.heartbeatWait, func() { + if c.IsAlive() { + c.onDisconnected() + } + }) + c.hbMu.Unlock() + }) + c.hbMu.Unlock() +} + +func (c *TcpClient) onHeartBeat() { + c.hbMu.Lock() + if c.waitingHBResponse { + c.waitingHBResponse = false + c.hbMu.Unlock() + return + } + c.hbMu.Unlock() + _ = c.Send(&packets.HeartBeat{}) +} + +func (c *TcpClient) stopHeartbeat() { + c.hbMu.Lock() + defer c.hbMu.Unlock() + if c.hbTimer != nil { + c.hbTimer.Stop() + c.hbTimer = nil + } +} + +func (c *TcpClient) onDisconnected() { + _ = c.Close() +} + +func (c *TcpClient) notifyDisconnected() { + c.disconnectMu.Lock() + listeners := append([]func(*TcpClient){}, c.disconnectListeners...) + c.disconnectListeners = nil + c.disconnectMu.Unlock() + + for _, listener := range listeners { + listener(c) + } +} + +func writeFull(w io.Writer, b []byte) error { + for len(b) > 0 { + n, err := w.Write(b) + if err != nil { + return err + } + if n == 0 { + return io.ErrShortWrite + } + b = b[n:] + } + return nil +} diff --git a/go/tcp_server.go b/go/tcp_server.go new file mode 100644 index 0000000..9c029b7 --- /dev/null +++ b/go/tcp_server.go @@ -0,0 +1,146 @@ +package toki_socket + +import ( + "context" + "crypto/tls" + "errors" + "fmt" + "net" + "sync" + + "google.golang.org/protobuf/proto" +) + +type TcpServer struct { + host string + port int + tlsConfig *tls.Config + newClient func(net.Conn) *TcpClient + clients []*TcpClient + mu sync.Mutex + listener net.Listener + started bool + OnClientConnected func(*TcpClient) +} + +func NewTcpServer(host string, port int, newClient func(net.Conn) *TcpClient) *TcpServer { + return &TcpServer{ + host: host, + port: port, + newClient: newClient, + OnClientConnected: func(*TcpClient) {}, + } +} + +func NewTcpServerTLS(host string, port int, tlsCfg *tls.Config, newClient func(net.Conn) *TcpClient) *TcpServer { + s := NewTcpServer(host, port, newClient) + s.tlsConfig = tlsCfg + return s +} + +func (s *TcpServer) Started() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.started +} + +func (s *TcpServer) IsSecure() bool { + return s.tlsConfig != nil +} + +func (s *TcpServer) Clients() []*TcpClient { + s.mu.Lock() + defer s.mu.Unlock() + return append([]*TcpClient{}, s.clients...) +} + +func (s *TcpServer) Start(ctx context.Context) error { + addr := fmt.Sprintf("%s:%d", s.host, s.port) + ln, err := net.Listen("tcp", addr) + if err != nil { + return err + } + if s.tlsConfig != nil { + ln = tls.NewListener(ln, s.tlsConfig) + } + + s.mu.Lock() + s.listener = ln + s.started = true + s.mu.Unlock() + + go func() { + <-ctx.Done() + _ = s.Stop() + }() + go s.acceptLoop() + return nil +} + +func (s *TcpServer) Stop() error { + s.mu.Lock() + if !s.started { + s.mu.Unlock() + return nil + } + s.started = false + ln := s.listener + clients := append([]*TcpClient{}, s.clients...) + s.clients = nil + s.mu.Unlock() + + var err error + if ln != nil { + err = ln.Close() + } + for _, client := range clients { + _ = client.Close() + } + return err +} + +func (s *TcpServer) Broadcast(m proto.Message) error { + clients := s.Clients() + var errs []error + for _, client := range clients { + if err := client.Send(m); err != nil { + errs = append(errs, err) + } + } + return errors.Join(errs...) +} + +func (s *TcpServer) acceptLoop() { + for { + conn, err := s.listener.Accept() + if err != nil { + s.mu.Lock() + started := s.started + s.mu.Unlock() + if !started || errors.Is(err, net.ErrClosed) { + return + } + continue + } + client := s.newClient(conn) + client.AddDisconnectListener(func(c *TcpClient) { + s.removeClient(c) + }) + s.mu.Lock() + s.clients = append(s.clients, client) + onConnected := s.OnClientConnected + s.mu.Unlock() + onConnected(client) + } +} + +func (s *TcpServer) removeClient(client *TcpClient) { + s.mu.Lock() + defer s.mu.Unlock() + for i, item := range s.clients { + if item == client { + s.clients = append(s.clients[:i], s.clients[i+1:]...) + return + } + } +} diff --git a/go/ws_client.go b/go/ws_client.go new file mode 100644 index 0000000..ebd0a6b --- /dev/null +++ b/go/ws_client.go @@ -0,0 +1,206 @@ +package toki_socket + +import ( + "context" + "crypto/tls" + "fmt" + "net/http" + "sync" + "time" + + "google.golang.org/protobuf/proto" + "nhooyr.io/websocket" + + "toki-labs.com/toki_socket/go/packets" +) + +const ( + DefaultHeartbeatIntervalSec = 30 + DefaultHeartbeatWaitSec = 10 +) + +type WsClient struct { + Communicator + conn *websocket.Conn + writeMu sync.Mutex // Defensive if WritePacket is called outside the queued writer. + readCtx context.Context + cancelRead context.CancelFunc + heartbeatInterval time.Duration + heartbeatWait time.Duration + waitingHBResponse bool + hbTimer *HeartbeatTimer + hbMu sync.Mutex + disconnectListeners []func(*WsClient) + disconnectMu sync.Mutex + wsCloseOnce sync.Once +} + +func NewWsClient(conn *websocket.Conn, intervalSec, waitSec int, parserMap ParserMap) *WsClient { + readCtx, cancelRead := context.WithCancel(context.Background()) + c := &WsClient{ + conn: conn, + readCtx: readCtx, + cancelRead: cancelRead, + heartbeatInterval: time.Duration(intervalSec) * time.Second, + heartbeatWait: time.Duration(waitSec) * time.Second, + } + c.Communicator.Initialize(c, parserMap) + c.SetWriteErrorHandler(func(error) { + c.onDisconnected() + }) + c.AddListener(TypeNameOf(&packets.HeartBeat{}), func(m proto.Message) { + c.onHeartBeat() + }) + go c.readLoop() + c.sendHeartBeat() + return c +} + +func DialWs(ctx context.Context, host string, port int, path string, parserMap ParserMap) (*WsClient, error) { + return DialWsWithHeartbeat(ctx, host, port, path, DefaultHeartbeatIntervalSec, DefaultHeartbeatWaitSec, parserMap) +} + +func DialWsWithHeartbeat(ctx context.Context, host string, port int, path string, intervalSec, waitSec int, parserMap ParserMap) (*WsClient, error) { + conn, _, err := websocket.Dial(ctx, fmt.Sprintf("ws://%s:%d%s", host, port, path), nil) + if err != nil { + return nil, err + } + return NewWsClient(conn, intervalSec, waitSec, parserMap), nil +} + +func DialWss(ctx context.Context, host string, port int, path string, tlsCfg *tls.Config, parserMap ParserMap) (*WsClient, error) { + return DialWssWithHeartbeat(ctx, host, port, path, tlsCfg, DefaultHeartbeatIntervalSec, DefaultHeartbeatWaitSec, parserMap) +} + +func DialWssWithHeartbeat(ctx context.Context, host string, port int, path string, tlsCfg *tls.Config, intervalSec, waitSec int, parserMap ParserMap) (*WsClient, error) { + opts := &websocket.DialOptions{} + if tlsCfg != nil { + opts.HTTPClient = &http.Client{ + Transport: &http.Transport{TLSClientConfig: tlsCfg}, + } + } + conn, _, err := websocket.Dial(ctx, fmt.Sprintf("wss://%s:%d%s", host, port, path), opts) + if err != nil { + return nil, err + } + return NewWsClient(conn, intervalSec, waitSec, parserMap), nil +} + +func (c *WsClient) readLoop() { + for c.IsAlive() { + msgType, b, err := c.conn.Read(c.readCtx) + if err != nil { + c.onDisconnected() + return + } + if msgType != websocket.MessageBinary { + continue + } + base := &packets.PacketBase{} + if err := proto.Unmarshal(b, base); err != nil { + c.onDisconnected() + return + } + c.OnReceivedData(base.GetTypeName(), base.GetData(), base.GetNonce(), base.GetResponseNonce()) + c.sendHeartBeat() + } +} + +func (c *WsClient) WritePacket(base *packets.PacketBase) error { + b, err := proto.Marshal(base) + if err != nil { + return err + } + c.writeMu.Lock() + defer c.writeMu.Unlock() + return c.conn.Write(context.Background(), websocket.MessageBinary, b) +} + +func (c *WsClient) AddDisconnectListener(handler func(*WsClient)) { + c.disconnectMu.Lock() + defer c.disconnectMu.Unlock() + c.disconnectListeners = append(c.disconnectListeners, handler) +} + +func (c *WsClient) RemoveDisconnectListeners() { + c.disconnectMu.Lock() + defer c.disconnectMu.Unlock() + c.disconnectListeners = nil +} + +func (c *WsClient) Close() error { + var err error + c.wsCloseOnce.Do(func() { + c.Communicator.shutdown() + c.stopHeartbeat() + c.cancelRead() + err = c.conn.Close(websocket.StatusNormalClosure, "") + c.notifyDisconnected() + }) + return err +} + +func (c *WsClient) sendHeartBeat() { + if !c.IsAlive() || c.heartbeatInterval <= 0 { + return + } + c.hbMu.Lock() + if c.hbTimer != nil { + c.hbTimer.Stop() + } + c.hbTimer = NewHeartbeatTimer(c.heartbeatInterval, func() { + if !c.IsAlive() { + return + } + c.hbMu.Lock() + c.waitingHBResponse = true + c.hbMu.Unlock() + _ = c.Send(&packets.HeartBeat{}) + c.hbMu.Lock() + if c.hbTimer != nil { + c.hbTimer.Stop() + } + c.hbTimer = NewHeartbeatTimer(c.heartbeatWait, func() { + if c.IsAlive() { + c.onDisconnected() + } + }) + c.hbMu.Unlock() + }) + c.hbMu.Unlock() +} + +func (c *WsClient) onHeartBeat() { + c.hbMu.Lock() + if c.waitingHBResponse { + c.waitingHBResponse = false + c.hbMu.Unlock() + return + } + c.hbMu.Unlock() + _ = c.Send(&packets.HeartBeat{}) +} + +func (c *WsClient) stopHeartbeat() { + c.hbMu.Lock() + defer c.hbMu.Unlock() + if c.hbTimer != nil { + c.hbTimer.Stop() + c.hbTimer = nil + } +} + +func (c *WsClient) onDisconnected() { + _ = c.Close() +} + +func (c *WsClient) notifyDisconnected() { + c.disconnectMu.Lock() + listeners := append([]func(*WsClient){}, c.disconnectListeners...) + c.disconnectListeners = nil + c.disconnectMu.Unlock() + + for _, listener := range listeners { + listener(c) + } +} diff --git a/go/ws_server.go b/go/ws_server.go new file mode 100644 index 0000000..4079092 --- /dev/null +++ b/go/ws_server.go @@ -0,0 +1,156 @@ +package toki_socket + +import ( + "context" + "crypto/tls" + "errors" + "fmt" + "net" + "net/http" + "sync" + + "google.golang.org/protobuf/proto" + "nhooyr.io/websocket" +) + +type WsServer struct { + host string + port int + path string + tlsConfig *tls.Config + newClient func(*websocket.Conn) *WsClient + clients []*WsClient + mu sync.Mutex + httpSrv *http.Server + listener net.Listener + started bool + OnClientConnected func(*WsClient) +} + +func NewWsServer(host string, port int, path string, newClient func(*websocket.Conn) *WsClient) *WsServer { + return &WsServer{ + host: host, + port: port, + path: path, + newClient: newClient, + OnClientConnected: func(*WsClient) {}, + } +} + +func NewWsServerTLS(host string, port int, path string, tlsCfg *tls.Config, newClient func(*websocket.Conn) *WsClient) *WsServer { + s := NewWsServer(host, port, path, newClient) + s.tlsConfig = tlsCfg + return s +} + +func (s *WsServer) Started() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.started +} + +func (s *WsServer) IsSecure() bool { + return s.tlsConfig != nil +} + +func (s *WsServer) Clients() []*WsClient { + s.mu.Lock() + defer s.mu.Unlock() + return append([]*WsClient{}, s.clients...) +} + +func (s *WsServer) Start(ctx context.Context) error { + mux := http.NewServeMux() + mux.HandleFunc(s.path, s.handleWebSocket) + s.httpSrv = &http.Server{ + Addr: fmt.Sprintf("%s:%d", s.host, s.port), + Handler: mux, + TLSConfig: s.tlsConfig, + } + + ln, err := net.Listen("tcp", s.httpSrv.Addr) + if err != nil { + return err + } + if s.tlsConfig != nil { + ln = tls.NewListener(ln, s.tlsConfig) + } + + s.mu.Lock() + s.listener = ln + s.started = true + s.mu.Unlock() + + go func() { + <-ctx.Done() + _ = s.Stop() + }() + go func() { + err := s.httpSrv.Serve(ln) + if err != nil && !errors.Is(err, http.ErrServerClosed) && !errors.Is(err, net.ErrClosed) { + _ = s.Stop() + } + }() + return nil +} + +func (s *WsServer) Stop() error { + s.mu.Lock() + if !s.started { + s.mu.Unlock() + return nil + } + s.started = false + srv := s.httpSrv + clients := append([]*WsClient{}, s.clients...) + s.clients = nil + s.mu.Unlock() + + var err error + if srv != nil { + err = srv.Close() + } + for _, client := range clients { + _ = client.Close() + } + return err +} + +func (s *WsServer) Broadcast(m proto.Message) error { + clients := s.Clients() + var errs []error + for _, client := range clients { + if err := client.Send(m); err != nil { + errs = append(errs, err) + } + } + return errors.Join(errs...) +} + +func (s *WsServer) handleWebSocket(w http.ResponseWriter, r *http.Request) { + conn, err := websocket.Accept(w, r, nil) + if err != nil { + return + } + client := s.newClient(conn) + client.AddDisconnectListener(func(c *WsClient) { + s.removeClient(c) + }) + + s.mu.Lock() + s.clients = append(s.clients, client) + onConnected := s.OnClientConnected + s.mu.Unlock() + onConnected(client) +} + +func (s *WsServer) removeClient(client *WsClient) { + s.mu.Lock() + defer s.mu.Unlock() + for i, item := range s.clients { + if item == client { + s.clients = append(s.clients[:i], s.clients[i+1:]...) + return + } + } +}