gitee.com/aurawing/surguard-go@v0.3.1-0.20240409071558-96509a61ecf3/device/postconfig_linux.go

package device

import (
	"bufio"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/coreos/go-iptables/iptables"
	"github.com/jviney/go-filelock"
	sysctl "github.com/lorenzosaino/go-sysctl"
)

const (
	ENV_SG_FWMARK        = "SG_FWMARK"
	ENV_SG_RT_ID         = "SG_RT_ID"
	ENV_SG_TUNNEL_FWMARK = "SG_TUNNEL_FWMARK"
)

var fwmark = 1024
var tunnelFwmark = 65535

var ip4t *iptables.IPTables
var deviceName string
var lock filelock.FileLock

func init() {
	lockfile := "/var/lock/surguard.lock"
	if _, err := os.Stat(lockfile); err != nil {
		if os.IsNotExist(err) {
			var file, err = os.Create(lockfile)
			if err != nil {
				panic(err)
			}
			file.Close()
		} else {
			panic(err)
		}
	}
	lock = filelock.FileLock{Path: lockfile, Timeout: time.Second * 5}
}

func GetDeviceIndex() int {
	return 0
}

func (device *Device) configTunDevice() error {
	ifBindInterface = false
	var err error
	//find tun device name
	deviceName, err = device.tun.device.Name()
	if err != nil {
		device.log.Errorf("configTunDevice: failed to find tun device name: %s", err)
		return err
	}

	//bring up tun device
	device.runCmd(false, "ip", "link", "set", "dev", deviceName, "up")
	for {
		if deviceState(device.state.state.Load()) == deviceStateUp {
			break
		} else {
			time.Sleep(500 * time.Millisecond)
		}
	}
	device.log.Verbosef("configTunDevice: device %s is up\n", deviceName)

	//close rp_filter of tun device
	rp_filter_all, err := sysctl.Get("net.ipv4.conf.all.rp_filter")
	if err != nil {
		device.log.Errorf("configTunDevice: failed to get net.ipv4.conf.all.rp_filter: %s", err)
		return err
	}
	if rp_filter_all != "0" {
		err = sysctl.Set("net.ipv4.conf.all.rp_filter", "0")
		if err != nil {
			device.log.Errorf("configTunDevice: failed to disable net.ipv4.conf.all.rp_filter: %s", err)
			return err
		}
		device.log.Verbosef("configTunDevice: set net.ipv4.conf.all.rp_filter to 0\n")
	}
	rp_filter_sg, err := sysctl.Get(fmt.Sprintf("net.ipv4.conf.%s.rp_filter", deviceName))
	if err != nil {
		device.log.Errorf("configTunDevice: failed to get net.ipv4.conf.%s.rp_filter: %s", deviceName, err)
		return err
	}
	if rp_filter_sg != "0" {
		err = sysctl.Set(fmt.Sprintf("net.ipv4.conf.%s.rp_filter", deviceName), "0")
		if err != nil {
			device.log.Errorf("configTunDevice: failed to disable net.ipv4.conf.%s.rp_filter: %s", deviceName, err)
			return err
		}
		device.log.Verbosef("configTunDevice: set net.ipv4.conf.%s.rp_filter to 0\n", deviceName)
	}

	//get fwmark and route table id
	rtid := 256
	fwmarkStr := os.Getenv(ENV_SG_FWMARK)
	tunnelFwmarkStr := os.Getenv(ENV_SG_TUNNEL_FWMARK)

	if fwmarkStr != "" {
		fwmark, err = strconv.Atoi(fwmarkStr)
		if err != nil {
			device.log.Errorf("configTunDevice: failed to parse fwmark %s: %s", fwmarkStr, err)
		}
	}
	if tunnelFwmarkStr != "" {
		tunnelFwmark, err = strconv.Atoi(tunnelFwmarkStr)
		if err != nil {
			device.log.Errorf("configTunDevice: failed to parse tunnel fwmark %s: %s", fwmarkStr, err)
		}
	}
	rtidStr := os.Getenv(ENV_SG_RT_ID)
	if rtidStr != "" {
		rtid, err = strconv.Atoi(rtidStr)
		if err != nil {
			device.log.Errorf("configTunDevice: failed to parse route table ID %s: %s", rtidStr, err)
		}
	}

	//modify route table
	err = device.modifyRtTables("/etc/iproute2/rt_tables", rtid)
	if err != nil {
		device.log.Errorf("configTunDevice: failed to modify route table ID %d: %s", rtid, err)
		return err
	}

	// set route policy
	device.runCmd(false, "ip", "route", "flush", "table", groupName)
	device.runCmd(false, "ip", "route", "add", "default", "dev", deviceName, "table", groupName)
	device.runCmd(true, "ip", "rule", "delete", "from", "all", "fwmark", strconv.Itoa(fwmark), "lookup", groupName)
	device.runCmd(false, "ip", "rule", "add", "from", "all", "fwmark", strconv.Itoa(fwmark), "lookup", groupName)
	device.log.Verbosef("configTunDevice: set route policy table %d %s for %s\n", rtid, groupName, deviceName)

	// //set up iptables
	// ip4t, err = iptables.NewWithProtocol(iptables.ProtocolIPv4)
	// if err != nil {
	// 	device.log.Errorf("PostConfig: failed to create new iptables instance: %s", err)
	// 	return err
	// }
	// exists, err := ip4t.ChainExists("mangle", groupName)
	// if err != nil {
	// 	device.log.Errorf("PostConfig: check chain %s failed: %s", groupName, err)
	// 	return err
	// }
	// if !exists {
	// 	err = ip4t.NewChain("mangle", groupName)
	// 	if err != nil {
	// 		device.log.Errorf("PostConfig: create chain %s failed: %s", groupName, err)
	// 		return err
	// 	}
	// }
	// err = ip4t.ClearChain("mangle", groupName)
	// if err != nil {
	// 	device.log.Errorf("PostConfig: clean chain %s failed: %s", groupName, err)
	// 	return err
	// }
	// err = ip4t.AppendUnique("mangle", "OUTPUT", "-j", groupName)
	// if err != nil {
	// 	device.log.Errorf("PostConfig: append chain to %s failed: %s", groupName, err)
	// 	return err
	// }
	return nil
}

