refactor IPSet

This commit is contained in:
Vladimir Avtsenov 2025-02-11 15:29:26 +03:00
parent 066eeb0ab7
commit 5fd28ae005
3 changed files with 15 additions and 15 deletions

View File

@ -21,17 +21,17 @@ type Group struct {
ipsetToLink *netfilterHelper.IPSetToLink ipsetToLink *netfilterHelper.IPSetToLink
} }
func (g *Group) AddIPv4(address net.IP, ttl time.Duration) error { func (g *Group) AddIP(address net.IP, ttl time.Duration) error {
ttlSeconds := uint32(ttl.Seconds()) ttlSeconds := uint32(ttl.Seconds())
return g.ipset.AddIP(address, &ttlSeconds) return g.ipset.AddIP(address, &ttlSeconds)
} }
func (g *Group) DelIPv4(address net.IP) error { func (g *Group) DelIP(address net.IP) error {
return g.ipset.Del(address) return g.ipset.DelIP(address)
} }
func (g *Group) ListIPv4() (map[string]*uint32, error) { func (g *Group) ListIP() (map[string]*uint32, error) {
return g.ipset.List() return g.ipset.ListIPs()
} }
func (g *Group) Enable() error { func (g *Group) Enable() error {

View File

@ -326,7 +326,7 @@ func (a *App) SyncGroup(group *Group) error {
} }
} }
currentAddresses, err := group.ListIPv4() currentAddresses, err := group.ListIP()
if err != nil { if err != nil {
return fmt.Errorf("failed to get old ipset list: %w", err) return fmt.Errorf("failed to get old ipset list: %w", err)
} }
@ -337,7 +337,7 @@ func (a *App) SyncGroup(group *Group) error {
continue continue
} }
ip := net.IP(addr) ip := net.IP(addr)
err = group.AddIPv4(ip, ttl) err = group.AddIP(ip, ttl)
if err != nil { if err != nil {
log.Error(). log.Error().
Str("address", ip.String()). Str("address", ip.String()).
@ -356,7 +356,7 @@ func (a *App) SyncGroup(group *Group) error {
continue continue
} }
ip := net.IP(addr) ip := net.IP(addr)
err = group.DelIPv4(ip) err = group.DelIP(ip)
if err != nil { if err != nil {
log.Error(). log.Error().
Str("address", ip.String()). Str("address", ip.String()).
@ -418,7 +418,7 @@ func (a *App) processARecord(aRecord dns.A) {
continue continue
} }
// TODO: Check already existed // TODO: Check already existed
err := group.AddIPv4(aRecord.A, ttlDuration) err := group.AddIP(aRecord.A, ttlDuration)
if err != nil { if err != nil {
log.Error(). log.Error().
Str("address", aRecord.A.String()). Str("address", aRecord.A.String()).
@ -467,7 +467,7 @@ func (a *App) processCNameRecord(cNameRecord dns.CNAME) {
continue continue
} }
for _, aRecord := range aRecords { for _, aRecord := range aRecords {
err := group.AddIPv4(aRecord.Address, now.Sub(aRecord.Deadline)) err := group.AddIP(aRecord.Address, now.Sub(aRecord.Deadline))
if err != nil { if err != nil {
log.Error(). log.Error().
Str("address", aRecord.Address.String()). Str("address", aRecord.Address.String()).

View File

@ -2,9 +2,10 @@ package netfilterHelper
import ( import (
"fmt" "fmt"
"github.com/vishvananda/netlink"
"net" "net"
"os" "os"
"github.com/vishvananda/netlink"
) )
type IPSet struct { type IPSet struct {
@ -23,7 +24,7 @@ func (r *IPSet) AddIP(addr net.IP, timeout *uint32) error {
return nil return nil
} }
func (r *IPSet) Del(addr net.IP) error { func (r *IPSet) DelIP(addr net.IP) error {
err := netlink.IpsetDel(r.SetName, &netlink.IPSetEntry{ err := netlink.IpsetDel(r.SetName, &netlink.IPSetEntry{
IP: addr, IP: addr,
}) })
@ -33,7 +34,7 @@ func (r *IPSet) Del(addr net.IP) error {
return nil return nil
} }
func (r *IPSet) List() (map[string]*uint32, error) { func (r *IPSet) ListIPs() (map[string]*uint32, error) {
list, err := netlink.IpsetList(r.SetName) list, err := netlink.IpsetList(r.SetName)
if err != nil { if err != nil {
return nil, err return nil, err
@ -62,9 +63,8 @@ func (nh *NetfilterHelper) IPSet(name string) (*IPSet, error) {
return nil, err return nil, err
} }
defaultTimeout := uint32(300)
err = netlink.IpsetCreate(ipset.SetName, "hash:net", netlink.IpsetCreateOptions{ err = netlink.IpsetCreate(ipset.SetName, "hash:net", netlink.IpsetCreateOptions{
Timeout: &defaultTimeout, Timeout: func(i uint32) *uint32 { return &i }(300),
}) })
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to create ipset: %w", err) return nil, fmt.Errorf("failed to create ipset: %w", err)