wgcl/internal/wg/wg.go

package wg

/*
 #cgo CFLAGS: -I./
 #include <wireguard.h>
 #include <arpa/inet.h>
 #include <stdlib.h>
 enum wg_device_flags dev_cfg = WGDEVICE_HAS_PRIVATE_KEY|WGDEVICE_HAS_LISTEN_PORT;
*/
import "C"
import (
	"bytes"
	"encoding/binary"
	"errors"
	"math/bits"
	"net/netip"
	"sync"
	"unsafe"

	"git.jasinco.work/wgcl/internal/logger"
	protovpn "git.jasinco.work/wgcl/proto/out"
	"github.com/vishvananda/netlink"
	"google.golang.org/protobuf/proto"
)

// TODO: inspect all api, complete error handling

type WGPeer struct {
	name    string
	privkey string //base64 string
	peer    C.wg_peer
	next    *WGPeer
}

type WG struct {
	dev        *C.wg_device
	lock       sync.RWMutex
	keep_alive bool
	peers      *WGPeer
	peers_len  int
	IP4Pool    *IPPool
}

var WGfailedadddev = errors.New("failed add wgdev")
var WGfailedsetdev = errors.New("failed set wgdev")
var WGfailedparseb64key = errors.New("failed parse base64 key")

func GetB64Key(key *C.wg_key) string {
	var psk_b64 C.wg_key_b64_string

	psk_ptr := (*C.uint8_t)(unsafe.Pointer(&(*key)[0]))
	psk_b64_ptr := (*C.char)(unsafe.Pointer(&psk_b64[0]))
	C.wg_key_to_base64(psk_b64_ptr, psk_ptr)
	return C.GoString(psk_b64_ptr)
}

func NewKey() string {
	var privkey C.wg_key

	privkey_as_ptr := (*C.uint8_t)(unsafe.Pointer(&privkey[0]))
	C.wg_generate_private_key(privkey_as_ptr)
	return GetB64Key(&privkey)
}

func NewPresharedKey() string {
	var psk C.wg_key
	psk_ptr := (*C.uint8_t)(unsafe.Pointer(&psk[0]))
	C.wg_generate_preshared_key(psk_ptr)
	return GetB64Key(&psk)
}

// device name and receive base64 encoded keys
func WG_init(name string, privkey string, listen uint16, ip4_pool *IPPool) (*WG, error) {
	var dev WG
	dev.dev = (*C.wg_device)(C.malloc(C.sizeof_wg_device))

	// hence it only call once at startup, I'd rather not to free it explicitly
	C.wg_key_from_base64((*C.uint8_t)(unsafe.Pointer(&dev.dev.private_key[0])), C.CString(privkey))

	// set device name
	ifname := C.CString(name)
	C.strncpy(&dev.dev.name[0], ifname, 16)

	dev.dev.flags = C.dev_cfg
	dev.dev.first_peer = (*C.wg_peer)(C.NULL)
	dev.dev.last_peer = (*C.wg_peer)(C.NULL)
	dev.dev.listen_port = (C.uint16_t)(listen)
	if C.wg_add_device((*C.char)(unsafe.Pointer(&dev.dev.name[0]))) < 0 {
		return nil, WGfailedadddev
	}
	if C.wg_set_device(dev.dev) < 0 {
		return nil, WGfailedsetdev
	}
	if C.wg_get_device(&dev.dev, (*C.char)(unsafe.Pointer(&dev.dev.name[0]))) < 0 {
		return nil, WGfailedsetdev
	}
	dev.peers = nil
	dev.IP4Pool = ip4_pool

	// setup interface
	var err error
	wg_if, err := netlink.LinkByIndex(int(dev.dev.ifindex))
	if err != nil {
		return nil, err
	}
	nl_addr, err := netlink.ParseAddr(ip4_pool.scope.String())
	if err != nil {
		return nil, err
	}
	err = netlink.AddrAdd(wg_if, nl_addr)
	if err != nil {
		return nil, err
	}
	err = netlink.LinkSetUp(wg_if)
	if err != nil {
		return nil, err
	}

	return &dev, nil
}

func (wg *WG) SyncStatus() {
	wg.lock.Lock()
	defer wg.lock.Unlock()
	if C.wg_get_device(&wg.dev, (*C.char)(unsafe.Pointer(&wg.dev.name[0]))) < 0 {
		logger.Logger.Info("Failed to sync wg status")
	}
}

func (wg *WG) Deinit() {
	C.wg_del_device((*C.char)(unsafe.Pointer(&wg.dev.name[0])))
}

func (wg *WG) WriteTxStart() {
	wg.lock.Lock()
}
func (wg *WG) WriteTxEnd() {
	wg.lock.Unlock()
}

func (wg *WG) AddPeer(peer *WGPeer) {
	peer.next = nil
	wg.peers_len += 1
	if wg.peers == nil {
		wg.peers = peer
		return
	}
	ptr := wg.peers
	for ptr.next != nil {
		ptr = ptr.next
	}
	ptr.next = peer
}