func (device *Device) initRules() error {
	// //find tun device name
	// deviceName, err := device.tun.device.Name()
	// if err != nil {
	// 	device.log.Errorf("initRules: failed to find tun device name: %s", err)
	// 	return err
	// }
	//set up iptables
	err := lock.Lock()
	if err != nil {
		device.log.Errorf("initRules: accquire global lock failed: %s", err)
		return err
	}
	defer lock.Unlock()
	ip4t, err = iptables.NewWithProtocol(iptables.ProtocolIPv4)
	if err != nil {
		device.log.Errorf("initRules: failed to create new iptables instance: %s", err)
		return err
	}
	//initialize IN-SURGUARD chain
	exists, err := ip4t.ChainExists("mangle", "IN-SURGUARD")
	if err != nil {
		device.log.Errorf("initRules: check chain IN-SURGUARD failed: %s", err)
		return err
	}
	if !exists {
		err = ip4t.NewChain("mangle", "IN-SURGUARD")
		if err != nil {
			device.log.Errorf("initRules: create chain IN-SURGUARD failed: %s", err)
			return err
		}
	}
	exists, err = ip4t.Exists("mangle", "PREROUTING", "-j", "IN-SURGUARD")
	if err != nil {
		device.log.Errorf("initRules: check if rule IN-SURGUARD in chain PREROUTING failed: %s", err)
		return err
	}
	if !exists {
		err = ip4t.AppendUnique("mangle", "PREROUTING", "-j", "IN-SURGUARD")
		if err != nil {
			device.log.Errorf("initRules: add chain reference IN-SURGUARD to chain PREROUTING failed: %s", err)
			return err
		}
	}
	// exists, err = ip4t.ChainExists("mangle", fmt.Sprintf("IN-%s", groupName))
	// if err != nil {
	// 	device.log.Errorf("initRules: check chain %s failed: %s", fmt.Sprintf("IN-%s", groupName), err)
	// 	return err
	// }
	// if !exists {
	// 	err = ip4t.NewChain("mangle", fmt.Sprintf("IN-%s", groupName))
	// 	if err != nil {
	// 		device.log.Errorf("initRules: create chain %s failed: %s", fmt.Sprintf("IN-%s", groupName), err)
	// 		return err
	// 	}
	// }
	// err = ip4t.ClearChain("mangle", fmt.Sprintf("IN-%s", groupName))
	// if err != nil {
	// 	device.log.Errorf("initRules: clean chain %s failed: %s", fmt.Sprintf("IN-%s", groupName), err)
	// 	return err
	// }
	exists, err = ip4t.Exists("mangle", "IN-SURGUARD", "-i", deviceName, "-j", "CONNMARK", "--set-mark", strconv.Itoa(fwmark))
	if err != nil {
		device.log.Errorf("initRules: checking if connmark setting rule in chain IN-SURGUARD failed: %s", err)
		return err
	}
	if !exists {
		err = ip4t.AppendUnique("mangle", "IN-SURGUARD", "-i", deviceName, "-j", "CONNMARK", "--set-mark", strconv.Itoa(fwmark))
		if err != nil {
			device.log.Errorf("initRules: create connmark setting rule in chain IN-SURGUARD failed: %s", err)
			return err
		}
	}
	// err = ip4t.AppendUnique("mangle", fmt.Sprintf("IN-%s", groupName), "-i", deviceName, "-j", "CONNMARK", "--set-mark", strconv.Itoa(fwmark))
	// if err != nil {
	// 	device.log.Errorf("initRules: create connmark setting rule in chain %s failed: %s", fmt.Sprintf("IN-%s", groupName), err)
	// 	return err
	// }
	// err = ip4t.AppendUnique("mangle", fmt.Sprintf("IN-%s", groupName), "-m", "connmark", "!", "--mark", "0", "-j", "RETURN")
	// if err != nil {
	// 	device.log.Errorf("initRules: create connmark checking rule in chain %s failed: %s", fmt.Sprintf("IN-%s", groupName), err)
	// 	return err
	// }
	// err = ip4t.AppendUnique("mangle", fmt.Sprintf("IN-%s", groupName), "-p", "udp", "--dport", strconv.Itoa(int(device.net.port)), "-j", "CONNMARK", "--set-mark", strconv.Itoa(tunnelFwmark))
	// if err != nil {
	// 	device.log.Errorf("initRules: create connmark setting rule in chain %s failed: %s", fmt.Sprintf("IN-%s", groupName), err)
	// 	return err
	// }
	// err = ip4t.Append("mangle", fmt.Sprintf("IN-%s", groupName), "-m", "connmark", "!", "--mark", "0", "-j", "RETURN")
	// if err != nil {
	// 	device.log.Errorf("initRules: create connmark checking rule in chain %s failed: %s", fmt.Sprintf("IN-%s", groupName), err)
	// 	return err
	// }
	//initialize PRE-SURGUARD chain
	exists, err = ip4t.ChainExists("mangle", "PRE-SURGUARD")
	if err != nil {
		device.log.Errorf("initRules: check chain PRE-SURGUARD failed: %s", err)
		return err
	}
	if !exists {
		err = ip4t.NewChain("mangle", "PRE-SURGUARD")
		if err != nil {
			device.log.Errorf("initRules: create chain PRE-SURGUARD failed: %s", err)
			return err
		}
	}
	exists, err = ip4t.Exists("mangle", "OUTPUT", "-j", "PRE-SURGUARD")
	if err != nil {
		device.log.Errorf("initRules: check if rule PRE-SURGUARD in chain OUTPUT failed: %s", err)
		return err
	}
	if !exists {
		err = ip4t.AppendUnique("mangle", "OUTPUT", "-j", "PRE-SURGUARD")
		if err != nil {
			device.log.Errorf("initRules: add chain reference PRE-SURGUARD to chain OUTPUT failed: %s", err)
			return err
		}
	}
	// exists, err = ip4t.Exists("mangle", "FORWARD", "-j", "PRE-SURGUARD")
	// if err != nil {
	// 	device.log.Errorf("initRules: check if rule PRE-SURGUARD in chain FORWARD failed: %s", err)
	// 	return err
	// }
	// if !exists {
	// 	err = ip4t.AppendUnique("mangle", "FORWARD", "-j", "PRE-SURGUARD")
	// 	if err != nil {
	// 		device.log.Errorf("initRules: add chain reference PRE-SURGUARD to chain FORWARD failed: %s", err)
	// 		return err
	// 	}
	// }
	exists, err = ip4t.Exists("mangle", "PRE-SURGUARD", "-j", "CONNMARK", "--restore-mark")
	if err != nil {
		device.log.Errorf("initRules: check if rule restore mark in chain PRE-SURGUARD failed: %s", err)
		return err
	}
	if !exists {
		err = ip4t.AppendUnique("mangle", "PRE-SURGUARD", "-j", "CONNMARK", "--restore-mark")
		if err != nil {
			device.log.Errorf("initRules: add rule restore mark in chain PRE-SURGUARD failed: %s", err)
			return err
		}
	}
	exists, err = ip4t.Exists("mangle", "PRE-SURGUARD", "-m", "mark", "!", "--mark", "0", "-j", "ACCEPT")
	if err != nil {
		device.log.Errorf("initRules: check if mark checking rule in chain PRE-SURGUARD failed: %s", err)
		return err
	}
	if !exists {
		err = ip4t.AppendUnique("mangle", "PRE-SURGUARD", "-m", "mark", "!", "--mark", "0", "-j", "ACCEPT")
		if err != nil {
			device.log.Errorf("initRules: create mark checking rule in chain PRE-SURGUARD failed: %s", err)
			return err
		}
	}
	exists, err = ip4t.Exists("mangle", "PRE-SURGUARD", "-m", "conntrack", "--ctstate", "RELATED,ESTABLISHED", "-j", "ACCEPT")
	if err != nil {
		device.log.Errorf("initRules: check if releated connection passing rule in chain PRE-SURGUARD failed: %s", err)
		return err
	}
	if !exists {
		err = ip4t.AppendUnique("mangle", "PRE-SURGUARD", "-m", "conntrack", "--ctstate", "RELATED,ESTABLISHED", "-j", "ACCEPT")
		if err != nil {
			device.log.Errorf("initRules: create releated connection passing rule in chain PRE-SURGUARD failed: %s", err)
			return err
		}
	}
	//initialize OUT-SURGUARD chain
	exists, err = ip4t.ChainExists("mangle", "OUT-SURGUARD")
	if err != nil {
		device.log.Errorf("initRules: check chain OUT-SURGUARD failed: %s", err)
		return err
	}
	if !exists {
		err = ip4t.NewChain("mangle", "OUT-SURGUARD")
		if err != nil {
			device.log.Errorf("initRules: create chain OUT-SURGUARD failed: %s", err)
			return err
		}
	}
	exists, err = ip4t.Exists("mangle", "OUTPUT", "-j", "OUT-SURGUARD")
	if err != nil {
		device.log.Errorf("initRules: check if rule OUT-SURGUARD in chain OUTPUT failed: %s", err)
		return err
	}
	if !exists {
		err = ip4t.AppendUnique("mangle", "OUTPUT", "-j", "OUT-SURGUARD")
		if err != nil {
			device.log.Errorf("initRules: add chain reference OUT-SURGUARD to chain OUTPUT failed: %s", err)
			return err
		}
	}
	// exists, err = ip4t.Exists("mangle", "FORWARD", "-j", "OUT-SURGUARD")
	// if err != nil {
	// 	device.log.Errorf("initRules: check if rule OUT-SURGUARD in chain FORWARD failed: %s", err)
	// 	return err
	// }
	// if !exists {
	// 	err = ip4t.AppendUnique("mangle", "FORWARD", "-j", "OUT-SURGUARD")
	// 	if err != nil {
	// 		device.log.Errorf("initRules: add chain reference OUT-SURGUARD to chain FORWARD failed: %s", err)
	// 		return err
	// 	}
	// }
	exists, err = ip4t.ChainExists("mangle", fmt.Sprintf("OUT-%s", groupName))
	if err != nil {
		device.log.Errorf("initRules: check chain %s failed: %s", fmt.Sprintf("OUT-%s", groupName), err)
		return err
	}
	if !exists {
		err = ip4t.NewChain("mangle", fmt.Sprintf("OUT-%s", groupName))
		if err != nil {
			device.log.Errorf("initRules: create chain %s failed: %s", fmt.Sprintf("OUT-%s", groupName), err)
			return err
		}
	}
	err = ip4t.ClearChain("mangle", fmt.Sprintf("OUT-%s", groupName))
	if err != nil {
		device.log.Errorf("initRules: clean chain %s failed: %s", fmt.Sprintf("OUT-%s", groupName), err)
		return err
	}
	exists, err = ip4t.Exists("mangle", "OUT-SURGUARD", "-j", fmt.Sprintf("OUT-%s", groupName))
	if err != nil {
		device.log.Errorf("initRules: check if rule %s in chain OUT-SURGUARD failed: %s", fmt.Sprintf("OUT-%s", groupName), err)
		return err
	}
	if !exists {
		err = ip4t.AppendUnique("mangle", "OUT-SURGUARD", "-j", fmt.Sprintf("OUT-%s", groupName))
		if err != nil {
			device.log.Errorf("initRules: add chain reference %s to chain OUT-SURGUARD failed: %s", fmt.Sprintf("OUT-%s", groupName), err)
			return err
		}
	}
	return nil
}

