fix: Accurate session count to avoid constant upward drift (#1348)

* fix: Accurate session count to avoid constant upward drift

The register handler incremented the session counter before `SetIDPublicKey` could fail, leaking +1 on duplicate IDs or bad keys. Now only incremented after successful registration.

The deregister handler decremented the counter before validating the request, leaking -1 on malformed or unauthorized requests. Removed the explicit decrement entirely (handled in logic below).

Cache eviction and TTL expiry silently removed sessions without decrementing the counter, causing monotonic growth. Unified all session decrements into a single cache removal callback that fires on deregistration, eviction, and cache close, filtering to only count entries with a SecretKey (true client sessions).

session_total was added as a metric so that even short lived sessions can be viewed in the metrics.

* address review nits in session-tracking tests

---------

Co-authored-by: Mzack9999 <mzack9999@protonmail.com>

Author: jentfoo

cmd/interactsh-server/main.go

@@ -14,6 +14,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	_ "net/http/pprof"
@@ -261,6 +262,11 @@ func main() {
 		}
 	}
 
+	serverOptions.Stats = &server.Metrics{}
+	storeOptions.OnRemoval = func() {
+		atomic.AddInt64(&serverOptions.Stats.Sessions, -1)
+	}
+
 	var err error
 	store, err = storage.New(&storeOptions)
 	if err != nil {
@@ -273,8 +279,6 @@ func main() {
 		_ = serverOptions.Storage.SetID(serverOptions.Token)
 	}
 
-	serverOptions.Stats = &server.Metrics{}
-
 	// If root-tld is enabled create a singleton unencrypted record in the store
 	if serverOptions.RootTLD {
 		for _, domain := range serverOptions.Domains {

pkg/server/http_server.go

@@ -388,13 +388,13 @@ func (h *HTTPServer) registerHandler(w http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	atomic.AddInt64(&h.options.Stats.Sessions, 1)
-
 	if err := h.options.Storage.SetIDPublicKey(r.CorrelationID, r.SecretKey, r.PublicKey); err != nil {
 		gologger.Warning().Msgf("Could not set id and public key for %s: %s\n", r.CorrelationID, err)
 		jsonError(w, fmt.Sprintf("could not set id and public key: %s", err), http.StatusBadRequest)
 		return
 	}
+	atomic.AddInt64(&h.options.Stats.Sessions, 1)
+	atomic.AddInt64(&h.options.Stats.SessionsTotal, 1)
 	jsonMsg(w, "registration successful", http.StatusOK)
 	gologger.Debug().Msgf("Registered correlationID %s for key\n", r.CorrelationID)
 }
@@ -409,8 +409,6 @@ type DeregisterRequest struct {
 
 // deregisterHandler is a handler for client deregister requests
 func (h *HTTPServer) deregisterHandler(w http.ResponseWriter, req *http.Request) {
-	atomic.AddInt64(&h.options.Stats.Sessions, -1)
-
 	r := &DeregisterRequest{}
 	if err := jsoniter.NewDecoder(req.Body).Decode(r); err != nil {
 		gologger.Warning().Msgf("Could not decode json body: %s\n", err)

pkg/server/http_server_test.go

@@ -1,13 +1,25 @@
 package server
 
 import (
+	"bytes"
+	"crypto/rand"
+	"crypto/rsa"
 	"crypto/tls"
+	"crypto/x509"
+	"encoding/base64"
+	"encoding/pem"
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
+	jsoniter "github.com/json-iterator/go"
+	"github.com/projectdiscovery/interactsh/pkg/storage"
+	"github.com/rs/xid"
 	"github.com/stretchr/testify/require"
 )
 
@@ -70,3 +82,72 @@ func TestWriteResponseFromDynamicRequest(t *testing.T) {
 		require.Equal(t, resp.Header.Get("Test"), "Another", "could not get correct result")
 	})
 }
+
+func TestSessionTotalMetric(t *testing.T) {
+	stats := &Metrics{}
+	removed := make(chan struct{})
+	closeOnce := sync.Once{}
+
+	store, err := storage.New(&storage.Options{
+		EvictionTTL: 5 * time.Minute,
+		OnRemoval: func() {
+			atomic.AddInt64(&stats.Sessions, -1)
+			closeOnce.Do(func() { close(removed) })
+		},
+	})
+	require.NoError(t, err)
+	defer func() { _ = store.Close() }()
+
+	h := &HTTPServer{
+		options: &Options{
+			Storage: store,
+			Stats:   stats,
+		},
+	}
+
+	// Generate a client key pair and registration payload.
+	priv, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+	pubBytes, err := x509.MarshalPKIXPublicKey(&priv.PublicKey)
+	require.NoError(t, err)
+	pubPem := pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubBytes})
+	pubB64 := base64.StdEncoding.EncodeToString(pubPem)
+
+	correlationID := xid.New().String()
+	secretKey := uuid.New().String()
+
+	// --- Register ---
+	regBody, err := jsoniter.Marshal(&RegisterRequest{
+		PublicKey:     pubB64,
+		SecretKey:     secretKey,
+		CorrelationID: correlationID,
+	})
+	require.NoError(t, err)
+	req := httptest.NewRequest("POST", "/register", bytes.NewReader(regBody))
+	w := httptest.NewRecorder()
+	h.registerHandler(w, req)
+	require.Equal(t, http.StatusOK, w.Code)
+
+	require.Equal(t, int64(1), atomic.LoadInt64(&stats.Sessions), "sessions should be 1 after register")
+	require.Equal(t, int64(1), atomic.LoadInt64(&stats.SessionsTotal), "sessions_total should be 1 after register")
+
+	// --- Deregister ---
+	deregBody, err := jsoniter.Marshal(&DeregisterRequest{
+		SecretKey:     secretKey,
+		CorrelationID: correlationID,
+	})
+	require.NoError(t, err)
+	req = httptest.NewRequest("POST", "/deregister", bytes.NewReader(deregBody))
+	w = httptest.NewRecorder()
+	h.deregisterHandler(w, req)
+	require.Equal(t, http.StatusOK, w.Code)
+
+	select {
+	case <-removed:
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out waiting for OnRemoval callback")
+	}
+
+	require.Equal(t, int64(0), atomic.LoadInt64(&stats.Sessions), "sessions should be 0 after deregister")
+	require.Equal(t, int64(1), atomic.LoadInt64(&stats.SessionsTotal), "sessions_total should remain 1 after deregister")
+}

pkg/server/metrics.go

@@ -9,17 +9,18 @@ import (
 )
 
 type Metrics struct {
-	Dns      uint64                `json:"dns"`
-	Ftp      uint64                `json:"ftp"`
-	Http     uint64                `json:"http"`
-	Ldap     uint64                `json:"ldap"`
-	Smb      uint64                `json:"smb"`
-	Smtp     uint64                `json:"smtp"`
-	Sessions int64                 `json:"sessions"`
-	Cache    *storage.CacheMetrics `json:"cache"`
-	Memory   *MemoryMetrics        `json:"memory"`
-	Cpu      *CpuStats             `json:"cpu"`
-	Network  *NetworkStats         `json:"network"`
+	Dns           uint64                `json:"dns"`
+	Ftp           uint64                `json:"ftp"`
+	Http          uint64                `json:"http"`
+	Ldap          uint64                `json:"ldap"`
+	Smb           uint64                `json:"smb"`
+	Smtp          uint64                `json:"smtp"`
+	Sessions      int64                 `json:"sessions"`
+	SessionsTotal int64                 `json:"sessions_total"`
+	Cache         *storage.CacheMetrics `json:"cache"`
+	Memory        *MemoryMetrics        `json:"memory"`
+	Cpu           *CpuStats             `json:"cpu"`
+	Network       *NetworkStats         `json:"network"`
 }
 
 func GetCacheMetrics(options *Options) *storage.CacheMetrics {

pkg/storage/option.go

@@ -15,6 +15,9 @@ type Options struct {
 	MaxSize                int
 	MaxSharedInteractions  int
 	EvictionStrategy       EvictionStrategy
+	// OnRemoval is called for each client session removed from cache
+	// (deregistration, TTL expiry, size eviction, or cache close).
+	OnRemoval func()
 }
 
 func (options *Options) UseDisk() bool {

pkg/storage/roundtrip_test.go

@@ -403,7 +403,7 @@ func TestStaleDataCleanupOnReRegistration(t *testing.T) {
 	_ = priv1
 }
 
-// TestCacheEvictionCleansLevelDB verifies the OnCacheRemovalCallback properly
+// TestCacheEvictionCleansLevelDB verifies the onCacheRemoval callback properly
 // deletes LevelDB entries when cache entries are evicted.
 func TestCacheEvictionCleansLevelDB(t *testing.T) {
 	tmpDir, err := os.MkdirTemp("", "interactsh-eviction-*")
@@ -451,7 +451,74 @@ func TestCacheEvictionCleansLevelDB(t *testing.T) {
 	// Small delay for async eviction callback
 	time.Sleep(50 * time.Millisecond)
 
-	// LevelDB entry should be cleaned up by OnCacheRemovalCallback
+	// LevelDB entry should be cleaned up by onCacheRemoval
 	_, err = db.db.Get([]byte(correlationID), nil)
 	require.Error(t, err, "LevelDB entry should be deleted after cache eviction")
 }
+
+// TestOnRemovalSessionTracking verifies that the OnRemoval callback fires
+// exactly once per client session on deregister and TTL eviction, and does
+// not fire for non-session entries created via SetID.
+func TestOnRemovalSessionTracking(t *testing.T) {
+	removed := make(chan struct{}, 10)
+	onRemoval := func() { removed <- struct{}{} }
+
+	db, err := New(&Options{
+		EvictionTTL:      50 * time.Millisecond,
+		EvictionStrategy: EvictionStrategyFixed,
+		OnRemoval:        onRemoval,
+	})
+	require.NoError(t, err)
+	defer func() { _ = db.Close() }()
+
+	waitRemoval := func(msg string) {
+		t.Helper()
+		select {
+		case <-removed:
+		case <-time.After(2 * time.Second):
+			t.Fatalf("timed out waiting for OnRemoval: %s", msg)
+		}
+	}
+
+	// --- Non-session entries (SetID) must not trigger OnRemoval ---
+	// Invalidate a SetID entry, then register+deregister a real session as a
+	// FIFO barrier: the cache event channel is ordered, so receiving the
+	// session's callback proves the SetID invalidation was already processed.
+	_ = db.SetID("token-entry")
+	db.cache.Invalidate("token-entry")
+
+	secret := uuid.New().String()
+	cid := xid.New().String()
+	_, pubKey := generateRSAKeyPair(t)
+	require.NoError(t, db.SetIDPublicKey(cid, secret, pubKey))
+	require.NoError(t, db.RemoveID(cid, secret))
+	waitRemoval("deregister barrier")
+	select {
+	case <-removed:
+		t.Fatal("SetID entry should not trigger OnRemoval")
+	default:
+	}
+
+	// --- TTL eviction must trigger OnRemoval ---
+	secret2 := uuid.New().String()
+	cid2 := xid.New().String()
+	_, pubKey2 := generateRSAKeyPair(t)
+	require.NoError(t, db.SetIDPublicKey(cid2, secret2, pubKey2))
+
+	// Periodically access the cache to trigger lazy eviction.
+	stop := make(chan struct{})
+	defer close(stop)
+	go func() {
+		ticker := time.NewTicker(10 * time.Millisecond)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-stop:
+				return
+			case <-ticker.C:
+				db.cache.GetIfPresent(cid2)
+			}
+		}
+	}()
+	waitRemoval("TTL eviction")
+}

pkg/storage/storagedb.go

@@ -50,9 +50,7 @@ func New(options *Options) (*StorageDB, error) {
 			cacheOptions = append(cacheOptions, cache.WithExpireAfterAccess(options.EvictionTTL))
 		}
 	}
-	if options.UseDisk() {
-		cacheOptions = append(cacheOptions, cache.WithRemovalListener(storageDB.OnCacheRemovalCallback))
-	}
+	cacheOptions = append(cacheOptions, cache.WithRemovalListener(storageDB.onCacheRemoval))
 	cacheDb := cache.New(cacheOptions...)
 	storageDB.cache = cacheDb
 
@@ -77,10 +75,21 @@ func New(options *Options) (*StorageDB, error) {
 	return storageDB, nil
 }
 
-func (s *StorageDB) OnCacheRemovalCallback(key cache.Key, value cache.Value) {
-	if k, ok := key.(string); ok {
+func (s *StorageDB) onCacheRemoval(key cache.Key, value cache.Value) {
+	k, ok := key.(string)
+	if !ok {
+		return
+	}
+	if s.Options.UseDisk() && s.db != nil {
 		_ = s.db.Delete([]byte(k), &opt.WriteOptions{})
 	}
+	// Only fire for client sessions (entries with a SecretKey),
+	// not for token/domain entries created via SetID.
+	if s.Options.OnRemoval != nil {
+		if cd, ok := value.(*CorrelationData); ok && cd.SecretKey != "" {
+			s.Options.OnRemoval()
+		}
+	}
 }
 
 func (s *StorageDB) GetCacheMetrics() (*CacheMetrics, error) {