sing-box/service/usermanager/manager.go

package usermanager

import (
	"bytes"
	"context"
	"crypto/subtle"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"github.com/sagernet/sing-box/adapter"
	boxService "github.com/sagernet/sing-box/adapter/service"
	C "github.com/sagernet/sing-box/constant"
	"github.com/sagernet/sing-box/log"
	"github.com/sagernet/sing-box/option"
	"github.com/sagernet/sing/common/bufio"
	E "github.com/sagernet/sing/common/exceptions"
	N "github.com/sagernet/sing/common/network"
	"github.com/sagernet/sing/service"
)

var _ adapter.UserManager = (*Manager)(nil)
var _ adapter.ConnectionTracker = (*Manager)(nil)

type authCacheEntry struct {
	username  string
	expiresAt time.Time
}

type userConnCounter struct {
	username string
	tx       int64
	rx       int64
	done     chan struct{}
}

type Manager struct {
	boxService.Adapter
	ctx    context.Context
	cancel context.CancelFunc
	logger log.ContextLogger

	authServer    string
	timeout       time.Duration
	cacheTTL      time.Duration
	refreshInt    time.Duration
	reportInt     time.Duration
	reportTraffic bool
	apiSecret     string
	apiListen     string

	httpClient *http.Client

	access    sync.RWMutex
	authCache map[string]*authCacheEntry
	userCreds map[string]string
	userConns map[string]map[*userConnCounter]struct{}

	protocolUsers   map[string][]adapter.UserInfo
	protocolLastSeq map[string]int64

	trafficAccess sync.Mutex
	traffic       map[string]*UserTrafficAtomic

	stopRefresh chan struct{}
	stopReport  chan struct{}
	stopAPI     chan struct{}
}

type UserTrafficAtomic struct {
	Tx atomic.Int64
	Rx atomic.Int64
}

func New(ctx context.Context, logger log.ContextLogger, tag string, options option.UserManagerOptions) (*Manager, error) {
	ctx, cancel := context.WithCancel(ctx)
	m := &Manager{
		Adapter:        boxService.NewAdapter(C.TypeUserManager, tag),
		ctx:            ctx,
		cancel:         cancel,
		logger:         logger,
		authServer:     options.AuthServer,
		timeout:        time.Duration(options.Timeout),
		cacheTTL:       time.Duration(options.CacheTTL),
		refreshInt:     time.Duration(options.RefreshInterval),
		reportInt:      time.Duration(options.ReportInterval),
		reportTraffic:  options.ReportTraffic,
		apiSecret:      options.APISecret,
		apiListen:      options.APIListen,
		httpClient: &http.Client{
			Timeout: time.Duration(options.Timeout),
		},
		authCache:       make(map[string]*authCacheEntry),
		userCreds:       make(map[string]string),
		userConns:       make(map[string]map[*userConnCounter]struct{}),
		protocolUsers:   make(map[string][]adapter.UserInfo),
		protocolLastSeq: make(map[string]int64),
		traffic:         make(map[string]*UserTrafficAtomic),
		stopRefresh:     make(chan struct{}),
		stopReport:      make(chan struct{}),
		stopAPI:         make(chan struct{}),
	}
	if m.timeout <= 0 {
		m.timeout = 10 * time.Second
	}
	if m.cacheTTL <= 0 {
		m.cacheTTL = 5 * time.Minute
	}
	if m.refreshInt <= 0 {
		m.refreshInt = 30 * time.Second
	}
	if m.reportInt <= 0 {
		m.reportInt = 60 * time.Second
	}
	return m, nil
}

func (m *Manager) Start(stage adapter.StartStage) error {
	if stage != adapter.StartStateStart {
		return nil
	}
	m.logger.Info("starting user manager")

	if err := m.refreshCredentials(); err != nil {
		m.logger.Warn("initial credential fetch failed: ", err)
	}

	router := service.FromContext[adapter.Router](m.ctx)
	if router != nil {
		router.AppendTracker(m)
		m.logger.Debug("registered as connection tracker")
	}

	go func() {
		ticker := time.NewTicker(m.refreshInt)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				if err := m.refreshCredentials(); err != nil {
					m.logger.Warn("credential refresh failed: ", err)
				}
			case <-m.stopRefresh:
				return
			}
		}
	}()

	if m.reportTraffic && m.authServer != "" {
		go func() {
			ticker := time.NewTicker(m.reportInt)
			defer ticker.Stop()
			select {
			case <-time.After(10 * time.Second):
			case <-m.stopReport:
				return
			}
			for {
				select {
				case <-ticker.C:
					m.reportTrafficToServer()
				case <-m.stopReport:
					return
				}
			}
		}()
	}

	if m.apiListen != "" {
		go m.serveAPI()
	}

	return nil
}