func (device *Device) zkChangedCallback(addedMap, deletedMap map[string]string) {
	for _, dels := range deletedMap {
		flag := true
		for _, del := range strings.Split(dels, ",") {
			strs := strings.Split(del, ":")
			ip4t.Delete("mangle", fmt.Sprintf("OUT-%s", groupName), "-d", strs[0], "-j", "MARK", "--set-mark", strconv.Itoa(fwmark))
			if flag {
				ip4t.Delete("mangle", fmt.Sprintf("OUT-%s", groupName), "-p", "udp", "-d", strs[0], "--dport", strs[1], "-j", "MARK", "--set-mark", strconv.Itoa(tunnelFwmark))
				flag = false
			}
			ip4t.Delete("mangle", fmt.Sprintf("OUT-%s", groupName), "-d", strs[0], "-m", "mark", "!", "--mark", "0", "-j", "RETURN")
		}
	}
	for _, adds := range addedMap {
		flag := true
		for _, add := range strings.Split(adds, ",") {
			strs := strings.Split(add, ":")
			ip4t.AppendUnique("mangle", fmt.Sprintf("OUT-%s", groupName), "-d", strs[0], "-j", "MARK", "--set-mark", strconv.Itoa(fwmark))
			if flag {
				ip4t.AppendUnique("mangle", fmt.Sprintf("OUT-%s", groupName), "-p", "udp", "-d", strs[0], "--dport", strs[1], "-j", "MARK", "--set-mark", strconv.Itoa(tunnelFwmark))
				flag = false
			}
			ip4t.AppendUnique("mangle", fmt.Sprintf("OUT-%s", groupName), "-d", strs[0], "-m", "mark", "!", "--mark", "0", "-j", "RETURN")
		}
	}
}