func (wg *WG) RemovePeer(peer *WGPeer) {
	if wg.peers == peer {
		wg.peers = peer.next
	}
	for ptr := wg.peers; ptr != nil; ptr = ptr.next {
		if ptr.next == peer {
			ptr.next = peer.next
			break
		}
	}
	wg.peers_len -= 1
}

func (wg *WG) Apply() error {

	if wg.peers != nil {
		first_ptr := (*C.wg_peer)(unsafe.Pointer(&wg.peers.peer))
		wg.dev.first_peer = first_ptr
	} else {
		wg.dev.first_peer = nil
		wg.dev.last_peer = nil
	}

	for ptr := wg.peers; ptr != nil; ptr = (*WGPeer)(ptr.next) {
		if ptr.next == nil {
			ptr.peer.next_peer = nil
			last_ptr := (*C.wg_peer)(unsafe.Pointer(&ptr.peer))
			wg.dev.last_peer = last_ptr
			break
		}
		ptr.peer.next_peer = &(*WGPeer)(ptr.next).peer
	}
	wgdev := (*C.wg_device)(unsafe.Pointer(wg.dev))
	if C.wg_set_device(wgdev) < 0 {
		return WGfailedsetdev
	}
	return nil
}
func (wg *WG) SearchPeerIP4(addr netip.Addr) *WGPeer {
	wg.lock.RLock()
	defer wg.lock.RUnlock()
	ip4_slice := addr.As4()
	for ptr := wg.peers; ptr != nil; ptr = ptr.next {
		if bytes.Equal(ip4_slice[:], ptr.peer.first_allowedip.anon0[:4]) {
			return ptr
		}
	}
	return nil
}
func (wg *WG) SearchPeerPubkey(pubkey string) (*WGPeer, error) {
	wg.lock.RLock()
	defer wg.lock.RUnlock()
	pubkey_cstr := C.CString(pubkey)
	defer C.free(unsafe.Pointer(pubkey_cstr))
	var pubkey_bin C.wg_key
	if C.wg_key_from_base64((*C.uint8_t)(&pubkey_bin[0]), pubkey_cstr) < 0 {
		return nil, WGfailedparseb64key
	}
	for ptr := wg.peers; ptr != nil; ptr = ptr.next {
		if C.wg_key_equal(&pubkey_bin, &ptr.peer.public_key) == 0 {
			return ptr, nil
		}
	}
	return nil, nil
}
func (w *WG) GetPeersJson() *protovpn.PeerResp {
	w.lock.RLock()
	defer w.lock.RUnlock()
	peers := make([]*protovpn.PeerStatus, w.peers_len)
	peers_ptr := 0
	for ptr := w.peers; ptr != nil; ptr = ptr.next {
		peers[peers_ptr] = ptr.ToPeerStatus()
		peers_ptr += 1
	}
	resp := protovpn.PeerResp_builder{Peers: peers}
	return resp.Build()
}

func WGPeer_init(name string, privkey string, remote netip.Addr) *WGPeer {
	var peer WGPeer
	peer.name = name
	peer.privkey = privkey

	var privkey_bin C.wg_key
	priv_cstr := C.CString(privkey)
	defer C.free(unsafe.Pointer(priv_cstr))
	C.wg_key_from_base64((*C.uint8_t)(unsafe.Pointer(&privkey_bin[0])), priv_cstr)
	C.wg_generate_public_key((*C.uint8_t)(unsafe.Pointer(&peer.peer.public_key[0])), (*C.uint8_t)(unsafe.Pointer(&privkey_bin[0])))
	peer.peer.flags = C.WGPEER_HAS_PUBLIC_KEY | C.WGPEER_REPLACE_ALLOWEDIPS

	remote_ip_arr := remote.As4()
	remote_ip := binary.LittleEndian.Uint32(remote_ip_arr[:])
	remote_endpoint := C.wg_allowed_ip_new_ip4((C.in_addr_t)(remote_ip), 32)

	peer.peer.first_allowedip = remote_endpoint
	peer.peer.last_allowedip = remote_endpoint

	return &peer
}

func WGPeer_free(peer *WGPeer) {
	C.free(unsafe.Pointer(peer.peer.first_allowedip))
}
func (p *WGPeer) Disable() {
	p.peer.flags |= C.WGPEER_REMOVE_ME
}
func (p *WGPeer) Enable() {
	p.peer.flags &= bits.Reverse32(C.WGPEER_REMOVE_ME)
}
func (p *WGPeer) ToPeerStatus() *protovpn.PeerStatus {
	ip4, _ := netip.AddrFromSlice(p.peer.first_allowedip.anon0[:4])
	return protovpn.PeerStatus_builder{
		Enabled: proto.Bool((p.peer.flags & C.WGPEER_REMOVE_ME) == 0),
		Pubkey:  proto.String(GetB64Key(&p.peer.public_key)),
		Nick:    proto.String(p.name),
		Ipv4:    proto.String(ip4.String()),
		Beacon:  proto.Int64(int64(p.peer.last_handshake_time.tv_sec)),
	}.Build()
}