278 lines
7.1 KiB
Go
278 lines
7.1 KiB
Go
package tcp
|
|
|
|
import (
|
|
"context"
|
|
"encoding/binary"
|
|
"io"
|
|
"net"
|
|
"testing"
|
|
"time"
|
|
|
|
"spark-server/internal/infrastructure/id"
|
|
"spark-server/internal/infrastructure/logger"
|
|
authlogic "spark-server/internal/module/auth/logic"
|
|
"spark-server/internal/transport/tcp/protocol"
|
|
"spark-server/internal/transport/tcp/session"
|
|
)
|
|
|
|
func TestHeartbeatWithoutAuth(t *testing.T) {
|
|
srv := startTCP(t, &Options{})
|
|
c := dial(t, srv.addr)
|
|
defer c.Close()
|
|
writeFrame(t, c, protocol.Frame{Version: protocol.CurrentVersion, Command: protocol.CmdHeartbeat, SequenceID: 3})
|
|
got := readFrame(t, c)
|
|
if got.Command != protocol.CmdHeartbeat || got.SequenceID != 3 {
|
|
t.Fatalf("heartbeat reply %+v", got)
|
|
}
|
|
}
|
|
|
|
func TestAuthGateAndGenerateMsgID(t *testing.T) {
|
|
ids, err := id.New(1)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
srv := startTCP(t, &Options{
|
|
Auth: authlogic.Static{Token: "secret", UserID: 42, DeviceID: 7},
|
|
IDs: ids,
|
|
})
|
|
c := dial(t, srv.addr)
|
|
defer c.Close()
|
|
|
|
writeFrame(t, c, protocol.Frame{
|
|
Version: protocol.CurrentVersion, Command: protocol.CmdGenerateMsgID, RequestID: 9, SequenceID: 1,
|
|
})
|
|
denied := readFrame(t, c)
|
|
eb, err := protocol.DecodeError(denied.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if eb.Code != protocol.CodeNotAuth {
|
|
t.Fatalf("want not-auth, got %+v", eb)
|
|
}
|
|
|
|
writeFrame(t, c, protocol.Frame{
|
|
Version: protocol.CurrentVersion, Command: protocol.CmdAuth, RequestID: 2, SequenceID: 2,
|
|
Body: protocol.EncodeAuthRequest(protocol.AuthRequest{Token: "secret", DeviceID: 7}),
|
|
})
|
|
authed := readFrame(t, c)
|
|
env, err := protocol.DecodeEnvelope(authed.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if authed.Command != protocol.CmdAuth || env.Code != protocol.CodeOK {
|
|
t.Fatalf("auth reply %+v env=%+v", authed, env)
|
|
}
|
|
|
|
writeFrame(t, c, protocol.Frame{
|
|
Version: protocol.CurrentVersion, Command: protocol.CmdAuth, RequestID: 8, SequenceID: 8,
|
|
Body: protocol.EncodeAuthRequest(protocol.AuthRequest{Token: "secret", DeviceID: 99, Language: "en"}),
|
|
})
|
|
mismatch := readFrame(t, c)
|
|
eb2, err := protocol.DecodeError(mismatch.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if eb2.Code != protocol.CodeUnauthorized || eb2.Message != "unauthorized" {
|
|
t.Fatalf("want device mismatch unauthorized, got %+v", eb2)
|
|
}
|
|
|
|
writeFrame(t, c, protocol.Frame{
|
|
Version: protocol.CurrentVersion, Command: protocol.CmdGenerateMsgID, RequestID: 3, SequenceID: 3,
|
|
})
|
|
msgid := readFrame(t, c)
|
|
if msgid.Command != protocol.CmdGenerateMsgID || len(msgid.Body) == 0 {
|
|
t.Fatalf("msgid reply %+v", msgid)
|
|
}
|
|
}
|
|
|
|
func TestAuthJWTSameService(t *testing.T) {
|
|
ids, err := id.New(1)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
j, err := authlogic.NewJWT(&authlogic.Options{Secret: "test-secret-test-secret", Issuer: "spark-im", TTL: time.Hour})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
tok, err := j.Issue(context.Background(), 42, 7, 1)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
srv := startTCP(t, &Options{Auth: j, IDs: ids})
|
|
c := dial(t, srv.addr)
|
|
defer c.Close()
|
|
writeFrame(t, c, protocol.Frame{
|
|
Version: protocol.CurrentVersion, Command: protocol.CmdAuth, SequenceID: 1,
|
|
Body: protocol.EncodeAuthRequest(protocol.AuthRequest{Token: tok.AccessToken, DeviceID: 7}),
|
|
})
|
|
authed := readFrame(t, c)
|
|
env, err := protocol.DecodeEnvelope(authed.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if authed.Command != protocol.CmdAuth || env.Code != protocol.CodeOK {
|
|
t.Fatalf("jwt auth %+v env=%+v", authed, env)
|
|
}
|
|
writeFrame(t, c, protocol.Frame{
|
|
Version: protocol.CurrentVersion, Command: protocol.CmdAuth, SequenceID: 2,
|
|
Body: protocol.EncodeAuthRequest(protocol.AuthRequest{Token: tok.AccessToken, DeviceID: 99}),
|
|
})
|
|
mismatch := readFrame(t, c)
|
|
eb, err := protocol.DecodeError(mismatch.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if eb.Code != protocol.CodeUnauthorized {
|
|
t.Fatalf("jwt device mismatch %+v", eb)
|
|
}
|
|
}
|
|
|
|
func TestProxyProtocolThenHeartbeat(t *testing.T) {
|
|
ids, err := id.New(1)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
srv := startTCP(t, &Options{
|
|
ProxyProtocol: true,
|
|
Auth: authlogic.Static{Token: "secret", UserID: 42, DeviceID: 7},
|
|
IDs: ids,
|
|
})
|
|
c := dial(t, srv.addr)
|
|
defer c.Close()
|
|
if _, err := c.Write([]byte("PROXY TCP4 203.0.113.10 192.0.2.1 443 9000\r\n")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
writeFrame(t, c, protocol.Frame{
|
|
Version: protocol.CurrentVersion, Command: protocol.CmdAuth, SequenceID: 1,
|
|
Body: protocol.EncodeAuthRequest(protocol.AuthRequest{Token: "secret", DeviceID: 7}),
|
|
})
|
|
authed := readFrame(t, c)
|
|
if authed.Command != protocol.CmdAuth {
|
|
t.Fatalf("auth after proxy %+v", authed)
|
|
}
|
|
deadline := time.Now().Add(2 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
sessions := srv.sessions.GetUserSessions(42)
|
|
if len(sessions) == 1 && sessions[0].RealIP() == "203.0.113.10:443" {
|
|
return
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
t.Fatalf("real ip not set, sessions=%v", dumpRealIP(srv.sessions.GetUserSessions(42)))
|
|
}
|
|
|
|
func dumpRealIP(ss []session.Session) []string {
|
|
out := make([]string, 0, len(ss))
|
|
for _, s := range ss {
|
|
out = append(out, s.RealIP())
|
|
}
|
|
return out
|
|
}
|
|
|
|
func startTCP(t *testing.T, o *Options) *Server {
|
|
t.Helper()
|
|
if o == nil {
|
|
o = &Options{}
|
|
}
|
|
if o.Addr == "" {
|
|
o.Addr = freeAddr(t)
|
|
}
|
|
if o.IdleTimeout == 0 {
|
|
o.IdleTimeout = time.Hour
|
|
}
|
|
o.KeepAlive = 30 * time.Second
|
|
log := logger.Nop()
|
|
pool := NewTaskPool(4, 32)
|
|
t.Cleanup(pool.Close)
|
|
s := NewServer(log, pool, session.NewManager(), o)
|
|
errCh := make(chan error, 1)
|
|
go func() { errCh <- s.Start() }()
|
|
select {
|
|
case <-s.ready:
|
|
case err := <-errCh:
|
|
t.Fatalf("start: %v", err)
|
|
case <-time.After(3 * time.Second):
|
|
t.Fatal("boot timeout")
|
|
}
|
|
t.Cleanup(func() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|
defer cancel()
|
|
_ = s.Stop(ctx)
|
|
select {
|
|
case <-errCh:
|
|
case <-time.After(3 * time.Second):
|
|
}
|
|
})
|
|
return s
|
|
}
|
|
|
|
func freeAddr(t *testing.T) string {
|
|
t.Helper()
|
|
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
addr := ln.Addr().String()
|
|
_ = ln.Close()
|
|
return addr
|
|
}
|
|
|
|
func dial(t *testing.T, addr string) net.Conn {
|
|
t.Helper()
|
|
var last error
|
|
deadline := time.Now().Add(2 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
c, err := net.DialTimeout("tcp", addr, 200*time.Millisecond)
|
|
if err == nil {
|
|
_ = c.SetDeadline(time.Now().Add(3 * time.Second))
|
|
return c
|
|
}
|
|
last = err
|
|
time.Sleep(20 * time.Millisecond)
|
|
}
|
|
t.Fatalf("dial %s: %v", addr, last)
|
|
return nil
|
|
}
|
|
|
|
func writeFrame(t *testing.T, c net.Conn, f protocol.Frame) {
|
|
t.Helper()
|
|
raw, err := protocol.Encode(f)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := c.Write(raw); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func readFrame(t *testing.T, c net.Conn) protocol.Frame {
|
|
t.Helper()
|
|
header := make([]byte, protocol.HeaderSize)
|
|
if _, err := io.ReadFull(c, header); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
bodyLen := int(binary.BigEndian.Uint32(header[19:23]))
|
|
raw := make([]byte, protocol.HeaderSize+bodyLen)
|
|
copy(raw, header)
|
|
if bodyLen > 0 {
|
|
if _, err := io.ReadFull(c, raw[protocol.HeaderSize:]); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
f, err := protocol.Decode(raw, 65536)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return f
|
|
}
|
|
|
|
func TestLimiterAllow(t *testing.T) {
|
|
l := newLimiter(1)
|
|
if !l.Allow() {
|
|
t.Fatal("first token")
|
|
}
|
|
if l.Allow() {
|
|
t.Fatal("burst exhausted")
|
|
}
|
|
}
|