func (device *Device) clearConfigOSSpecific() error {
	err := lock.Lock()
	if err != nil {
		device.log.Errorf("clearConfigOSSpecific: accquire global lock failed: %s", err)
		return err
	}
	defer lock.Unlock()
	device.runCmd(false, "ip", "route", "flush", "table", groupName)
	device.runCmd(true, "ip", "rule", "delete", "from", "all", "fwmark", strconv.Itoa(fwmark), "lookup", groupName)
	// err := ip4t.ClearChain("mangle", fmt.Sprintf("IN-%s", groupName))
	// if err != nil {
	// 	device.log.Errorf("PostConfig: clean chain %s failed: %s", fmt.Sprintf("IN-%s", groupName), err)
	// 	return err
	// }

	err = ip4t.Delete("mangle", "IN-SURGUARD", "-i", deviceName, "-j", "CONNMARK", "--set-mark", strconv.Itoa(fwmark))
	if err != nil {
		device.log.Errorf("clearConfigOSSpecific: delete connmark setting rule in chain IN-SURGUARD failed: %s", err)
		return err
	}
	// err = ip4t.DeleteChain("mangle", fmt.Sprintf("IN-%s", groupName))
	// if err != nil {
	// 	device.log.Errorf("PostConfig: delete chain %s failed: %s", fmt.Sprintf("IN-%s", groupName), err)
	// 	return err
	// }
	err = ip4t.ClearChain("mangle", fmt.Sprintf("OUT-%s", groupName))
	if err != nil {
		device.log.Errorf("clearConfigOSSpecific: clean chain %s failed: %s", fmt.Sprintf("OUT-%s", groupName), err)
		return err
	}
	err = ip4t.Delete("mangle", "OUT-SURGUARD", "-j", fmt.Sprintf("OUT-%s", groupName))
	if err != nil {
		device.log.Errorf("clearConfigOSSpecific: delete chain reference %s in OUT-SURGUARD chain failed: %s", fmt.Sprintf("OUT-%s", groupName), err)
		return err
	}
	err = ip4t.DeleteChain("mangle", fmt.Sprintf("OUT-%s", groupName))
	if err != nil {
		device.log.Errorf("clearConfigOSSpecific: delete chain %s failed: %s", fmt.Sprintf("OUT-%s", groupName), err)
		return err
	}
	return nil
}

