From abf8977f55b1be7b51140cfc9bfb402e40efe7c0 Mon Sep 17 00:00:00 2001
From: Stanislav Chzhen <s.chzhen@adguard.com>
Date: Wed, 29 Nov 2023 15:29:36 +0300
Subject: [PATCH] ipset: header data query

---
 internal/ipset/ipset_linux.go               | 42 ++++++++++++++++++++-
 internal/ipset/ipset_linux_internal_test.go |  5 +++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/internal/ipset/ipset_linux.go b/internal/ipset/ipset_linux.go
index 68f53dcb40b..92ecd1d6d73 100644
--- a/internal/ipset/ipset_linux.go
+++ b/internal/ipset/ipset_linux.go
@@ -101,6 +101,7 @@ func (qc *queryConn) listAll() (sets []props, err error) {
 type ipsetConn interface {
 	Add(name string, entries ...*ipset.Entry) (err error)
 	Close() (err error)
+	Header(name string) (p *ipset.HeaderPolicy, err error)
 	listAll() (sets []props, err error)
 }
 
@@ -112,6 +113,9 @@ type props struct {
 	// name of the ipset.
 	name string
 
+	// typeName of the ipset.
+	typeName string
+
 	// family of the IP addresses in the ipset.
 	family netfilter.ProtoFamily
 
@@ -148,6 +152,8 @@ func (p *props) parseAttribute(a netfilter.Attribute) {
 	case ipset.AttrSetName:
 		// Trim the null character.
 		p.name = string(bytes.Trim(a.Data, "\x00"))
+	case ipset.AttrTypeName:
+		p.typeName = string(bytes.Trim(a.Data, "\x00"))
 	case ipset.AttrFamily:
 		p.family = netfilter.ProtoFamily(a.Data[0])
 	default:
@@ -288,6 +294,35 @@ func (m *manager) parseIpsetConfig(ipsetConf []string) (err error) {
 	return nil
 }
 
+// ipsetProps returns the properties of an ipset with the given name.
+//
+// Additional header data query.  See
+// https://github.com/AdguardTeam/AdGuardHome/issues/6420.
+func (m *manager) ipsetProps(p props) (err error) {
+	// The family doesn't seem to matter when we use a header query, so
+	// query only the IPv4 one.
+	//
+	// TODO(a.garipov): Find out if this is a bug or a feature.
+	var res *ipset.HeaderPolicy
+	res, err = m.ipv4Conn.Header(p.name)
+	if err != nil {
+		return err
+	}
+
+	if res == nil || res.Family == nil {
+		return errors.Error("empty response or no family data")
+	}
+
+	family := netfilter.ProtoFamily(res.Family.Value)
+	if family != netfilter.ProtoIPv4 && family != netfilter.ProtoIPv6 {
+		return fmt.Errorf("unexpected ipset family %q", family)
+	}
+
+	p.family = family
+
+	return nil
+}
+
 // ipsets returns currently known ipsets.
 func (m *manager) ipsets(names []string) (sets []props, err error) {
 	for _, n := range names {
@@ -297,7 +332,12 @@ func (m *manager) ipsets(names []string) (sets []props, err error) {
 		}
 
 		if p.family != netfilter.ProtoIPv4 && p.family != netfilter.ProtoIPv6 {
-			return nil, fmt.Errorf("%q unexpected ipset family %q", p.name, p.family)
+			log.Debug("%q %q unexpected ipset family %q", p.name, p.typeName, p.family)
+
+			err = m.ipsetProps(p)
+			if err != nil {
+				return nil, fmt.Errorf("%q %q making header query: %w", p.name, p.typeName, err)
+			}
 		}
 
 		sets = append(sets, p)
diff --git a/internal/ipset/ipset_linux_internal_test.go b/internal/ipset/ipset_linux_internal_test.go
index 84e2565083b..f22d93c151c 100644
--- a/internal/ipset/ipset_linux_internal_test.go
+++ b/internal/ipset/ipset_linux_internal_test.go
@@ -47,6 +47,11 @@ func (c *fakeConn) Close() (err error) {
 	return nil
 }
 
+// Header implements the [ipsetConn] interface for *fakeConn.
+func (c *fakeConn) Header(_ string) (_ *ipset.HeaderPolicy, _ error) {
+	return nil, nil
+}
+
 // listAll implements the [ipsetConn] interface for *fakeConn.
 func (c *fakeConn) listAll() (sets []props, err error) {
 	return c.sets, nil