hvpn-node3/link.go

package hvpnnode3

import (
	"errors"
	"fmt"
	"log/slog"
	"net"
	"sync"
	"time"

	"gitea.hbanafa.com/HeshamTB/hvpn-node3/proto"
	"github.com/vishvananda/netlink"
	"golang.zx2c4.com/wireguard/wgctrl"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
)

var SINGLE_IP_MASK net.IPMask = net.IPv4Mask(255, 255, 255, 255)
var PEER_KEEP_ALIVE_INTERVAL time.Duration = time.Second * 25

type WGLink struct {
    *netlink.LinkAttrs
    *wgctrl.Client
    IPPool
    Endpoint string // Endpoint for clients to use when connecting to this link. Has no effect.
    lock *sync.Mutex
}

// Retruns an existing or create a WGLink
func InitWGLink(ifName string, privateKey *wgtypes.Key, port int, endpoint string) (*WGLink, error){
    attrs := netlink.NewLinkAttrs()
    attrs.Name = ifName
    wg := WGLink{LinkAttrs: &attrs, Endpoint: endpoint}
    link, err := netlink.LinkByName(ifName)
    if err != nil {
        switch err.(type) {
        case netlink.LinkNotFoundError:
            if err := netlink.LinkAdd(&wg); err != nil {
                return nil, err
            }
        default:
            return nil, err
        }
    } else {
        wg.LinkAttrs = link.Attrs()
    }
    
    err = wg.initClient()
    if err != nil {
        return nil, err
    }

    err = wg.ConfigureDevice(
        ifName,
        wgtypes.Config{
            PrivateKey: privateKey,
            ListenPort: &port,
        },
    )
    if err != nil {
        return nil, err
    }

    wg.lock = &sync.Mutex{}

    return &wg, netlink.LinkSetUp(&wg)
}

func (WGLink) Type() string {
    return "wireguard"
}

func (wg *WGLink) Attrs() *netlink.LinkAttrs {
    return wg.LinkAttrs
}

func (wg *WGLink) Close() error {
    return netlink.LinkDel(wg)
}

func (wg *WGLink) initClient() error {
    client, err := wgctrl.New()
    if client == nil {
        return errors.New("Could not initialize new Wireguard Client")
    }
    wg.Client = client
    return err
}

// Adds a peer to the wireguard netlink.
func (wg *WGLink) AddPeer(publicKey string) (*wgtypes.Peer, error) {
    slog.Debug(fmt.Sprintf("Trying to add peer %s", publicKey))
    pubKey, err := wgtypes.ParseKey(publicKey)
    if err != nil {
        return nil, err
    }

    exists, err := wg.exists(pubKey)
    if err != nil {
        return nil, err
    }
    if exists {
        return nil, proto.PeerExistsErr{PublicKey: publicKey}
    }
    
    peerIp, err := wg.Allocate()
    if err != nil {
        return nil, err
    }

    peerConfig := wgtypes.PeerConfig{
        PublicKey: pubKey,
        AllowedIPs: []net.IPNet{
            {
                IP: peerIp,
                Mask: SINGLE_IP_MASK,
            },
        },
        PersistentKeepaliveInterval: &PEER_KEEP_ALIVE_INTERVAL,
        ReplaceAllowedIPs: true,
    }

    wgConf := wgtypes.Config{
        ReplacePeers: false,
        Peers: []wgtypes.PeerConfig{peerConfig},
    }

    err = wg.ApplyConfig(wgConf)
    if err != nil {
        return nil, err
    }
    
    return wg.getPeer(pubKey)
}

func (wg *WGLink) DeletePeer(publickey string) error {
    pkey, err := wgtypes.ParseKey(publickey)
    if err != nil {
        return err
    }
    return wg.deletePeer(pkey)
}

func (wg *WGLink) deletePeer(publickey wgtypes.Key) error {
    rmCfg := createARemovePeerCfg(publickey)
    ip, err := wg.getPeerIP(publickey)
    if err != nil {
        return err
    }

    err = wg.ApplyConfig(rmCfg)
    if err != nil {
        return err
    }

    err = wg.Free(ip)
    if err != nil {
        return err
    }

    return nil
}

func (wg *WGLink) getPeerIP(publickey wgtypes.Key) (net.IP, error) {
    peer, err := wg.getPeer(publickey)
    if err != nil { return nil, err }
    return peer.AllowedIPs[0].IP, nil
}

func (wg *WGLink) Exists(publicKey string) (bool, error) {
    pubkey, err := wgtypes.ParseKey(publicKey)
    if err != nil {
        return false, err
    }
    return wg.exists(pubkey)
}

func (wg *WGLink) exists(pubkey wgtypes.Key) (bool, error) {

    dev, err := wg.Client.Device(wg.Name)
    if err != nil {
        return false, err
    }

    for _, peer := range dev.Peers {
        if peer.PublicKey == pubkey {
            return true, nil
        }
    }
    return false, nil

}

// Calls wgctrl.ConfigureDevice wrapped with a Mutex to ensure sync
func (wg *WGLink) ApplyConfig(cfg wgtypes.Config) error {
    wg.lock.Lock()
    defer wg.lock.Unlock()
    return wg.ConfigureDevice(wg.Name, cfg)
}

func (wg *WGLink) GetPeer(publickey string) (*wgtypes.Peer, error) {
    k, err := wgtypes.ParseKey(publickey)
    if err != nil {
        return nil, err
    }
    return wg.getPeer(k)
}

func (wg *WGLink) getPeer(pubkey wgtypes.Key) (*wgtypes.Peer, error) {
    
    peers, err := wg.GetAllPeers()
    if err != nil {
        return nil, err
    }

    for _, peer := range peers {
        if peer.PublicKey == pubkey {
            return &peer, nil
        }
    }
    return nil, proto.PeerDoesNotExist
}

func (wg *WGLink) GetAllPeers() ([]wgtypes.Peer, error) {
    dev, err := wg.Device(wg.Name)
    if err != nil {
        return nil, err
    }
    return dev.Peers, nil
}

func createARemovePeerCfg(publickey wgtypes.Key) wgtypes.Config {
    rmPeerCfg := wgtypes.PeerConfig{
        Remove: true,
        PublicKey: publickey,
    }
    return wgtypes.Config{
        Peers: []wgtypes.PeerConfig{rmPeerCfg},
        ReplacePeers: false,
    }
}

func ValidKey(key string) bool {
    _, err := wgtypes.ParseKey(key)
    if err != nil { return false }
    return true
}