func (device *Device) modifyRtTables(filename string, rtid int) error {
	//lock := filelock.FileLock{Path: filename, Timeout: time.Second * 5}
	err := lock.Lock()
	if err != nil {
		device.log.Errorf("modifyRtTables: accquire global lock failed: %s", err)
		return err
	}
	defer lock.Unlock()
	output := make([]string, 0)

	fileData, err := os.ReadFile(filename)
	if err != nil {
		device.log.Errorf("modifyRtTables: can not read file %s", filename)
		return err
	}

	scanner := bufio.NewScanner(strings.NewReader(string(fileData)))
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, fmt.Sprintf("%d\t%s", rtid, groupName)) {
			return nil
		} else if !strings.Contains(line, groupName) {
			output = append(output, line)
		}
	}

	output = append(output, fmt.Sprintf("%d\t%s", rtid, groupName))

	err = os.WriteFile(filename, []byte(strings.Join(output, "\n")), 0644)
	if err != nil {
		device.log.Errorf("modifyRtTables: can not write back file %s", filename)
		return err
	}

	return nil
}

func (device *Device) runCmd(ignoreErr bool, args ...string) {
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	err := cmd.Run()
	if nil != err {
		if !ignoreErr {
			device.log.Errorf("Error running %s: %s\n", args[0], err)
		} else {
			device.log.Verbosef("Warning: %s\n", err.Error())
		}
	}
}