diff --git a/internal/web/service/host.go b/internal/web/service/host.go index 4da687ebf..4343e9387 100644 --- a/internal/web/service/host.go +++ b/internal/web/service/host.go @@ -1,6 +1,7 @@ package service import ( + "encoding/json" "slices" "sort" "strconv" @@ -182,19 +183,51 @@ func buildHostRows(groupId string, req *entity.HostGroup) []*model.Host { return rows } +// A node's group is stored once per address and per adopted inbound, so the +// response cap does not bound the write: the group's size and a per-node row budget do. +const ( + maxAdoptedHostGroupBytes = 16 << 10 + maxAdoptedHostRows = 2000 +) + +// adoptableHostGroups drops the node-reported groups too large to adopt. +func adoptableHostGroups(groups []*entity.HostGroup) []*entity.HostGroup { + return slices.DeleteFunc(slices.Clone(groups), func(g *entity.HostGroup) bool { + if g == nil { + return true + } + wire, err := json.Marshal(g) + return err != nil || len(wire) > maxAdoptedHostGroupBytes + }) +} + +// adoptedHostRowBudget is how many more host rows a node's inbounds may adopt. +func adoptedHostRowBudget(db *gorm.DB, nodeID int) (int, error) { + var held int64 + err := db.Model(&model.Host{}). + Where("inbound_id IN (?)", db.Model(&model.Inbound{}).Select("id").Where("node_id = ?", nodeID)). + Count(&held).Error + return max(maxAdoptedHostRows-int(held), 0), err +} + // adoptedHostRows projects a node's host groups onto a freshly adopted central -// inbound so TLS/SNI/fingerprint overrides survive the node-to-master import. -func adoptedHostRows(groups []*entity.HostGroup, nodeInboundId, centralInboundId int) []*model.Host { +// inbound, building at most budget rows; it also returns how many groups that left out. +func adoptedHostRows(groups []*entity.HostGroup, nodeInboundId, centralInboundId, budget int) ([]*model.Host, int) { var rows []*model.Host + skipped := 0 for _, g := range groups { if g == nil || !slices.Contains(g.InboundIds, nodeInboundId) { continue } + if len(rows)+max(len(g.Hosts), 1) > budget { + skipped++ + continue + } scoped := *g scoped.InboundIds = []int{centralInboundId} rows = append(rows, buildHostRows(g.GroupId, &scoped)...) } - return rows + return rows, skipped } func validateInboundsExist(tx *gorm.DB, inboundIds []int) error { diff --git a/internal/web/service/inbound_node.go b/internal/web/service/inbound_node.go index 76c352389..3c543b68c 100644 --- a/internal/web/service/inbound_node.go +++ b/internal/web/service/inbound_node.go @@ -654,6 +654,16 @@ func (s *InboundService) setRemoteTrafficLocked(nodeID int, snap *runtime.Traffi } } + // The host rows the node's inbounds already hold count against this merge's adoption. + hostGroups := adoptableHostGroups(snap.HostGroups) + hostGroupsSkipped := len(snap.HostGroups) - len(hostGroups) + hostRowBudget := 0 + if len(hostGroups) > 0 { + if hostRowBudget, err = adoptedHostRowBudget(db, nodeID); err != nil { + return false, err + } + } + tx := db.Begin() committed := false defer func() { @@ -821,9 +831,13 @@ func (s *InboundService) setRemoteTrafficLocked(nodeID int, snap *runtime.Traffi if newIb.Tag != snapIb.Tag { tagToCentral[newIb.Tag] = &newIb } - if rows := adoptedHostRows(snap.HostGroups, snapIb.Id, newIb.Id); len(rows) > 0 { + rows, skipped := adoptedHostRows(hostGroups, snapIb.Id, newIb.Id, hostRowBudget) + hostGroupsSkipped += skipped + if len(rows) > 0 { if err := tx.Create(&rows).Error; err != nil { logger.Warningf("setRemoteTraffic: adopt host rows for tag %q failed: %v", newIb.Tag, err) + } else { + hostRowBudget -= len(rows) } } newInboundIDs[newIb.Id] = struct{}{} @@ -901,6 +915,13 @@ func (s *InboundService) setRemoteTrafficLocked(nodeID int, snap *runtime.Traffi } } + if hostGroupsSkipped > 0 { + logger.Warningf( + "setRemoteTraffic: node %d: left out %d host group adoption(s) — a group is limited to %d bytes and the node's inbounds to %d host rows", + nodeID, hostGroupsSkipped, maxAdoptedHostGroupBytes, maxAdoptedHostRows, + ) + } + // The inbound each reported email now lives under, so a swept inbound's // accumulator rows follow the email instead of being dropped and re-seeded at 0. snapEmailHome := make(map[string]int, len(snapEmailsAll)) diff --git a/internal/web/service/node_host_adopt_test.go b/internal/web/service/node_host_adopt_test.go index 53906744b..bac184472 100644 --- a/internal/web/service/node_host_adopt_test.go +++ b/internal/web/service/node_host_adopt_test.go @@ -1,6 +1,8 @@ package service import ( + "fmt" + "strings" "testing" "github.com/mhsanaei/3x-ui/v3/internal/database" @@ -124,3 +126,71 @@ func TestSetRemoteTraffic_SweptInboundTakesItsDependentsAlong(t *testing.T) { t.Fatalf("balancer members after the sweep = %v, want [%d]", balancer.InboundIds, kept.Id) } } + +// floodSnapshot reports inboundCount new inbounds (ids and ports from firstPort) +// and makes every group a member of all of them. +func floodSnapshot(tagPrefix string, firstPort, inboundCount int, groups ...*entity.HostGroup) *runtime.TrafficSnapshot { + snap := &runtime.TrafficSnapshot{HostGroups: groups} + ids := make([]int, inboundCount) + for i := range ids { + ids[i] = firstPort + i + snap.Inbounds = append(snap.Inbounds, &model.Inbound{ + Id: firstPort + i, Tag: fmt.Sprintf("%s-%d", tagPrefix, i), Enable: true, Port: firstPort + i, Protocol: model.VLESS, Settings: `{"clients":[]}`, + }) + } + for _, g := range groups { + g.InboundIds = ids + } + return snap +} + +func countHostRows(t *testing.T, query string, args ...any) int64 { + t.Helper() + var n int64 + if err := database.GetDB().Model(&model.Host{}).Where(query, args...).Count(&n).Error; err != nil { + t.Fatalf("count hosts: %v", err) + } + return n +} + +// A group is stored once per address and per inbound, so a small hosts/list answer +// must not be able to write an unbounded number of rows or bytes, in one merge or many. +func TestSetRemoteTraffic_BoundsAdoptedHostRows(t *testing.T) { + setupConflictDB(t) + + const nodeID = 7 + if err := database.GetDB().Create(&model.Node{ + Id: nodeID, Name: "flood-node", Address: "10.0.0.7", Port: 2053, ApiToken: "t", Guid: "flood-node-guid", + }).Error; err != nil { + t.Fatalf("create node: %v", err) + } + + addresses := make([]string, 50) + for i := range addresses { + addresses[i] = fmt.Sprintf("h%d.example", i) + } + snap := floodSnapshot("flood", 30000, 100, + &entity.HostGroup{GroupId: "g-big", Hosts: []string{"cdn.example"}, Remark: "big", Sni: strings.Repeat("s", 100_000)}, + &entity.HostGroup{GroupId: "g-many", Hosts: addresses, Remark: "many", Sni: "sni.example"}, + ) + + svc := InboundService{} + if _, err := svc.setRemoteTrafficLocked(nodeID, snap, false, false); err != nil { + t.Fatalf("setRemoteTrafficLocked: %v", err) + } + if big := countHostRows(t, "group_id = ?", "g-big"); big != 0 { + t.Fatalf("a 100000-byte group was stored %d times", big) + } + if total := countHostRows(t, "1 = 1"); total != maxAdoptedHostRows { + t.Fatalf("one merge adopted %d host rows, want %d", total, maxAdoptedHostRows) + } + + more := floodSnapshot("more", 31000, 10, &entity.HostGroup{GroupId: "g-more", Hosts: addresses, Remark: "more"}) + more.Inbounds = append(more.Inbounds, snap.Inbounds...) + if _, err := svc.setRemoteTrafficLocked(nodeID, more, false, false); err != nil { + t.Fatalf("second setRemoteTrafficLocked: %v", err) + } + if total := countHostRows(t, "1 = 1"); total != maxAdoptedHostRows { + t.Fatalf("a second merge raised the node's host rows to %d, want %d", total, maxAdoptedHostRows) + } +}