func (m *Manager) Close() error {
	m.cancel()
	close(m.stopRefresh)
	close(m.stopReport)
	close(m.stopAPI)
	return nil
}

func (m *Manager) Authenticate(ctx context.Context, protocol string, credential string, addr string) (string, error) {
	if credential == "" {
		return "", E.New("empty credential")
	}

	m.access.RLock()
	entry, cached := m.authCache[credential]
	m.access.RUnlock()
	if cached && time.Now().Before(entry.expiresAt) {
		return entry.username, nil
	}

	username, err := m.remoteAuth(protocol, credential, addr)
	if err != nil {
		return "", err
	}

	m.access.Lock()
	m.authCache[credential] = &authCacheEntry{
		username:  username,
		expiresAt: time.Now().Add(m.cacheTTL),
	}
	m.userCreds[username] = credential
	m.access.Unlock()

	return username, nil
}

type authRequest struct {
	Protocol   string `json:"protocol"`
	Credential string `json:"credential"`
	Addr       string `json:"addr,omitempty"`
}

type authResponse struct {
	Ok bool   `json:"ok"`
	ID string `json:"id,omitempty"`
}

func (m *Manager) remoteAuth(protocol string, credential string, addr string) (string, error) {
	if m.authServer == "" {
		return "", E.New("auth server not configured")
	}
	req := authRequest{
		Protocol:   protocol,
		Credential: credential,
		Addr:       addr,
	}
	data, err := json.Marshal(req)
	if err != nil {
		return "", E.Cause(err, "marshal auth request")
	}
	httpReq, err := http.NewRequestWithContext(m.ctx, "POST", m.authServer+"/api/auth", bytes.NewReader(data))
	if err != nil {
		return "", E.Cause(err, "create auth request")
	}
	httpReq.Header.Set("Content-Type", "application/json")
	resp, err := m.httpClient.Do(httpReq)
	if err != nil {
		return "", E.Cause(err, "auth request failed")
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusForbidden {
		return "", E.New("authentication rejected")
	}
	if resp.StatusCode != http.StatusOK {
		return "", E.New("auth server returned ", resp.StatusCode)
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", E.Cause(err, "read auth response")
	}
	var authResp authResponse
	if err := json.Unmarshal(body, &authResp); err != nil {
		return "", E.Cause(err, "parse auth response")
	}
	if !authResp.Ok {
		return "", E.New("authentication rejected")
	}
	if authResp.ID == "" {
		return "", E.New("auth response missing user ID")
	}
	return authResp.ID, nil
}

func (m *Manager) GetCredentials(protocol string) ([]adapter.UserInfo, error) {
	m.access.RLock()
	defer m.access.RUnlock()
	users := m.protocolUsers[protocol]
	if users == nil {
		return nil, nil
	}
	result := make([]adapter.UserInfo, len(users))
	copy(result, users)
	return result, nil
}

type credentialsResponse struct {
	Credentials []credentialEntry `json:"credentials"`
}

type credentialEntry struct {
	Protocol   string `json:"protocol"`
	Credential string `json:"credential"`
	Username   string `json:"username"`
}

func (m *Manager) refreshCredentials() error {
	if m.authServer == "" {
		return nil
	}
	httpReq, err := http.NewRequestWithContext(m.ctx, "GET", m.authServer+"/api/credentials", nil)
	if err != nil {
		return err
	}
	resp, err := m.httpClient.Do(httpReq)
	if err != nil {
		return E.Cause(err, "fetch credentials")
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return E.New("credentials server returned ", resp.StatusCode)
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return E.Cause(err, "read credentials response")
	}
	var credResp credentialsResponse
	if err := json.Unmarshal(body, &credResp); err != nil {
		return E.Cause(err, "parse credentials response")
	}

	newCache := make(map[string]*authCacheEntry)
	newUserCreds := make(map[string]string)
	byProtocol := make(map[string][]adapter.UserInfo)

	for _, c := range credResp.Credentials {
		newCache[c.Credential] = &authCacheEntry{
			username:  c.Username,
			expiresAt: time.Now().Add(m.cacheTTL),
		}
		newUserCreds[c.Username] = c.Credential
		byProtocol[c.Protocol] = append(byProtocol[c.Protocol], adapter.UserInfo{
			Username:   c.Username,
			Credential: c.Credential,
		})
	}

	m.access.Lock()
	for k, v := range newCache {
		if _, exists := m.authCache[k]; !exists {
			m.authCache[k] = v
		}
	}
	for k, v := range newUserCreds {
		m.userCreds[k] = v
	}
	m.protocolUsers = byProtocol
	m.access.Unlock()

	return nil
}

func (m *Manager) ReportTraffic(username string, tx int64, rx int64) {
	if username == "" {
		return
	}
	m.trafficAccess.Lock()
	t, ok := m.traffic[username]
	if !ok {
		t = &UserTrafficAtomic{}
		m.traffic[username] = t
	}
	m.trafficAccess.Unlock()
	if tx > 0 {
		t.Tx.Add(tx)
	}
	if rx > 0 {
		t.Rx.Add(rx)
	}
}

func (m *Manager) GetTraffic(username string) (int64, int64) {
	m.trafficAccess.Lock()
	t, ok := m.traffic[username]
	m.trafficAccess.Unlock()
	if !ok {
		return 0, 0
	}
	return t.Tx.Load(), t.Rx.Load()
}

func (m *Manager) ListTraffic() map[string]adapter.UserTraffic {
	m.trafficAccess.Lock()
	defer m.trafficAccess.Unlock()
	result := make(map[string]adapter.UserTraffic, len(m.traffic))
	for username, t := range m.traffic {
		result[username] = adapter.UserTraffic{
			Tx: t.Tx.Load(),
			Rx: t.Rx.Load(),
		}
	}
	return result
}

func (m *Manager) KickUser(username string) {
	m.access.Lock()
	defer m.access.Unlock()
	cred, ok := m.userCreds[username]
	if ok {
		delete(m.authCache, cred)
		delete(m.userCreds, username)
	}
	for proto, users := range m.protocolUsers {
		filtered := make([]adapter.UserInfo, 0, len(users))
		for _, u := range users {
			if u.Username != username {
				filtered = append(filtered, u)
			}
		}
		m.protocolUsers[proto] = filtered
	}
	m.logger.Info("kicked user ", username)
}

// ConnectionTracker implementation

func (m *Manager) RoutedConnection(ctx context.Context, conn net.Conn, metadata adapter.InboundContext, matchedRule adapter.Rule, matchOutbound adapter.Outbound) net.Conn {
	username := metadata.User
	if username == "" {
		return conn
	}
	counter := &userConnCounter{
		username: username,
		done:     make(chan struct{}),
	}
	m.access.Lock()
	if m.userConns[username] == nil {
		m.userConns[username] = make(map[*userConnCounter]struct{})
	}
	m.userConns[username][counter] = struct{}{}
	m.access.Unlock()

	wrapped := bufio.NewCounterConn(conn,
		[]N.CountFunc{func(n int64) { counter.tx += n }},
		[]N.CountFunc{func(n int64) { counter.rx += n }},
	)

	return &trackedConn{
		ExtendedConn: wrapped,
		onClose: func() {
			m.ReportTraffic(username, counter.tx, counter.rx)
			m.access.Lock()
			if conns, ok := m.userConns[username]; ok {
				delete(conns, counter)
				if len(conns) == 0 {
					delete(m.userConns, username)
				}
			}
			m.access.Unlock()
		},
	}
}

func (m *Manager) RoutedPacketConnection(ctx context.Context, conn N.PacketConn, metadata adapter.InboundContext, matchedRule adapter.Rule, matchOutbound adapter.Outbound) N.PacketConn {
	username := metadata.User
	if username == "" {
		return conn
	}
	counter := &userConnCounter{
		username: username,
		done:     make(chan struct{}),
	}
	m.access.Lock()
	if m.userConns[username] == nil {
		m.userConns[username] = make(map[*userConnCounter]struct{})
	}
	m.userConns[username][counter] = struct{}{}
	m.access.Unlock()

	wrapped := bufio.NewCounterPacketConn(conn,
		[]N.CountFunc{func(n int64) { counter.tx += n }},
		[]N.CountFunc{func(n int64) { counter.rx += n }},
	)

	return &trackedPacketConn{
		PacketConn: wrapped,
		onClose: func() {
			m.ReportTraffic(username, counter.tx, counter.rx)
			m.access.Lock()
			if conns, ok := m.userConns[username]; ok {
				delete(conns, counter)
				if len(conns) == 0 {
					delete(m.userConns, username)
				}
			}
			m.access.Unlock()
		},
	}
}

type trackedConn struct {
	N.ExtendedConn
	onClose func()
}

func (c *trackedConn) Close() error {
	c.onClose()
	return c.ExtendedConn.Close()
}

func (c *trackedConn) Upstream() any {
	return c.ExtendedConn
}

func (c *trackedConn) ReaderReplaceable() bool {
	return true
}

func (c *trackedConn) WriterReplaceable() bool {
	return true
}

type trackedPacketConn struct {
	N.PacketConn
	onClose func()
}

func (c *trackedPacketConn) Close() error {
	c.onClose()
	return c.PacketConn.Close()
}

func (c *trackedPacketConn) Upstream() any {
	return c.PacketConn
}

func (c *trackedPacketConn) ReaderReplaceable() bool {
	return true
}

func (c *trackedPacketConn) WriterReplaceable() bool {
	return true
}

// Traffic report

type trafficReport struct {
	Traffic map[string]adapter.UserTraffic `json:"traffic"`
}

func (m *Manager) reportTrafficToServer() {
	traffic := m.ListTraffic()
	if len(traffic) == 0 {
		return
	}
	report := trafficReport{Traffic: traffic}
	data, err := json.Marshal(report)
	if err != nil {
		m.logger.Warn("marshal traffic report: ", err)
		return
	}
	httpReq, err := http.NewRequestWithContext(m.ctx, "POST", m.authServer+"/api/traffic", bytes.NewReader(data))
	if err != nil {
		m.logger.Warn("create traffic report: ", err)
		return
	}
	httpReq.Header.Set("Content-Type", "application/json")
	resp, err := m.httpClient.Do(httpReq)
	if err != nil {
		m.logger.Warn("traffic report failed: ", err)
		return
	}
	resp.Body.Close()
	m.logger.Debug("traffic reported, users=", len(traffic))
}

// HTTP API

func (m *Manager) serveAPI() {
	mux := http.NewServeMux()
	mux.HandleFunc("/traffic", m.handleTraffic)
	mux.HandleFunc("/traffic/", m.handleTrafficUser)
	mux.HandleFunc("/kick/", m.handleKick)

	apiSecret := m.apiSecret

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if apiSecret != "" {
			auth := r.Header.Get("Authorization")
			expected := "Bearer " + apiSecret
			if subtle.ConstantTimeCompare([]byte(auth), []byte(expected)) != 1 {
				http.Error(w, "Forbidden", http.StatusForbidden)
				return
			}
		}
		mux.ServeHTTP(w, r)
	})

	server := &http.Server{
		Addr:    m.apiListen,
		Handler: handler,
	}

	listener, err := net.Listen("tcp", m.apiListen)
	if err != nil {
		m.logger.Error("api listen: ", err)
		return
	}
	m.logger.Info("user manager API listening on ", m.apiListen)

	err = server.Serve(listener)
	if err != nil && !errors.Is(err, http.ErrServerClosed) {
		m.logger.Error("api serve: ", err)
	}
}

func (m *Manager) handleTraffic(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	traffic := m.ListTraffic()
	writeJSON(w, http.StatusOK, map[string]any{"users": traffic})
}

func (m *Manager) handleTrafficUser(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	username := r.URL.Path[len("/traffic/"):]
	if username == "" {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}
	tx, rx := m.GetTraffic(username)
	writeJSON(w, http.StatusOK, map[string]any{"tx": tx, "rx": rx})
}

func (m *Manager) handleKick(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	username := r.URL.Path[len("/kick/"):]
	if username == "" {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}
	m.KickUser(username)
	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}

func (m *Manager) GetCachedUsers() []string {
	m.access.RLock()
	defer m.access.RUnlock()
	users := make([]string, 0, len(m.userCreds))
	for u := range m.userCreds {
		users = append(users, u)
	}
	return users
}

var _ = fmt.Sprintf
var _ = io.Discard