MWSE/internal/services/ippressure.go

package services

import (
	"fmt"
	"sync"

	"git.saqut.com/saqut/mwse/internal/protocol"
	"git.saqut.com/saqut/mwse/internal/ws"
)

// shortCodeAlphabet is the 22-letter set the original used (note: J, Q, U, W are
// intentionally absent). Three letters give 22^3 = 10,648 codes.
const shortCodeAlphabet = "ABCDEFGHIKLMNOPRSTVXYZ"

// Announcer receives address allocation events. In the multi-process Node
// deployment these were forwarded to a parent via process.send for the live
// traffic panel. For the single-node 0.1.0 core the default is a no-op; a cluster
// integration can supply a real implementation later.
type Announcer interface {
	Announce(kind, action, clientID string, value any)
}

type noopAnnouncer struct{}

func (noopAnnouncer) Announce(string, string, string, any) {}

// IPPressure allocates three kinds of unique virtual address to clients. A single
// mutex guards all three tables; allocation is infrequent relative to messaging,
// so finer-grained locking would add complexity for no real gain.
type IPPressure struct {
	ann Announcer

	mu         sync.Mutex
	busyNumber map[int]string    // number -> clientID
	busyCode   map[string]string // shortcode -> clientID
	busyIP     map[string]string // ip -> clientID
}

// NewIPPressure builds an allocator. A nil announcer becomes a no-op.
func NewIPPressure(ann Announcer) *IPPressure {
	if ann == nil {
		ann = noopAnnouncer{}
	}
	return &IPPressure{
		ann:        ann,
		busyNumber: make(map[int]string),
		busyCode:   make(map[string]string),
		busyIP:     make(map[string]string),
	}
}

// ---- number (starts at 24, counts up) -----------------------------------

func (p *IPPressure) lockNumber(clientID string) int {
	p.mu.Lock()
	defer p.mu.Unlock()
	for n := 24; ; n++ {
		if _, busy := p.busyNumber[n]; !busy {
			p.busyNumber[n] = clientID
			p.ann.Announce("AP_NUMBER", "LOCK", clientID, n)
			return n
		}
	}
}

func (p *IPPressure) releaseNumber(n int) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if clientID, ok := p.busyNumber[n]; ok {
		p.ann.Announce("AP_NUMBER", "RELEASE", clientID, n)
		delete(p.busyNumber, n)
	}
}

func (p *IPPressure) whoisNumber(n int) (string, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	id, ok := p.busyNumber[n]
	return id, ok
}

// ---- short code (three letters from the restricted alphabet) ------------

func (p *IPPressure) lockCode(clientID string) string {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, a := range shortCodeAlphabet {
		for _, b := range shortCodeAlphabet {
			for _, d := range shortCodeAlphabet {
				code := string([]rune{a, b, d})
				if _, busy := p.busyCode[code]; !busy {
					p.busyCode[code] = clientID
					p.ann.Announce("AP_SHORTCODE", "LOCK", clientID, code)
					return code
				}
			}
		}
	}
	return "" // address space exhausted
}

func (p *IPPressure) releaseCode(code string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if clientID, ok := p.busyCode[code]; ok {
		p.ann.Announce("AP_SHORTCODE", "RELEASE", clientID, code)
		delete(p.busyCode, code)
	}
}

func (p *IPPressure) whoisCode(code string) (string, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	id, ok := p.busyCode[code]
	return id, ok
}

// ---- ip address (10.0.0.1 upward) ---------------------------------------

func (p *IPPressure) lockIP(clientID string) string {
	p.mu.Lock()
	defer p.mu.Unlock()
	a, b, cc, d := 10, 0, 0, 1
	for {
		ip := fmt.Sprintf("%d.%d.%d.%d", a, b, cc, d)
		if _, busy := p.busyIP[ip]; !busy {
			p.busyIP[ip] = clientID
			p.ann.Announce("AP_IPADDRESS", "LOCK", clientID, ip)
			return ip
		}
		switch {
		case d != 255:
			d++
		case cc != 255:
			d, cc = 0, cc+1
		case b != 255:
			d, cc, b = 0, 0, b+1
		case a != 255:
			d, cc, b, a = 0, 0, 0, a+1
		default:
			return "" // address space exhausted
		}
	}
}

func (p *IPPressure) releaseIP(ip string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if clientID, ok := p.busyIP[ip]; ok {
		p.ann.Announce("AP_IPADDRESS", "RELEASE", clientID, ip)
		delete(p.busyIP, ip)
	}
}

func (p *IPPressure) whoisIP(ip string) (string, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	id, ok := p.busyIP[ip]
	return id, ok
}

