410 lines
12 KiB
Go
410 lines
12 KiB
Go
package rpc
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.yuntuo.hk/k2safe/go-omni/ws"
|
|
)
|
|
|
|
type memoryTransport struct {
|
|
ctx context.Context
|
|
other *Peer
|
|
}
|
|
|
|
func (m *memoryTransport) Context() context.Context { return m.ctx }
|
|
func (m *memoryTransport) Send(ctx context.Context, msg ws.Message) error {
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
return m.other.receive(msg)
|
|
}
|
|
func (m *memoryTransport) TrySend(msg ws.Message) error { return m.other.receive(msg) }
|
|
func (m *memoryTransport) Close(int, string) error { return nil }
|
|
func (m *memoryTransport) Abort() { m.other.Stop() }
|
|
|
|
func pair(t testing.TB, cfg Config, setup func(*Peer, *Peer)) (*Peer, *Peer) {
|
|
t.Helper()
|
|
a, err := New(cfg)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
b, err := New(cfg)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if setup != nil {
|
|
setup(a, b)
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
if err = a.bind(&memoryTransport{ctx, b}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err = b.bind(&memoryTransport{ctx, a}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() {
|
|
cancel()
|
|
a.Stop()
|
|
b.Stop()
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
if err := a.Wait(ctx); err != nil {
|
|
t.Error(err)
|
|
}
|
|
if err := b.Wait(ctx); err != nil {
|
|
t.Error(err)
|
|
}
|
|
})
|
|
return a, b
|
|
}
|
|
func handle(t testing.TB, p *Peer, name string, h Handler) {
|
|
t.Helper()
|
|
if err := p.Handle(name, h); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func TestBidirectionalNestedCalls(t *testing.T) {
|
|
a, _ := pair(t, Config{}, func(a, b *Peer) {
|
|
handle(t, a, "double", func(_ context.Context, v json.RawMessage) (any, error) {
|
|
var n int
|
|
_ = json.Unmarshal(v, &n)
|
|
return n * 2, nil
|
|
})
|
|
handle(t, b, "nested", func(ctx context.Context, v json.RawMessage) (any, error) {
|
|
var n int
|
|
_ = json.Unmarshal(v, &n)
|
|
var doubled int
|
|
err := b.Call(ctx, "double", n, &doubled)
|
|
return doubled + 1, err
|
|
})
|
|
})
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < 20; i++ {
|
|
wg.Add(1)
|
|
go func(n int) {
|
|
defer wg.Done()
|
|
var result int
|
|
if err := a.Call(context.Background(), "nested", n, &result); err != nil || result != n*2+1 {
|
|
t.Errorf("result %d, error %v", result, err)
|
|
}
|
|
}(i)
|
|
}
|
|
wg.Wait()
|
|
}
|
|
|
|
func TestCancellationBusyAndCompletionBarrier(t *testing.T) {
|
|
entered := make(chan struct{})
|
|
exit := make(chan struct{})
|
|
release := make(chan struct{})
|
|
a, b := pair(t, Config{MaxConcurrent: 1}, func(a, b *Peer) {
|
|
handle(t, b, "block", func(ctx context.Context, _ json.RawMessage) (any, error) {
|
|
close(entered)
|
|
<-ctx.Done()
|
|
close(exit)
|
|
<-release
|
|
return nil, ctx.Err()
|
|
})
|
|
})
|
|
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
|
defer cancel()
|
|
result := make(chan error, 1)
|
|
go func() { result <- a.Call(ctx, "block", nil, nil) }()
|
|
<-entered
|
|
if err := a.Call(context.Background(), "block", nil, nil); err == nil {
|
|
t.Fatal("capacity not enforced")
|
|
} else {
|
|
var pub *Error
|
|
if !errors.As(err, &pub) || pub.Code != "BUSY" {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
if err := <-result; !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Fatal(err)
|
|
}
|
|
<-exit
|
|
b.Stop()
|
|
wait, c := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
|
defer c()
|
|
if err := b.Wait(wait); !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Fatal("handler not joined", err)
|
|
}
|
|
close(release)
|
|
if err := b.Wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func TestErrorsDoNotExposeImplementation(t *testing.T) {
|
|
a, _ := pair(t, Config{}, func(a, b *Peer) {
|
|
handle(t, b, "private", func(context.Context, json.RawMessage) (any, error) { return nil, errors.New("password=secret") })
|
|
handle(t, b, "panic", func(context.Context, json.RawMessage) (any, error) { panic("token=private") })
|
|
handle(t, b, "public", func(context.Context, json.RawMessage) (any, error) {
|
|
return nil, &Error{"DENIED", "public explanation"}
|
|
})
|
|
})
|
|
for _, name := range []string{"private", "panic", "missing", "public"} {
|
|
err := a.Call(context.Background(), name, nil, nil)
|
|
var pub *Error
|
|
if !errors.As(err, &pub) {
|
|
t.Fatal(err)
|
|
}
|
|
if strings.Contains(pub.Message, "secret") || strings.Contains(pub.Message, "private") {
|
|
t.Fatal("leaked error", pub.Message)
|
|
}
|
|
if name == "public" && pub.Message != "public explanation" {
|
|
t.Fatal(pub)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestReplyLateDuplicateAndProtocolLimits(t *testing.T) {
|
|
p, _ := New(Config{})
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
q, _ := New(Config{})
|
|
_ = p.bind(&memoryTransport{ctx, q})
|
|
defer p.Stop()
|
|
// Late and duplicated replies must not block the reader or create allocations.
|
|
for i := 0; i < 10; i++ {
|
|
if err := p.receive(ws.Message{Type: ws.Text, Data: []byte(`{"v":1,"type":"reply","id":"old","result":null}`)}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
invalid := []string{
|
|
`{"v":2,"type":"event","method":"x","params":null}`,
|
|
`{"v":1,"type":"call","id":"x","id":"y","method":"x","params":null}`,
|
|
`{"v":1,"type":"event","method":"x","params":{"role":1,"role":2}}`,
|
|
`{"v":1,"type":"reply","id":"a","result":null,"error":{"code":"BAD","message":"x"}}`,
|
|
`{"v":1,"type":"cancel","id":"a","params":null}`,
|
|
`{"v":1,"type":"call","id":"x","method":"x","params":null,"extra":true}`,
|
|
`{"v":1,"type":"event","method":"x","params":` + strings.Repeat("[", 34) + "0" + strings.Repeat("]", 34) + `}`,
|
|
}
|
|
for _, data := range invalid {
|
|
if _, err := decode([]byte(data), 1<<20); !errors.Is(err, ErrProtocol) {
|
|
t.Fatalf("accepted %s", data)
|
|
}
|
|
}
|
|
if _, err := decode([]byte(`{"v":1,"type":"event","method":"x","params":null}`), 4); err == nil {
|
|
t.Fatal("size ignored")
|
|
}
|
|
if err := p.receive(ws.Message{Type: ws.Binary, Data: []byte("x")}); !errors.Is(err, ErrProtocol) {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func TestPendingLimitAndDisconnect(t *testing.T) {
|
|
entered := make(chan struct{})
|
|
a, b := pair(t, Config{MaxPending: 1}, func(a, b *Peer) {
|
|
handle(t, b, "wait", func(ctx context.Context, _ json.RawMessage) (any, error) {
|
|
close(entered)
|
|
<-ctx.Done()
|
|
return nil, ctx.Err()
|
|
})
|
|
})
|
|
done := make(chan error, 1)
|
|
go func() { done <- a.Call(context.Background(), "wait", nil, nil) }()
|
|
<-entered
|
|
if err := a.Call(context.Background(), "wait", nil, nil); !errors.Is(err, ErrBusy) {
|
|
t.Fatal(err)
|
|
}
|
|
a.Stop()
|
|
b.Stop()
|
|
if err := <-done; !errors.Is(err, ErrClosed) {
|
|
t.Fatal(err)
|
|
}
|
|
if err := a.bind(&memoryTransport{context.Background(), b}); !errors.Is(err, ErrClosed) {
|
|
t.Fatal("peer rebound", err)
|
|
}
|
|
}
|
|
|
|
func TestSubscriptionsRestoreAndRemoval(t *testing.T) {
|
|
var mu sync.Mutex
|
|
var keys []string
|
|
setup := func(a, b *Peer) {
|
|
handle(t, b, "rpc.subscribe", func(_ context.Context, raw json.RawMessage) (any, error) {
|
|
var sub Subscription
|
|
if err := json.Unmarshal(raw, &sub); err != nil {
|
|
return nil, err
|
|
}
|
|
mu.Lock()
|
|
keys = append(keys, sub.Key)
|
|
mu.Unlock()
|
|
return nil, nil
|
|
})
|
|
handle(t, b, "rpc.unsubscribe", func(context.Context, json.RawMessage) (any, error) { return nil, nil })
|
|
}
|
|
a, _ := pair(t, Config{}, setup)
|
|
s, err := NewSubscriptions(SubscriptionConfig{MaxSubscriptions: 2})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_ = s.Subscribe(context.Background(), nil, "b", map[string]int{"n": 1})
|
|
_ = s.Subscribe(context.Background(), nil, "a", nil)
|
|
if err = s.Subscribe(context.Background(), nil, "c", nil); !errors.Is(err, ErrBusy) {
|
|
t.Fatal(err)
|
|
}
|
|
if n, err := s.Restore(context.Background(), a); n != 2 || err != nil {
|
|
t.Fatal(n, err)
|
|
}
|
|
if err = s.Unsubscribe(context.Background(), a, "a"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
next, _ := pair(t, Config{}, setup)
|
|
if n, err := s.Restore(context.Background(), next); n != 1 || err != nil {
|
|
t.Fatal(n, err)
|
|
}
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
if strings.Join(keys, ",") != "a,b,b" {
|
|
t.Fatal(keys)
|
|
}
|
|
}
|
|
|
|
func TestRealTCPRPCAndShutdown(t *testing.T) {
|
|
opened := make(chan *Peer, 1)
|
|
server, err := ws.NewServer(ws.ServerConfig{Subprotocols: []string{Subprotocol}, RequireSubprotocol: true, Prepare: func(context.Context, *http.Request) (ws.Handlers, error) {
|
|
p, _ := New(Config{})
|
|
handle(t, p, "echo", func(_ context.Context, raw json.RawMessage) (any, error) { return raw, nil })
|
|
opened <- p
|
|
return p.Handlers(), nil
|
|
}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
httpServer := httptest.NewServer(server.Handler(ws.Handlers{}))
|
|
defer httpServer.Close()
|
|
client, err := ws.NewClient(ws.ClientConfig{Subprotocols: []string{Subprotocol}, RequireSubprotocol: true, ProxyMode: "direct"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer func() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
_ = client.Shutdown(ctx)
|
|
_ = server.Shutdown(ctx)
|
|
}()
|
|
p, _ := New(Config{})
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|
defer cancel()
|
|
session, err := client.Dial(ctx, "ws"+strings.TrimPrefix(httpServer.URL, "http"), p.Handlers())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err = p.WaitReady(ctx); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
remote := <-opened
|
|
var result map[string]int
|
|
if err = p.Call(ctx, "echo", map[string]int{"n": 42}, &result); err != nil || result["n"] != 42 {
|
|
t.Fatal(result, err)
|
|
}
|
|
if err = p.Shutdown(ctx); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err = session.Wait(ctx); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err = remote.Wait(ctx); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func BenchmarkPeerCall(b *testing.B) {
|
|
a, _ := pair(b, Config{}, func(a, p *Peer) {
|
|
handle(b, p, "echo", func(_ context.Context, raw json.RawMessage) (any, error) { return raw, nil })
|
|
})
|
|
b.ReportAllocs()
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
var out int
|
|
if err := a.Call(context.Background(), "echo", 42, &out); err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|
|
}
|
|
func BenchmarkProtocolDecode(b *testing.B) {
|
|
data := []byte(`{"v":1,"type":"call","id":"abc","method":"status.get","params":{"detail":true},"timeoutMs":1000}`)
|
|
b.ReportAllocs()
|
|
b.SetBytes(int64(len(data)))
|
|
for i := 0; i < b.N; i++ {
|
|
if _, err := decode(data, 1<<20); err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|
|
}
|
|
func FuzzProtocolDecode(f *testing.F) {
|
|
for _, s := range []string{`{"v":1,"type":"event","method":"x","params":null}`, `{"v":1,"type":"reply","id":"x","result":42}`, `null`, `{"v":1,"v":2}`} {
|
|
f.Add([]byte(s))
|
|
}
|
|
f.Fuzz(func(t *testing.T, data []byte) { _, _ = decode(data, 4096) })
|
|
}
|
|
|
|
type panickingResult struct{}
|
|
|
|
func (panickingResult) MarshalJSON() ([]byte, error) { panic("secret encoder state") }
|
|
func TestResultEncodingPanicAndOversize(t *testing.T) {
|
|
a, _ := pair(t, Config{MaxMessageBytes: 256}, func(a, b *Peer) {
|
|
handle(t, b, "panicResult", func(context.Context, json.RawMessage) (any, error) { return panickingResult{}, nil })
|
|
handle(t, b, "large", func(context.Context, json.RawMessage) (any, error) { return strings.Repeat("x", 1024), nil })
|
|
})
|
|
for _, method := range []string{"panicResult", "large"} {
|
|
err := a.Call(context.Background(), method, nil, nil)
|
|
var e *Error
|
|
if !errors.As(err, &e) {
|
|
t.Fatal(err)
|
|
}
|
|
if method == "panicResult" && e.Code != "INTERNAL" {
|
|
t.Fatal(e)
|
|
}
|
|
if method == "large" && e.Code != "RESULT_TOO_LARGE" {
|
|
t.Fatal(e)
|
|
}
|
|
}
|
|
}
|
|
func TestEventErrorObservation(t *testing.T) {
|
|
observed := make(chan error, 1)
|
|
a, _ := pair(t, Config{OnEventError: func(_ context.Context, err error) { observed <- err; panic("observer panic") }}, func(a, b *Peer) {
|
|
handle(t, b, "fail", func(context.Context, json.RawMessage) (any, error) { return nil, errors.New("local failure") })
|
|
})
|
|
if err := a.Notify(context.Background(), "fail", nil); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
select {
|
|
case err := <-observed:
|
|
if err.Error() != "local failure" {
|
|
t.Fatal(err)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatal("event error not observed")
|
|
}
|
|
}
|
|
func TestDuplicateIncomingIDRejected(t *testing.T) {
|
|
entered := make(chan struct{})
|
|
a, b := pair(t, Config{}, func(a, b *Peer) {
|
|
handle(t, b, "wait", func(ctx context.Context, _ json.RawMessage) (any, error) {
|
|
close(entered)
|
|
<-ctx.Done()
|
|
return nil, ctx.Err()
|
|
})
|
|
})
|
|
raw := ws.Message{Type: ws.Text, Data: []byte(`{"v":1,"type":"call","id":"same","method":"wait","params":null}`)}
|
|
if err := b.receive(raw); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
<-entered
|
|
if err := b.receive(raw); !errors.Is(err, ErrProtocol) {
|
|
t.Fatal("duplicate active ID accepted", err)
|
|
}
|
|
a.Stop()
|
|
b.Stop()
|
|
}
|