caddy-sip-guardian/dialog_state.go

package sipguardian

import (
	"sync"
	"time"

	"go.uber.org/zap"
)

// DialogState tracks the state of a SIP dialog for topology hiding
type DialogState struct {
	// Dialog identifiers
	CallID      string
	FromTag     string
	ToTag       string
	LocalBranch string // Via branch we added

	// Original values (for response rewriting)
	OriginalCallID  string // if anonymized
	OriginalVia     string // original top Via header
	OriginalContact string // original Contact header

	// Proxy values (what we replaced with)
	ProxyContact string
	ProxyVia     string

	// Client information (for routing responses)
	ClientHost string
	ClientPort int

	// Timestamps
	Created      time.Time
	LastActivity time.Time

	// State tracking
	IsConfirmed bool // 2xx received for dialog-creating request
	IsTerminated bool
}

// DialogManager manages dialog state for topology hiding
type DialogManager struct {
	dialogs    map[string]*DialogState
	byBranch   map[string]*DialogState // Quick lookup by Via branch
	cleanupTTL time.Duration
	logger     *zap.Logger
	mu         sync.RWMutex
}

// NewDialogManager creates a new dialog manager
func NewDialogManager(logger *zap.Logger, cleanupTTL time.Duration) *DialogManager {
	if cleanupTTL == 0 {
		cleanupTTL = 10 * time.Minute
	}

	dm := &DialogManager{
		dialogs:    make(map[string]*DialogState),
		byBranch:   make(map[string]*DialogState),
		cleanupTTL: cleanupTTL,
		logger:     logger,
	}

	return dm
}

// dialogKey generates a unique key for a dialog
func dialogKey(callID, fromTag string) string {
	return callID + ":" + fromTag
}

// CreateDialog creates a new dialog state for an outgoing request
func (dm *DialogManager) CreateDialog(msg *SIPMessage, clientHost string, clientPort int) *DialogState {
	dm.mu.Lock()
	defer dm.mu.Unlock()

	callID := msg.GetCallID()
	fromTag := msg.GetFromTag()
	branch := msg.GetViaBranch()

	state := &DialogState{
		CallID:       callID,
		FromTag:      fromTag,
		LocalBranch:  branch,
		ClientHost:   clientHost,
		ClientPort:   clientPort,
		Created:      time.Now(),
		LastActivity: time.Now(),
	}

	// Store original Via if we're going to modify it
	if via := msg.GetHeader("Via"); via != nil {
		state.OriginalVia = via.Value
	}

	// Store original Contact
	if contact := msg.GetHeader("Contact"); contact != nil {
		state.OriginalContact = contact.Value
	}

	// Store by dialog key and by branch
	key := dialogKey(callID, fromTag)
	dm.dialogs[key] = state

	if branch != "" {
		dm.byBranch[branch] = state
	}

	dm.logger.Debug("Dialog created",
		zap.String("call_id", callID),
		zap.String("from_tag", fromTag),
		zap.String("branch", branch),
	)

	return state
}

// GetDialogByCallID retrieves a dialog by Call-ID and From-tag
func (dm *DialogManager) GetDialogByCallID(callID, fromTag string) *DialogState {
	dm.mu.RLock()
	defer dm.mu.RUnlock()

	key := dialogKey(callID, fromTag)
	return dm.dialogs[key]
}

// GetDialogByBranch retrieves a dialog by Via branch (for response routing)
func (dm *DialogManager) GetDialogByBranch(branch string) *DialogState {
	dm.mu.RLock()
	defer dm.mu.RUnlock()

	return dm.byBranch[branch]
}

// UpdateDialog updates an existing dialog with new information
func (dm *DialogManager) UpdateDialog(callID, fromTag string, toTag string) {
	dm.mu.Lock()
	defer dm.mu.Unlock()

	key := dialogKey(callID, fromTag)
	if state, ok := dm.dialogs[key]; ok {
		state.LastActivity = time.Now()
		if toTag != "" && state.ToTag == "" {
			state.ToTag = toTag
		}
	}
}

// ConfirmDialog marks a dialog as confirmed (2xx received)
func (dm *DialogManager) ConfirmDialog(callID, fromTag string) {
	dm.mu.Lock()
	defer dm.mu.Unlock()

	key := dialogKey(callID, fromTag)
	if state, ok := dm.dialogs[key]; ok {
		state.IsConfirmed = true
		state.LastActivity = time.Now()
	}
}

// TerminateDialog marks a dialog as terminated
func (dm *DialogManager) TerminateDialog(callID, fromTag string) {
	dm.mu.Lock()
	defer dm.mu.Unlock()

	key := dialogKey(callID, fromTag)
	if state, ok := dm.dialogs[key]; ok {
		state.IsTerminated = true
		state.LastActivity = time.Now()
	}
}

// RemoveDialog removes a dialog from the manager
func (dm *DialogManager) RemoveDialog(callID, fromTag string) {
	dm.mu.Lock()
	defer dm.mu.Unlock()

	key := dialogKey(callID, fromTag)
	if state, ok := dm.dialogs[key]; ok {
		// Remove from branch index
		if state.LocalBranch != "" {
			delete(dm.byBranch, state.LocalBranch)
		}
		delete(dm.dialogs, key)

		dm.logger.Debug("Dialog removed",
			zap.String("call_id", callID),
			zap.String("from_tag", fromTag),
		)
	}
}

