Fix: several latent bugs in entry caching, private input and output writers

Agent-Logs-Url: https://github.com/Loyalsoldier/geoip/sessions/051395fa-8215-474c-b4a5-423d84a33e0c

Co-authored-by: Loyalsoldier <10487845+Loyalsoldier@users.noreply.github.com>

Author: Claude

lib/container.go

@@ -93,11 +93,13 @@ func (c *container) Add(entry *Entry, opts ...IgnoreIPOption) error {
 				val.ipv6Builder = new(netipx.IPSetBuilder)
 			}
 			val.ipv6Builder.AddSet(ipv6set)
+			val.resetIPv6Set()
 		case IPv6:
 			if !val.hasIPv4Builder() {
 				val.ipv4Builder = new(netipx.IPSetBuilder)
 			}
 			val.ipv4Builder.AddSet(ipv4set)
+			val.resetIPv4Set()
 		default:
 			if !val.hasIPv4Builder() {
 				val.ipv4Builder = new(netipx.IPSetBuilder)
@@ -107,14 +109,18 @@ func (c *container) Add(entry *Entry, opts ...IgnoreIPOption) error {
 			}
 			val.ipv4Builder.AddSet(ipv4set)
 			val.ipv6Builder.AddSet(ipv6set)
+			val.resetIPv4Set()
+			val.resetIPv6Set()
 		}
 
 	case false:
 		switch ignoreIPType {
 		case IPv4:
 			entry.ipv4Builder = nil
+			entry.resetIPv4Set()
 		case IPv6:
 			entry.ipv6Builder = nil
+			entry.resetIPv6Set()
 		}
 		c.entries[name] = entry
 	}
@@ -159,11 +165,13 @@ func (c *container) Remove(entry *Entry, rCase CaseRemove, opts ...IgnoreIPOptio
 				val.ipv6Builder = new(netipx.IPSetBuilder)
 			}
 			val.ipv6Builder.RemoveSet(ipv6set)
+			val.resetIPv6Set()
 		case IPv6:
 			if !val.hasIPv4Builder() {
 				val.ipv4Builder = new(netipx.IPSetBuilder)
 			}
 			val.ipv4Builder.RemoveSet(ipv4set)
+			val.resetIPv4Set()
 		default:
 			if !val.hasIPv4Builder() {
 				val.ipv4Builder = new(netipx.IPSetBuilder)
@@ -173,14 +181,18 @@ func (c *container) Remove(entry *Entry, rCase CaseRemove, opts ...IgnoreIPOptio
 			}
 			val.ipv4Builder.RemoveSet(ipv4set)
 			val.ipv6Builder.RemoveSet(ipv6set)
+			val.resetIPv4Set()
+			val.resetIPv6Set()
 		}
 
 	case CaseRemoveEntry:
 		switch ignoreIPType {
 		case IPv4:
 			val.ipv6Builder = nil
+			val.resetIPv6Set()
 		case IPv6:
 			val.ipv4Builder = nil
+			val.resetIPv4Set()
 		default:
 			delete(c.entries, name)
 		}

lib/entry.go

@@ -43,6 +43,18 @@ func (e *Entry) hasIPv6Set() bool {
 	return e.ipv6Set != nil
 }
 
+// resetIPv4Set discards the cached IPv4 set, so that it will be rebuilt
+// from the builder on the next use.
+func (e *Entry) resetIPv4Set() {
+	e.ipv4Set = nil
+}
+
+// resetIPv6Set discards the cached IPv6 set, so that it will be rebuilt
+// from the builder on the next use.
+func (e *Entry) resetIPv6Set() {
+	e.ipv6Set = nil
+}
+
 func (e *Entry) GetIPv4Set() (*netipx.IPSet, error) {
 	if err := e.buildIPSet(); err != nil {
 		return nil, err
@@ -253,11 +265,13 @@ func (e *Entry) add(prefix *netip.Prefix, ipType IPType) error {
 			e.ipv4Builder = new(netipx.IPSetBuilder)
 		}
 		e.ipv4Builder.AddPrefix(*prefix)
+		e.resetIPv4Set()
 	case IPv6:
 		if !e.hasIPv6Builder() {
 			e.ipv6Builder = new(netipx.IPSetBuilder)
 		}
 		e.ipv6Builder.AddPrefix(*prefix)
+		e.resetIPv6Set()
 	default:
 		return ErrInvalidIPType
 	}
@@ -270,10 +284,12 @@ func (e *Entry) remove(prefix *netip.Prefix, ipType IPType) error {
 	case IPv4:
 		if e.hasIPv4Builder() {
 			e.ipv4Builder.RemovePrefix(*prefix)
+			e.resetIPv4Set()
 		}
 	case IPv6:
 		if e.hasIPv6Builder() {
 			e.ipv6Builder.RemovePrefix(*prefix)
+			e.resetIPv6Set()
 		}
 	default:
 		return ErrInvalidIPType

plugin/maxmind/maxmind_country_csv_in.go

@@ -167,6 +167,10 @@ func (g *GeoLite2CountryCSVIn) getCountryCode() (map[string]string, error) {
 	}
 
 	ccMap := make(map[string]string)
+	if len(lines) == 0 {
+		return nil, fmt.Errorf("❌ [type %s | action %s] empty country code file: %s", g.Type, g.Action, g.CountryCodeFile)
+	}
+
 	for _, line := range lines[1:] {
 		if len(line) < 5 {
 			return nil, fmt.Errorf("❌ [type %s | action %s] invalid record: %v", g.Type, g.Action, line)

plugin/mihomo/mrs_out.go

@@ -164,13 +164,18 @@ func (m *MRSOut) writeFile(filename string, ipRanges []netipx.IPRange) error {
 	if err != nil {
 		return err
 	}
+
 	defer f.Close()
 
 	err = m.convertToMrs(ipRanges, f)
 	if err != nil {
 		return err
 	}
 
+	if err := f.Close(); err != nil {
+		return err
+	}
+
 	log.Printf("✅ [%s] %s --> %s", m.Type, filename, m.OutputDir)
 
 	return nil
@@ -181,7 +186,11 @@ func (m *MRSOut) convertToMrs(ipRanges []netipx.IPRange, w io.Writer) (err error
 	if err != nil {
 		return err
 	}
-	defer encoder.Close()
+	defer func() {
+		if closeErr := encoder.Close(); closeErr != nil && err == nil {
+			err = closeErr
+		}
+	}()
 
 	// header
 	_, err = encoder.Write(mrsMagicBytes[:])

plugin/singbox/srs_out.go

@@ -190,6 +190,10 @@ func (s *SRSOut) writeFile(filename string, ruleset *option.PlainRuleSet) error
 		return err
 	}
 
+	if err := f.Close(); err != nil {
+		return err
+	}
+
 	log.Printf("✅ [%s] %s --> %s", s.Type, filename, s.OutputDir)
 
 	return nil

plugin/special/private.go

@@ -84,10 +84,7 @@ func (p *Private) GetDescription() string {
 }
 
 func (p *Private) Input(container lib.Container) (lib.Container, error) {
-	entry, found := container.GetEntry(entryNamePrivate)
-	if !found {
-		entry = lib.NewEntry(entryNamePrivate)
-	}
+	entry := lib.NewEntry(entryNamePrivate)
 
 	for _, cidr := range privateCIDRs {
 		if err := entry.AddPrefix(cidr); err != nil {