DragonCoreSSH-NewWEB/quota_hardening_test.go

package main

import (
	"context"
	"errors"
	"io"
	"net"
	"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 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 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 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")
	}
}