keksvpn/pkg/wgfunc/server.go

package wgfunc

import (
	"errors"
	"fmt"
	"github.com/vishvananda/netlink"
	"golang.zx2c4.com/wireguard/wgctrl"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Server struct {
	name string
	iPRanges []string

	wgDev *netlink.GenericLink
	wgConfig wgtypes.Config
	wgClient *wgctrl.Client
	wgUpdate sync.Mutex

	port int //Wireguard Server Port
	privateKey string //Wireguard Server Private Key
}

type ServerConfig struct{
	Name string
	IPRanges []string
	Port int //Wireguard Server Port
	PrivateKey string //Wireguard Server Private Key
}

func NewServer(config ServerConfig) (Server, error){
	server := Server{
		name: config.Name,
		iPRanges: config.IPRanges,
		port: config.Port,
		privateKey: config.PrivateKey,
	}

	return server, nil
}

type PeerConfig struct {
	PublicKey string
	PreSharedKey string
	AllowedIPs []string
	Endpoint string
}

func (s *Server) Down() (error) {
	return s.removeNetworkDevice()
}

func (s *Server) Up() (error)  {
	err := s.createNetworkDevice()
	if err != nil {
		return err
	}

	err = s.addIPToNetwork()
	if err != nil{
		err2 := s.removeNetworkDevice()
		if err2 != nil {
			return fmt.Errorf("Error deleting Network Device %w because of IP range error %s.", err2, err)
		}
		return err
	}

	err = s.wgSetup()
	if err != nil{
		err2 := s.removeNetworkDevice()
		if err2 != nil {
			return fmt.Errorf("Error deleting Network Device %w because of Wiregaurd error %s.", err2, err)
		}
		return err
	}

	err = s.upNetworkDevice()
	if err != nil{ //todo: test
		err2 := s.removeNetworkDevice()
		if err2 != nil {
			return fmt.Errorf("Error deleting Network Device %w because of Wiregaurd error %s.", err2, err)
		}
		return err
	}

	return nil
}

func (s *Server) AddPeer(config PeerConfig) (error) {
	s.wgUpdate.Lock()
	defer s.wgUpdate.Unlock()
	if s.wgConfig.PrivateKey == nil {
		return errors.New("server not up")
	}

	client, err := wgctrl.New()
	if err != nil {
		return err
	}

	pubKey, err := wgtypes.ParseKey(config.PublicKey)
	if err != nil {
		return err
	}

	var ipList []net.IPNet

	for _, ip := range config.AllowedIPs {
		_, net, err := net.ParseCIDR(ip)
		if err != nil {
			return err
		}
		ipList = append(ipList, *net)
	}

	d := 30 * time.Second

	var endpoint *net.UDPAddr
	if config.Endpoint != "" {
		//todo: allow hostnames
		details := strings.Split(config.Endpoint, ":")
		ip := net.ParseIP(details[0])
		if ip == nil{
			return errors.New("Invalide endpoint ip")
		}

		port, err := strconv.Atoi(details[1])
		if err != nil {
			return fmt.Errorf("Cant convert port to int: %w", err)
		}

		udpaddr := net.UDPAddr{
			IP:   ip,
			Port: port,
		}
		endpoint = &udpaddr
	}

	var preKey *wgtypes.Key
	if config.PreSharedKey != "" {
		presharedKey, err := wgtypes.ParseKey(config.PreSharedKey )
		if err != nil {
			return err
		}
		preKey = &presharedKey
	}



	peerConfig := wgtypes.PeerConfig{
		PublicKey:                   pubKey,
		Remove:                      false,
		UpdateOnly:                  false,
		PresharedKey:                preKey,
		Endpoint:                    endpoint,
		PersistentKeepaliveInterval: &d,
		ReplaceAllowedIPs:           true,
		AllowedIPs:                  ipList,
	}

	s.wgConfig.Peers = append(s.wgConfig.Peers, peerConfig)

	err = client.ConfigureDevice(s.name, s.wgConfig)
	if err != nil {
		return err
	}

	return nil
}

func (s *Server) RemovePeer(publicKey string) (error) {
	s.wgUpdate.Lock()
	defer s.wgUpdate.Unlock()

	if s.wgConfig.PrivateKey == nil {
		return errors.New("server not up")
	}

	client, err := wgctrl.New()
	if err != nil {
		return err
	}

	var newlist []wgtypes.PeerConfig

	for _ , peer := range s.wgConfig.Peers {
		if peer.PublicKey.String() != publicKey {
			newlist = append(newlist, peer)
		}
	}

	s.wgConfig.Peers = newlist
	err = client.ConfigureDevice(s.name, s.wgConfig)
	if err != nil {
		return err
	}
	return nil
}

func (s *Server) createNetworkDevice() (error) {
	la := netlink.NewLinkAttrs()
	la.Name = s.name

	wgDev := &netlink.GenericLink{
		LinkAttrs: la,
		LinkType:  "wireguard",
	}

	err := netlink.LinkAdd(wgDev)
	if err != nil {
		return err
	}

	s.wgDev = wgDev

	return nil
}

func (s *Server) upNetworkDevice() (error) {
	err := netlink.LinkSetUp(s.wgDev) //todo: test
	if err != nil {
		return err
	}
	return nil
}

func (s *Server) removeNetworkDevice() (error) {
	la := netlink.NewLinkAttrs()
	la.Name = s.name

	wgDev := &netlink.GenericLink{
		LinkAttrs: la,
		LinkType:  "wireguard",
	}

	err := netlink.LinkDel(wgDev)
	return err
}

func (s *Server) addIPToNetwork() (error) {
	var ips []*netlink.Addr
	for _, iprange := range s.iPRanges {
		ip, err := netlink.ParseAddr(iprange)
		if err != nil {
			return err
		}
		ips = append(ips, ip)
	}

	for _, ip := range ips {
		err := netlink.AddrAdd(s.wgDev, ip)
		if err != nil { //ToDo: Check how to test this case
			return err
		}
	}
	return nil
}

func (s *Server) wgSetup() (error) {
	if s.wgConfig.PrivateKey != nil {
		return errors.New("Wireguard already setedup")
	}

	client, err := wgctrl.New()
	if err != nil {
		return err
	}
	s.wgClient = client

	privkey, err := wgtypes.ParseKey(s.privateKey)
	if err != nil {
		return err
	}


	s.wgConfig = wgtypes.Config{
		PrivateKey:   &privkey,
		ListenPort:   &s.port,
		FirewallMark: nil,
		ReplacePeers: true,
		Peers:        nil,
	}

	s.wgClient.ConfigureDevice(s.name, s.wgConfig)

	return nil
}