internal/mtprotoedge/admission_test.go
Rama main · réplica de lectura del repositorio original.
package mtprotoedge
import (
"context"
"errors"
"net"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto/codec"
"telesrv/internal/store"
"telesrv/internal/store/memory"
)
func TestAdmissionConnectionLimitsAndIdempotentRelease(t *testing.T) {
a := newAdmissionController(2, 1, 1)
ip1a := &net.TCPAddr{IP: net.ParseIP("203.0.113.1"), Port: 1000}
ip1b := &net.TCPAddr{IP: net.ParseIP("203.0.113.1"), Port: 1001}
ip2 := &net.TCPAddr{IP: net.ParseIP("203.0.113.2"), Port: 1000}
ip3 := &net.TCPAddr{IP: net.ParseIP("203.0.113.3"), Port: 1000}
release1, ok := a.acquireConnection(ip1a)
if !ok {
t.Fatal("first connection rejected")
}
if _, ok := a.acquireConnection(ip1b); ok {
t.Fatal("second connection from same IP bypassed per-IP cap")
}
release2, ok := a.acquireConnection(ip2)
if !ok {
t.Fatal("second IP connection rejected below global cap")
}
if _, ok := a.acquireConnection(ip3); ok {
t.Fatal("third connection bypassed global cap")
}
release1()
release1() // idempotent release must not drive the counter negative.
releaseAgain, ok := a.acquireConnection(ip1b)
if !ok {
t.Fatal("released per-IP/global slot was not reusable")
}
releaseAgain()
release2()
a.mu.Lock()
defer a.mu.Unlock()
if a.connections != 0 || len(a.byIP) != 0 {
t.Fatalf("admission counters after release = %d/%v, want 0/empty", a.connections, a.byIP)
}
}
func TestAdmissionHandshakeLimitAndRelease(t *testing.T) {
a := newAdmissionController(-1, -1, 1)
release, ok := a.tryAcquireHandshake()
if !ok {
t.Fatal("first handshake rejected")
}
if _, ok := a.tryAcquireHandshake(); ok {
t.Fatal("second handshake bypassed semaphore")
}
release()
release() // idempotent
release2, ok := a.tryAcquireHandshake()
if !ok {
t.Fatal("released handshake slot was not reusable")
}
release2()
}
type oneConnListener struct {
conn net.Conn
once sync.Once
}
func (l *oneConnListener) Accept() (net.Conn, error) {
var conn net.Conn
l.once.Do(func() {
conn = l.conn
})
if conn == nil {
return nil, net.ErrClosed
}
return conn, nil
}
func (l *oneConnListener) Close() error { return l.conn.Close() }
func (l *oneConnListener) Addr() net.Addr { return l.conn.LocalAddr() }
func TestAdmissionListenerTracksUntilPhysicalClose(t *testing.T) {
serverSide, clientSide := net.Pipe()
defer clientSide.Close()
a := newAdmissionController(1, 1, 1)
ln := a.wrapListener(&oneConnListener{conn: serverSide})
conn, err := ln.Accept()
if err != nil {
t.Fatalf("Accept: %v", err)
}
a.mu.Lock()
active := a.connections
a.mu.Unlock()
if active != 1 {
t.Fatalf("active after Accept = %d, want 1", active)
}
_ = conn.Close()
_ = conn.Close()
a.mu.Lock()
active = a.connections
a.mu.Unlock()
if active != 0 {
t.Fatalf("active after physical Close = %d, want 0", active)
}
}
type temporaryAcceptTestError struct{}
func (temporaryAcceptTestError) Error() string { return "temporary accept failure" }
func (temporaryAcceptTestError) Timeout() bool { return false }
func (temporaryAcceptTestError) Temporary() bool { return true }
type temporaryThenConnListener struct {
conn net.Conn
closed chan struct{}
closeOnce sync.Once
calls atomic.Int32
}
type connThenErrorListener struct {
conn net.Conn
err error
closeOnce sync.Once
calls atomic.Int32
}
func (l *connThenErrorListener) Accept() (net.Conn, error) {
if l.calls.Add(1) == 1 {
return l.conn, nil
}
return nil, l.err
}
func (l *connThenErrorListener) Close() error {
var err error
l.closeOnce.Do(func() {
if l.conn != nil {
err = l.conn.Close()
}
})
return err
}
func (l *connThenErrorListener) Addr() net.Addr {
if l.conn != nil {
return l.conn.LocalAddr()
}
return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345}
}
type fixedErrorListener struct {
err error
addr net.Addr
}
func (l *fixedErrorListener) Accept() (net.Conn, error) { return nil, l.err }
func (*fixedErrorListener) Close() error { return nil }
func (l *fixedErrorListener) Addr() net.Addr { return l.addr }
func (l *temporaryThenConnListener) Accept() (net.Conn, error) {
call := l.calls.Add(1)
if call == 1 {
return nil, temporaryAcceptTestError{}
}
if call == 2 {
return l.conn, nil
}
<-l.closed
return nil, net.ErrClosed
}
func (l *temporaryThenConnListener) Close() error {
l.closeOnce.Do(func() {
close(l.closed)
_ = l.conn.Close()
})
return nil
}
func (l *temporaryThenConnListener) Addr() net.Addr {
return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345}
}
func TestAcceptLoopRetriesTemporaryError(t *testing.T) {
serverSide, clientSide := net.Pipe()
defer clientSide.Close()
ln := &temporaryThenConnListener{conn: serverSide, closed: make(chan struct{})}
srv := New(Options{HandshakeIdleTimeout: 100 * time.Millisecond})
ctx, cancel := context.WithCancel(context.Background())
done := make(chan error, 1)
go func() { done <- srv.acceptLoop(ctx, ln, false) }()
deadline := time.Now().Add(time.Second)
for ln.calls.Load() < 3 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if ln.calls.Load() < 3 {
cancel()
<-done
t.Fatalf("accept calls = %d, want temporary retry then next accept", ln.calls.Load())
}
cancel()
select {
case err := <-done:
if err != nil {
t.Fatalf("acceptLoop after temporary error: %v", err)
}
case <-time.After(time.Second):
t.Fatal("acceptLoop did not stop after cancel")
}
}
func TestAcceptLoopPermanentErrorCancelsAcceptedConnectionsBeforeWait(t *testing.T) {
serverSide, clientSide := net.Pipe()
defer clientSide.Close()
wantErr := errors.New("permanent accept failure")
ln := &connThenErrorListener{conn: serverSide, err: wantErr}
srv := New(Options{HandshakeIdleTimeout: time.Hour})
done := make(chan error, 1)
go func() {
done <- srv.acceptLoop(context.Background(), ln, false)
}()
select {
case err := <-done:
if !errors.Is(err, wantErr) {
t.Fatalf("acceptLoop error = %v, want %v", err, wantErr)
}
case <-time.After(time.Second):
t.Fatal("acceptLoop waited for an accepted connection before canceling it")
}
_ = clientSide.SetReadDeadline(time.Now().Add(time.Second))
var one [1]byte
if _, err := clientSide.Read(one[:]); err == nil {
t.Fatal("accepted connection remained open after permanent accept failure")
}
}
func TestServeMixedStopsAllComponentsWhenOneReturnsCleanly(t *testing.T) {
srv := New(Options{WebSocket: true})
ln := &fixedErrorListener{
err: net.ErrClosed,
addr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 2398},
}
done := make(chan error, 1)
go func() {
done <- srv.serveMixed(context.Background(), ln)
}()
select {
case err := <-done:
if err != nil {
t.Fatalf("serveMixed error = %v, want nil closed-listener shutdown", err)
}
case <-time.After(time.Second):
t.Fatal("serveMixed did not stop remaining components after one clean exit")
}
}
type countingAuthKeyStore struct {
store.AuthKeyStore
gets atomic.Int32
revalidates atomic.Int32
}
func (s *countingAuthKeyStore) Get(ctx context.Context, id [8]byte) (store.AuthKeyData, bool, error) {
s.gets.Add(1)
return s.AuthKeyStore.Get(ctx, id)
}
func (s *countingAuthKeyStore) Revalidate(ctx context.Context, id [8]byte) (store.AuthKeyData, bool, error) {
s.revalidates.Add(1)
return s.AuthKeyStore.Revalidate(ctx, id)
}
func TestUnknownAuthKeyRespondsOnceThenCloses(t *testing.T) {
keys := &countingAuthKeyStore{AuthKeyStore: memory.NewAuthKeyStore()}
addr, _, _ := startTestServer(t, Options{AuthKeys: keys})
conn := dialTransportOnly(t, addr)
var request bin.Buffer
request.PutLong(0x0102030405060708)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := conn.Send(ctx, &request); err != nil {
t.Fatalf("send unknown auth key: %v", err)
}
var response bin.Buffer
err := conn.Recv(ctx, &response)
var protocolErr *codec.ProtocolErr
if !errors.As(err, &protocolErr) || protocolErr.Code != codec.CodeAuthKeyNotFound {
t.Fatalf("first recv err = %T %v, want protocol -404", err, err)
}
if got := keys.gets.Load(); got != 1 {
t.Fatalf("AuthKeyStore.Get calls = %d, want 1", got)
}
response.Reset()
err = conn.Recv(ctx, &response)
if err == nil {
t.Fatal("connection remained readable after terminal -404")
}
if got := keys.gets.Load(); got != 1 {
t.Fatalf("AuthKeyStore.Get calls after close = %d, want 1", got)
}
}