package main import ( "context" "errors" "io" "net" "net/http" "net/http/httptest" "strings" "sync" "sync/atomic" "testing" "time" "golang.org/x/time/rate" ) func TestNativeClientMaxConnectionsAndBatchedActiveDelta(t *testing.T) { oldStore := statsStore statsStore = &Store{} defer func() { statsStore = oldStore }() const uuid = "11111111-1111-1111-1111-111111111111" m := &XrayManager{ nativeQuotaByUUID: map[string]*xrayNativeQuotaState{ uuid: {maxConns: 1, generation: 1}, }, } state := m.nativeQuotaState(uuid) release, acquiredState, ok := m.acquireNativeClientConnection(uuid, "user@example") if !ok || release == nil { t.Fatal("first native connection was rejected") } if acquiredState != state { t.Fatal("connection lease did not retain the authenticated policy state") } if _, _, ok := m.acquireNativeClientConnection(uuid, "user@example"); ok { t.Fatal("connection above max_conns was accepted") } m.nativeDBMu.Lock() pending := m.nativeActivePending[uuid] m.nativeDBMu.Unlock() if pending.Delta != 1 || !pending.Connected || pending.State != state { t.Fatalf("connect was not queued for batch persistence: %+v", pending) } release() release() // idempotent release must not underflow counters. m.nativeDBMu.Lock() pending = m.nativeActivePending[uuid] m.nativeDBMu.Unlock() if pending.Delta != 0 || !pending.Connected { t.Fatalf("connect/disconnect batch should net to zero and retain last-active: %+v", pending) } state.mu.Lock() active := state.activeConns state.mu.Unlock() if active != 0 { t.Fatalf("active connection count = %d, want 0", active) } release2, _, ok := m.acquireNativeClientConnection(uuid, "user@example") if !ok { t.Fatal("slot was not reusable after release") } release2() } func TestNativeOnlineUsersAreKeyedByUUID(t *testing.T) { oldStore := statsStore statsStore = nil defer func() { statsStore = oldStore }() m := &XrayManager{} m.recordNativeConnect("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", "shared@example", nil) m.recordNativeConnect("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb", "shared@example", nil) if got := m.CountOnlineUsers(); got != 2 { t.Fatalf("online UUID count = %d, want 2 for two UUIDs sharing one email", got) } m.statsMu.RLock() _, first := m.statsByEmail["aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"] _, second := m.statsByEmail["bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"] m.statsMu.RUnlock() if !first || !second { t.Fatal("native runtime stats were not stored under canonical UUID keys") } } func TestNativeExpiryRejectsAndDisconnectsClient(t *testing.T) { const uuid = "cccccccc-cccc-cccc-cccc-cccccccccccc" state := &xrayNativeQuotaState{hasExpiry: true, expiresAt: time.Now().Add(time.Hour), generation: 1} m := &XrayManager{nativeQuotaByUUID: map[string]*xrayNativeQuotaState{uuid: state}} closer := &closeTrackingReader{} release, _, ok := m.acquireNativeClientConnection(uuid, "expiry@example", closer) if !ok { t.Fatal("unexpired client was rejected") } m.disconnectNativeClient(uuid) if !closer.closed.Load() { t.Fatal("active native client was not closed during revocation") } release() state.mu.Lock() state.expiresAt = time.Now().Add(-time.Second) state.mu.Unlock() if reason := m.nativeClientAccessDenied(uuid); reason != "expired" { t.Fatalf("expired client denial = %q, want expired", reason) } if _, _, ok := m.acquireNativeClientConnection(uuid, "expiry@example"); ok { t.Fatal("expired client acquired a new connection") } } func TestQuotaAndExpiryValidationRejectsUnsafeInput(t *testing.T) { if err := validateQuotaConfig(1, "typo", 1); err == nil { t.Fatal("unknown quota action was accepted") } if err := validateQuotaConfig(1, quotaActionBlock, -1); err == nil { t.Fatal("negative throttle setting was accepted") } if _, err := parseOptionalXrayExpiry("not-a-date"); err == nil { t.Fatal("invalid Xray expiry was accepted") } if exp, err := parseOptionalXrayExpiry(""); err != nil || exp != nil { t.Fatalf("empty Xray expiry = (%v, %v), want nil, nil", exp, err) } } func TestRemoveNativeQuotaPolicyPrunesPendingMaps(t *testing.T) { m := &XrayManager{ nativeQuotaByUUID: map[string]*xrayNativeQuotaState{ "gone": {generation: 1}, }, nativeTrafficPending: map[string]xrayPendingTraffic{ "gone": {Uplink: 10}, }, nativeActivePending: map[string]xrayPendingActive{ "gone": {Delta: 1}, }, } m.removeNativeQuotaPolicy("gone") if m.nativeQuotaState("gone") != nil { t.Fatal("quota policy was not removed") } m.nativeDBMu.Lock() _, trafficExists := m.nativeTrafficPending["gone"] _, activeExists := m.nativeActivePending["gone"] m.nativeDBMu.Unlock() if trafficExists || activeExists { t.Fatal("deleted UUID remained in a pending persistence map") } } func TestNativeCounterIsBoundedAndReleaseIsIdempotent(t *testing.T) { var active atomicInt64ForTest release1, ok := acquireNativeCounter(&active.Int64, 2) if !ok { t.Fatal("first slot rejected") } release2, ok := acquireNativeCounter(&active.Int64, 2) if !ok { t.Fatal("second slot rejected") } if _, ok := acquireNativeCounter(&active.Int64, 2); ok { t.Fatal("slot above limit accepted") } release1() release1() if got := active.Load(); got != 1 { t.Fatalf("active after double release = %d, want 1", got) } release2() if got := active.Load(); got != 0 { t.Fatalf("active after releases = %d, want 0", got) } } // Embedding keeps the test declaration readable while still passing the exact // atomic.Int64 type required by acquireNativeCounter. type atomicInt64ForTest struct{ Int64 atomic.Int64 } func (a *atomicInt64ForTest) Load() int64 { return a.Int64.Load() } func TestTrackedNativeConnectionsAreClosedOnShutdown(t *testing.T) { oldAccepting := nativeTransportAccepting.Load() defer nativeTransportAccepting.Store(oldAccepting) beginNativeTransportAccepting() before := nativeTransportConnections.Load() serverSide, clientSide := net.Pipe() defer clientSide.Close() wrapped, ok := wrapTrackedNativeTransportConn(serverSide) if !ok { t.Fatal("tracked connection was rejected") } if got := nativeTransportConnections.Load(); got != before+1 { t.Fatalf("transport count = %d, want %d", got, before+1) } stopNativeTransportAccepting() closeAllNativeTransportConnections() _ = clientSide.SetReadDeadline(time.Now().Add(time.Second)) if _, err := clientSide.Read(make([]byte, 1)); err == nil { t.Fatal("peer remained open after native shutdown") } if err := wrapped.Close(); err != nil && !errors.Is(err, net.ErrClosed) { t.Fatalf("second close returned unexpected error: %v", err) } if got := nativeTransportConnections.Load(); got != before { t.Fatalf("transport count after shutdown = %d, want %d", got, before) } } func TestCloseAllXHTTPSessionsReleasesGlobalSlots(t *testing.T) { oldLimit := nativeTuneXHTTPMaxSessions.Load() nativeTuneXHTTPMaxSessions.Store(8) defer nativeTuneXHTTPMaxSessions.Store(oldLimit) before := nativeXHTTPSessions.Load() ib := &nativeInbound{xhttpMaxBufferedPosts: 2} for _, id := range []string{"one", "two"} { if sess := ib.upsertXHTTPSession(httptest.NewRecorder(), id); sess == nil { t.Fatalf("session %q was rejected", id) } } if got := nativeXHTTPSessions.Load(); got != before+2 { t.Fatalf("global XHTTP sessions = %d, want %d", got, before+2) } ib.closeAllXHTTPSessions() if got := nativeXHTTPSessions.Load(); got != before { t.Fatalf("global XHTTP sessions after close = %d, want %d", got, before) } ib.xhttpMu.Lock() remaining := len(ib.xhttpSessions) ib.xhttpMu.Unlock() if remaining != 0 { t.Fatalf("inbound retained %d XHTTP sessions", remaining) } } func TestNegativeXHTTPSessionLimitMeansUnlimited(t *testing.T) { old := nativeTuneXHTTPMaxSessions.Load() nativeTuneXHTTPMaxSessions.Store(0) defer nativeTuneXHTTPMaxSessions.Store(old) if got := (&nativeInbound{}).xhttpMaxActiveSessions(); got != 0 { t.Fatalf("unlimited XHTTP session limit normalized to %d", got) } } func TestNativeProtocolGuardsRemainFinite(t *testing.T) { oldRequests := nativeTuneMaxXHTTPRequests.Load() defer nativeTuneMaxXHTTPRequests.Store(oldRequests) // Zero is the internal representation of an explicitly disabled application // request counter. HTTP/2 must still retain a finite per-connection guard. nativeTuneMaxXHTTPRequests.Store(0) if got := nativeHTTP2MaxConcurrentStreams(); got != defaultNativeHTTP2MaxStreams { t.Fatalf("HTTP/2 stream guard = %d, want %d", got, defaultNativeHTTP2MaxStreams) } nativeTuneMaxXHTTPRequests.Store(32) if got := nativeHTTP2MaxConcurrentStreams(); got != 32 { t.Fatalf("HTTP/2 stream guard did not honor lower request cap: %d", got) } if got := nativeMuxMaxSessionLimit(); got != 64 { t.Fatalf("per-transport Mux session guard = %d, want 64", got) } } func TestXHTTPHandlerDoesNotApplyWebRequestCeiling(t *testing.T) { oldLimit := nativeTuneMaxXHTTPRequests.Load() oldActive := nativeXHTTPRequests.Load() nativeTuneMaxXHTTPRequests.Store(1) nativeXHTTPRequests.Store(1) defer func() { nativeTuneMaxXHTTPRequests.Store(oldLimit) nativeXHTTPRequests.Store(oldActive) }() ib := &nativeInbound{transport: "xhttp", path: "/"} req := httptest.NewRequest(http.MethodOptions, "/", nil) rec := httptest.NewRecorder() ib.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("XHTTP OPTIONS at legacy request ceiling = %d, want 200", rec.Code) } } func TestLegacyXHTTPTuningMigratesToVPNDefaults(t *testing.T) { got := normalizeNativeXrayTuning(&XrayNativeTuning{ MuxGlobalSessions: 8192, MaxConcurrentConnections: 4096, MaxConcurrentXHTTPRequests: 8192, XHTTPMaxSessions: 4096, }) if got.MuxGlobalSessions != defaultNativeMuxGlobalSessions || got.MaxConcurrentConnections != defaultNativeMaxConnections || got.MaxConcurrentXHTTPRequests != defaultNativeMaxXHTTPRequests || got.XHTTPMaxSessions != defaultNativeXHTTPMaxSessions { t.Fatalf("legacy tuning was not migrated: %+v", got) } } func TestXHTTPMetadataLengthIsBoundedBeforeSessionAllocation(t *testing.T) { ib := &nativeInbound{transport: "xhttp", path: "/"} req := httptest.NewRequest("GET", "/"+strings.Repeat("a", nativeXHTTPMaxSessionIDBytes+1), nil) rec := httptest.NewRecorder() ib.ServeHTTP(rec, req) if rec.Code != 400 { t.Fatalf("oversized XHTTP session id status = %d, want 400", rec.Code) } ib.xhttpMu.Lock() sessions := len(ib.xhttpSessions) ib.xhttpMu.Unlock() if sessions != 0 { t.Fatalf("oversized metadata allocated %d sessions", sessions) } } func TestXHTTPUploadMemoryIsReleasedOnReadAndClose(t *testing.T) { before := nativeXHTTPBufferedBytes.Load() q := newNativeXHTTPUploadQueue(4, 8) lease, ok := acquireNativeXHTTPMemory(8) if !ok { t.Fatal("failed to reserve XHTTP test memory") } lease.shrink(4) if err := q.push(context.Background(), nativeXHTTPPacket{Payload: []byte("test"), Seq: 0}, lease); err != nil { lease.release() t.Fatalf("queue push failed: %v", err) } lease.release() // transferred leases are a no-op for the producer. if got := nativeXHTTPBufferedBytes.Load(); got != before+4 { t.Fatalf("buffered bytes after push = %d, want %d", got, before+4) } buf := make([]byte, 4) if n, err := q.Read(buf); err != nil || n != 4 || string(buf) != "test" { t.Fatalf("queue read = (%d, %v, %q), want (4, nil, test)", n, err, string(buf)) } if got := nativeXHTTPBufferedBytes.Load(); got != before { t.Fatalf("buffered bytes after read = %d, want %d", got, before) } lease, ok = acquireNativeXHTTPMemory(3) if !ok { t.Fatal("failed to reserve second XHTTP test memory") } if err := q.push(context.Background(), nativeXHTTPPacket{Payload: []byte("xyz"), Seq: 2}, lease); err != nil { lease.release() t.Fatalf("second queue push failed: %v", err) } lease.release() q.close() if got := nativeXHTTPBufferedBytes.Load(); got != before { t.Fatalf("buffered bytes after close = %d, want %d", got, before) } } func TestXHTTPUploadQueueEnforcesPerSessionByteBudget(t *testing.T) { before := nativeXHTTPBufferedBytes.Load() q := newNativeXHTTPUploadQueue(4, 4) defer q.close() lease, ok := acquireNativeXHTTPMemory(5) if !ok { t.Fatal("failed to reserve XHTTP test memory") } defer lease.release() err := q.push(context.Background(), nativeXHTTPPacket{Payload: make([]byte, 5)}, lease) if !errors.Is(err, errNativeXHTTPUploadBufferFull) { t.Fatalf("oversized queue push error = %v, want buffer limit", err) } lease.release() if got := nativeXHTTPBufferedBytes.Load(); got != before { t.Fatalf("rejected payload retained %d bytes, baseline %d", got, before) } } func TestXHTTPUploadQueueBackpressuresInsteadOfRejectingBurst(t *testing.T) { before := nativeXHTTPBufferedBytes.Load() q := newNativeXHTTPUploadQueue(2, 4) defer q.close() first, ok := acquireNativeXHTTPMemory(4) if !ok { t.Fatal("failed to reserve first XHTTP payload") } if err := q.push(context.Background(), nativeXHTTPPacket{Payload: []byte("one!"), Seq: 0}, first); err != nil { first.release() t.Fatalf("first queue push failed: %v", err) } first.release() second, ok := acquireNativeXHTTPMemory(4) if !ok { t.Fatal("failed to reserve second XHTTP payload") } done := make(chan error, 1) go func() { done <- q.push(context.Background(), nativeXHTTPPacket{Payload: []byte("two!"), Seq: 1}, second) }() select { case err := <-done: second.release() t.Fatalf("second burst packet did not backpressure: %v", err) case <-time.After(25 * time.Millisecond): } buf := make([]byte, 4) if n, err := q.Read(buf); err != nil || n != 4 || string(buf) != "one!" { second.release() t.Fatalf("first queue read = (%d, %v, %q)", n, err, string(buf)) } select { case err := <-done: if err != nil { second.release() t.Fatalf("backpressured packet failed after space released: %v", err) } second.release() case <-time.After(time.Second): second.release() t.Fatal("backpressured packet did not resume") } q.close() if got := nativeXHTTPBufferedBytes.Load(); got != before { t.Fatalf("backpressure test leaked %d buffered bytes (baseline %d)", got, before) } } func TestXHTTPBodyReservationUsesActualContentLength(t *testing.T) { ib := &nativeInbound{xhttpMaxEachPostBytes: 1_000_000} req := httptest.NewRequest(http.MethodPost, "/session/0", strings.NewReader("small")) if got := ib.xhttpUploadReservationBytes(req); got != 5 { t.Fatalf("body reservation = %d, want actual payload length 5", got) } req.ContentLength = -1 if got := ib.xhttpUploadReservationBytes(req); got != 1_000_000 { t.Fatalf("chunked body reservation = %d, want configured maximum", got) } } func TestNativeQuotaResetWaitsForInFlightTraffic(t *testing.T) { const uuid = "22222222-2222-2222-2222-222222222222" state := &xrayNativeQuotaState{usedBytes: 123, generation: 1} m := &XrayManager{nativeQuotaByUUID: map[string]*xrayNativeQuotaState{uuid: state}} state.trafficMu.RLock() done := make(chan struct{}) go func() { m.resetNativeQuotaUsage(uuid) close(done) }() select { case <-done: state.trafficMu.RUnlock() t.Fatal("traffic reset crossed an in-flight writer boundary") case <-time.After(25 * time.Millisecond): } state.trafficMu.RUnlock() select { case <-done: case <-time.After(time.Second): t.Fatal("traffic reset did not complete after writer released") } state.mu.Lock() used, generation := state.usedBytes, state.generation state.mu.Unlock() if used != 0 || generation != 2 { t.Fatalf("reset state = used %d generation %d, want 0/2", used, generation) } } func TestNativeRateWaitCanBeCanceled(t *testing.T) { lim := rate.NewLimiter(1, 1) if !lim.AllowN(time.Now(), 1) { t.Fatal("failed to consume initial limiter token") } ctx, cancel := context.WithCancel(context.Background()) cancel() if err := waitNativeRate(ctx, lim, 1); !errors.Is(err, context.Canceled) { t.Fatalf("waitNativeRate error = %v, want context.Canceled", err) } } func TestSSHDirtyQueueDoesNotScanInactiveUsers(t *testing.T) { sshTrafficDirtyMu.Lock() old := sshTrafficDirty sshTrafficDirty = make(map[string]*UserState) sshTrafficDirtyMu.Unlock() defer func() { sshTrafficDirtyMu.Lock() sshTrafficDirty = old sshTrafficDirtyMu.Unlock() }() active := &UserState{Cfg: UserConfig{Username: "active"}} inactive := &UserState{Cfg: UserConfig{Username: "inactive"}} markSSHUserTrafficDirty(active) dirty := takeSSHUserTrafficDirty() if len(dirty) != 1 || dirty["active"] != active { t.Fatalf("dirty queue = %#v", dirty) } if _, found := dirty[inactive.Cfg.Username]; found { t.Fatal("inactive user appeared in dirty queue") } if next := takeSSHUserTrafficDirty(); len(next) != 0 { t.Fatalf("dirty queue was not drained: %#v", next) } } func TestOldNativeConnectionCannotDecrementReplacementAccount(t *testing.T) { oldStore := statsStore statsStore = &Store{} defer func() { statsStore = oldStore }() const uuid = "replacement-active-user" oldState := &xrayNativeQuotaState{maxConns: 1, generation: 1} m := &XrayManager{nativeQuotaByUUID: map[string]*xrayNativeQuotaState{uuid: oldState}} release, acquiredState, ok := m.acquireNativeClientConnection(uuid, "old@example") if !ok || acquiredState != oldState { t.Fatal("failed to acquire old account connection") } // Discard the old account's successful connect delta so this assertion only // measures what happens when that old connection later disconnects. m.nativeDBMu.Lock() m.nativeActivePending = nil m.nativeDBMu.Unlock() newState := &xrayNativeQuotaState{maxConns: 1, generation: 1} m.nativeQuotaMu.Lock() m.nativeQuotaByUUID[uuid] = newState m.nativeQuotaMu.Unlock() release() m.nativeDBMu.Lock() pending := m.nativeActivePending[uuid] m.nativeDBMu.Unlock() if pending.Delta != 0 || pending.State != nil { t.Fatalf("old disconnect was queued against replacement account: %+v", pending) } newState.mu.Lock() active := newState.activeConns newState.mu.Unlock() if active != 0 { t.Fatalf("replacement account active count changed to %d", active) } } func TestOldNativeTrafficCannotAttachToReplacementAccount(t *testing.T) { oldStore := statsStore statsStore = &Store{} defer func() { statsStore = oldStore }() const uuid = "replacement-traffic-user" oldState := &xrayNativeQuotaState{generation: 1} newState := &xrayNativeQuotaState{generation: 1} m := &XrayManager{nativeQuotaByUUID: map[string]*xrayNativeQuotaState{uuid: oldState}} meter := newTrafficMeter(uuid, "old@example", true, oldState) meter.n = 1234 m.nativeQuotaMu.Lock() m.nativeQuotaByUUID[uuid] = newState m.nativeQuotaMu.Unlock() oldMgr := xrayMgr xrayMgr = m defer func() { xrayMgr = oldMgr }() meter.flush() m.nativeDBMu.Lock() pending := m.nativeTrafficPending[uuid] m.nativeDBMu.Unlock() if pending.Uplink != 0 || pending.Downlink != 0 || pending.State != nil { t.Fatalf("old traffic was queued against replacement account: %+v", pending) } m.statsMu.RLock() stat := m.statsByEmail["old@example"] m.statsMu.RUnlock() if stat.Uplink != 0 || stat.Downlink != 0 { t.Fatalf("old traffic resurfaced in runtime stats: %+v", stat) } } func TestNativeFlusherDropsMismatchedPolicyIdentity(t *testing.T) { oldStore := statsStore statsStore = &Store{} defer func() { statsStore = oldStore }() const uuid = "identity-prune-user" oldState := &xrayNativeQuotaState{generation: 1} newState := &xrayNativeQuotaState{generation: 1} m := &XrayManager{ nativeQuotaByUUID: map[string]*xrayNativeQuotaState{uuid: newState}, nativeTrafficPending: map[string]xrayPendingTraffic{ uuid: {Email: "old@example", Uplink: 99, State: oldState}, }, nativeActivePending: map[string]xrayPendingActive{ uuid: {Email: "old@example", Delta: -1, State: oldState}, }, } m.flushNativeStatsToDB() m.nativeDBMu.Lock() defer m.nativeDBMu.Unlock() if len(m.nativeTrafficPending) != 0 || len(m.nativeActivePending) != 0 { t.Fatalf("mismatched pending deltas survived prune: traffic=%v active=%v", m.nativeTrafficPending, m.nativeActivePending) } } func TestNativeMuxFinishRunsOnce(t *testing.T) { var calls atomic.Int32 s := &nativeMuxSession{ closed: make(chan struct{}), uplink: make(chan nativeMuxUplinkItem), onClose: func(*nativeMuxSession) { calls.Add(1) }, } var wg sync.WaitGroup for i := 0; i < 32; i++ { wg.Add(1) go func() { defer wg.Done() s.finish() }() } wg.Wait() if got := calls.Load(); got != 1 { t.Fatalf("mux onClose called %d times, want 1", got) } } type closeTrackingReader struct { closed atomic.Bool } func (r *closeTrackingReader) Read([]byte) (int, error) { return 0, io.EOF } func (r *closeTrackingReader) Close() error { r.closed.Store(true) return nil } func TestNativeXHTTPQueueCloseClosesQueuedStreamReader(t *testing.T) { q := newNativeXHTTPUploadQueue(1, 1024) r := &closeTrackingReader{} if err := q.push(context.Background(), nativeXHTTPPacket{Reader: r}, nil); err != nil { t.Fatalf("queue stream reader: %v", err) } q.close() if !r.closed.Load() { t.Fatal("queued stream reader was not closed during queue shutdown") } }