// registerIPPressure wires the alloc/realloc/release/whois handlers and the
// disconnect cleanup. The allocator instance is returned for tests.
func registerIPPressure(hub *ws.Hub, ann Announcer) *IPPressure {
	p := NewIPPressure(ann)

	// --- IP address ---
	hub.Register("alloc/APIPAddress", func(c *ws.Client, m protocol.Message) any {
		if ip := c.APIP(); ip != "" {
			return map[string]any{"status": "success", "ip": ip}
		}
		ip := p.lockIP(c.ID)
		if ip == "" {
			return map[string]any{"status": "fail"}
		}
		c.SetAPIP(ip)
		return map[string]any{"status": "success", "ip": ip}
	})
	hub.Register("realloc/APIPAddress", func(c *ws.Client, m protocol.Message) any {
		old := c.APIP()
		if old == "" {
			return map[string]any{"status": "fail"}
		}
		// Allocate the new address before freeing the old one, so realloc actually
		// yields a different address (the old stays reserved during the search).
		ip := p.lockIP(c.ID)
		if ip == "" {
			return map[string]any{"status": "fail"} // exhausted; keep the old one
		}
		p.releaseIP(old)
		c.SetAPIP(ip)
		return map[string]any{"status": "success", "ip": ip}
	})
	hub.Register("release/APIPAddress", func(c *ws.Client, m protocol.Message) any {
		p.releaseIP(c.APIP())
		c.SetAPIP("")
		return success()
	})
	hub.Register("whois/APIPAddress", func(c *ws.Client, m protocol.Message) any {
		if id, ok := p.whoisIP(m.Str("whois")); ok {
			return map[string]any{"status": "success", "socket": id}
		}
		return map[string]any{"status": "fail"}
	})

	// --- number ---
	hub.Register("alloc/APNumber", func(c *ws.Client, m protocol.Message) any {
		if n := c.APNumber(); n != 0 {
			return map[string]any{"status": "success", "number": n}
		}
		n := p.lockNumber(c.ID)
		c.SetAPNumber(n)
		return map[string]any{"status": "success", "number": n}
	})
	hub.Register("realloc/APNumber", func(c *ws.Client, m protocol.Message) any {
		old := c.APNumber()
		if old == 0 {
			return map[string]any{"status": "fail"}
		}
		n := p.lockNumber(c.ID) // old stays reserved, so n != old
		p.releaseNumber(old)
		c.SetAPNumber(n)
		return map[string]any{"status": "success", "number": n}
	})
	hub.Register("release/APNumber", func(c *ws.Client, m protocol.Message) any {
		p.releaseNumber(c.APNumber())
		c.SetAPNumber(0)
		return success()
	})
	hub.Register("whois/APNumber", func(c *ws.Client, m protocol.Message) any {
		if id, ok := p.whoisNumber(m.Int("whois")); ok {
			return map[string]any{"status": "success", "socket": id}
		}
		return map[string]any{"status": "fail"}
	})

	// --- short code ---
	hub.Register("alloc/APShortCode", func(c *ws.Client, m protocol.Message) any {
		if code := c.APShortCode(); code != "" {
			return map[string]any{"status": "success", "code": code}
		}
		code := p.lockCode(c.ID)
		if code == "" {
			return map[string]any{"status": "fail"}
		}
		c.SetAPShortCode(code)
		return map[string]any{"status": "success", "code": code}
	})
	hub.Register("realloc/APShortCode", func(c *ws.Client, m protocol.Message) any {
		old := c.APShortCode()
		if old == "" {
			return map[string]any{"status": "fail"}
		}
		code := p.lockCode(c.ID)
		if code == "" {
			return map[string]any{"status": "fail"} // exhausted; keep the old one
		}
		p.releaseCode(old)
		c.SetAPShortCode(code)
		return map[string]any{"status": "success", "code": code}
	})
	hub.Register("release/APShortCode", func(c *ws.Client, m protocol.Message) any {
		p.releaseCode(c.APShortCode())
		c.SetAPShortCode("")
		return success()
	})
	hub.Register("whois/APShortCode", func(c *ws.Client, m protocol.Message) any {
		if id, ok := p.whoisCode(m.Str("whois")); ok {
			return map[string]any{"status": "success", "socket": id}
		}
		return map[string]any{"status": "fail"}
	})

	// Release every address a client held when it disconnects.
	hub.OnDisconnect(func(c *ws.Client) {
		if c.APIP() != "" {
			p.releaseIP(c.APIP())
		}
		if c.APNumber() != 0 {
			p.releaseNumber(c.APNumber())
		}
		if c.APShortCode() != "" {
			p.releaseCode(c.APShortCode())
		}
	})

	return p
}