// StoreOriginals stores original header values before rewriting
func (dm *DialogManager) StoreOriginals(callID, fromTag string, via, contact string, originalCallID string) {
	dm.mu.Lock()
	defer dm.mu.Unlock()

	key := dialogKey(callID, fromTag)
	if state, ok := dm.dialogs[key]; ok {
		if via != "" {
			state.OriginalVia = via
		}
		if contact != "" {
			state.OriginalContact = contact
		}
		if originalCallID != "" {
			state.OriginalCallID = originalCallID
		}
	}
}

// StoreProxyValues stores the values we used to replace originals
func (dm *DialogManager) StoreProxyValues(callID, fromTag string, via, contact string) {
	dm.mu.Lock()
	defer dm.mu.Unlock()

	key := dialogKey(callID, fromTag)
	if state, ok := dm.dialogs[key]; ok {
		if via != "" {
			state.ProxyVia = via
		}
		if contact != "" {
			state.ProxyContact = contact
		}
	}
}

// Cleanup removes stale dialogs
func (dm *DialogManager) Cleanup() int {
	dm.mu.Lock()
	defer dm.mu.Unlock()

	cutoff := time.Now().Add(-dm.cleanupTTL)
	removed := 0

	for key, state := range dm.dialogs {
		// Remove terminated dialogs after TTL
		if state.IsTerminated && state.LastActivity.Before(cutoff) {
			if state.LocalBranch != "" {
				delete(dm.byBranch, state.LocalBranch)
			}
			delete(dm.dialogs, key)
			removed++
			continue
		}

		// Remove unconfirmed dialogs that have been idle too long
		if !state.IsConfirmed && state.LastActivity.Before(cutoff) {
			if state.LocalBranch != "" {
				delete(dm.byBranch, state.LocalBranch)
			}
			delete(dm.dialogs, key)
			removed++
			continue
		}

		// Remove very old dialogs regardless of state (prevent memory leak)
		if state.Created.Before(time.Now().Add(-24 * time.Hour)) {
			if state.LocalBranch != "" {
				delete(dm.byBranch, state.LocalBranch)
			}
			delete(dm.dialogs, key)
			removed++
		}
	}

	if removed > 0 {
		dm.logger.Debug("Cleaned up stale dialogs", zap.Int("count", removed))
	}

	return removed
}

// GetStats returns statistics about the dialog manager
func (dm *DialogManager) GetStats() map[string]interface{} {
	dm.mu.RLock()
	defer dm.mu.RUnlock()

	confirmed := 0
	terminated := 0
	pending := 0

	for _, state := range dm.dialogs {
		if state.IsTerminated {
			terminated++
		} else if state.IsConfirmed {
			confirmed++
		} else {
			pending++
		}
	}

	return map[string]interface{}{
		"total_dialogs":      len(dm.dialogs),
		"confirmed_dialogs":  confirmed,
		"terminated_dialogs": terminated,
		"pending_dialogs":    pending,
		"branch_index_size":  len(dm.byBranch),
	}
}

// TransactionState tracks a single SIP transaction (simpler than full dialog)
type TransactionState struct {
	Branch       string
	Method       string
	ClientHost   string
	ClientPort   int
	OriginalVia  string
	Created      time.Time
}

// TransactionManager manages transaction state for stateless topology hiding
type TransactionManager struct {
	transactions map[string]*TransactionState
	cleanupTTL   time.Duration
	logger       *zap.Logger
	mu           sync.RWMutex
}

// NewTransactionManager creates a new transaction manager
func NewTransactionManager(logger *zap.Logger, cleanupTTL time.Duration) *TransactionManager {
	if cleanupTTL == 0 {
		cleanupTTL = 32 * time.Second // SIP transaction timeout
	}

	return &TransactionManager{
		transactions: make(map[string]*TransactionState),
		cleanupTTL:   cleanupTTL,
		logger:       logger,
	}
}

// CreateTransaction creates a new transaction state
func (tm *TransactionManager) CreateTransaction(branch, method, clientHost string, clientPort int, originalVia string) *TransactionState {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	state := &TransactionState{
		Branch:      branch,
		Method:      method,
		ClientHost:  clientHost,
		ClientPort:  clientPort,
		OriginalVia: originalVia,
		Created:     time.Now(),
	}

	tm.transactions[branch] = state
	return state
}

// GetTransaction retrieves a transaction by branch
func (tm *TransactionManager) GetTransaction(branch string) *TransactionState {
	tm.mu.RLock()
	defer tm.mu.RUnlock()

	return tm.transactions[branch]
}

// RemoveTransaction removes a transaction
func (tm *TransactionManager) RemoveTransaction(branch string) {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	delete(tm.transactions, branch)
}

// Cleanup removes expired transactions
func (tm *TransactionManager) Cleanup() int {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	cutoff := time.Now().Add(-tm.cleanupTTL)
	removed := 0

	for branch, state := range tm.transactions {
		if state.Created.Before(cutoff) {
			delete(tm.transactions, branch)
			removed++
		}
	}

	return removed
}

// GetStats returns statistics about the transaction manager
func (tm *TransactionManager) GetStats() map[string]interface{} {
	tm.mu.RLock()
	defer tm.mu.RUnlock()

	return map[string]interface{}{
		"active_transactions": len(tm.transactions),
	}
}