hvpn-node3/link.go

package hvpnnode3

import (
	"errors"
	"fmt"
	"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
    meta *PeerMeta
    Endpoint string // Endpoint for clients to use when connecting to this link. Has no effect.
    StartedAT time.Time
    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.meta = NewPeerMeta()
    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
}

func (wg *WGLink) SetIP() error {
    ip, err := wg.Allocate()
    if err != nil {
        return err
    }
    ipnet := net.IPNet{
        IP: ip,
        Mask: wg.Network().Mask,
    }
    netlinkIP, err := netlink.ParseAddr(ipnet.String())
    
    err = netlink.AddrAdd(wg, netlinkIP)
    if err != nil {
        return err
    }

    return nil
}

// Adds a peer to the wireguard netlink.
func (wg *WGLink) AddPeer(publicKey string) (*wgtypes.Peer, error) {
    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
    }

    wg.meta.AddPeer(pubKey)
    
    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
    }

    wg.meta.DeletePeer(publickey)

    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 (wg *WGLink) GetPublicKey() (*wgtypes.Key, error) {
    dev, err := wg.Device(wg.Name)
    if err != nil {
        return nil, err
    }
    return &dev.PublicKey, nil
}

// A peer config has it's own private key and the server node (peer section) endpoint and public key
type NewPeer struct {
    Peer wgtypes.Peer // The server as a peer 
    PrivateKey wgtypes.Key
    Address net.IPNet
    DNS []net.IP
    MTU uint

    Endpoint string // Domain name and port
}

// Create a new peer. Used for server-side clinet config generation
// New peer is added to the wireguard device
func (wg *WGLink) CreateNewPeer() (*NewPeer, error) {
    dev, err := wg.Device(wg.Name)
    if err != nil {
        return nil, err
    }

    privKey, err := wgtypes.GeneratePrivateKey()
    if err != nil {
        return nil, err
    }

    serverPubkey, err := wg.GetPublicKey()
    if err != nil {
        return nil, err
    }
    
    wgPeer, err := wg.AddPeer(privKey.PublicKey().String())
    if err != nil {
        return nil, err
    }

    peer := NewPeer{
        Address: wgPeer.AllowedIPs[0],
        DNS: WG_CLIENT_DNS,
        PrivateKey: privKey,
        MTU: WG_CLIENT_MTU,
        Peer: wgtypes.Peer{
            PublicKey: *serverPubkey,
        },
        Endpoint: fmt.Sprintf("%s:%d", wg.Endpoint, dev.ListenPort),
    }

    return &peer, 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
}