vppn/peer/init.go

package peer

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/netip"
	"os"
	"path/filepath"

	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"vppn/m"
)

// LocalState is the persisted identity for this peer, written on first run and
// loaded on every subsequent run.
type LocalState struct {
	PrivKey  wgtypes.Key
	VPNIP    netip.Addr
	VPNNet   netip.Prefix
	WGPort   uint16
	IsRelay  bool
	IsPublic bool
}

// localStateJSON is the on-disk representation.
type localStateJSON struct {
	PrivKey  string       `json:"priv_key"` // standard base64
	VPNIP    netip.Addr   `json:"vpn_ip"`
	VPNNet   netip.Prefix `json:"vpn_net"`
	WGPort   uint16       `json:"wg_port"`
	IsRelay  bool         `json:"is_relay"`
	IsPublic bool         `json:"is_public"`
}

// LoadOrInit loads LocalState from path, or registers with the hub and creates
// the file if it doesn't exist.
func LoadOrInit(statePath, hubURL, apiKey string) (LocalState, error) {
	if data, err := os.ReadFile(statePath); err == nil {
		return parseLocalState(data)
	}

	privKey, err := wgtypes.GeneratePrivateKey()
	if err != nil {
		return LocalState{}, fmt.Errorf("generate key: %w", err)
	}

	state, err := initFromHub(hubURL, apiKey, privKey)
	if err != nil {
		return LocalState{}, err
	}

	if err := saveLocalState(statePath, state); err != nil {
		return LocalState{}, fmt.Errorf("save state: %w", err)
	}
	return state, nil
}

func initFromHub(hubURL, apiKey string, privKey wgtypes.Key) (LocalState, error) {
	pubKey := privKey.PublicKey()
	body, _ := json.Marshal(m.PeerInitArgs{WGPubKey: pubKey[:]})

	req, err := http.NewRequest(http.MethodPost, hubURL+"/peer/init/", bytes.NewReader(body))
	if err != nil {
		return LocalState{}, err
	}
	req.SetBasicAuth("", apiKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return LocalState{}, fmt.Errorf("hub init: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return LocalState{}, fmt.Errorf("hub init: HTTP %d", resp.StatusCode)
	}

	var r m.PeerInitResp
	if err := json.NewDecoder(resp.Body).Decode(&r); err != nil {
		return LocalState{}, fmt.Errorf("hub init decode: %w", err)
	}

	if len(r.Network) != 4 {
		return LocalState{}, fmt.Errorf("hub init: invalid network %v", r.Network)
	}

	netAddr := netip.AddrFrom4([4]byte(r.Network))
	octets := netAddr.As4()
	octets[3] = r.PeerIP
	vpnIP := netip.AddrFrom4(octets)
	vpnNet := netip.PrefixFrom(netAddr, 24)

	var isRelay, isPublic bool
	var wgPort uint16
	if self := r.NetworkState.Peers[r.PeerIP]; self != nil {
		isRelay = self.Relay
		isPublic = len(self.Addr4) > 0 || len(self.Addr6) > 0
		wgPort = self.Port
	}

	return LocalState{
		PrivKey:  privKey,
		VPNIP:    vpnIP,
		VPNNet:   vpnNet,
		WGPort:   wgPort,
		IsRelay:  isRelay,
		IsPublic: isPublic,
	}, nil
}

func parseLocalState(data []byte) (LocalState, error) {
	var j localStateJSON
	if err := json.Unmarshal(data, &j); err != nil {
		return LocalState{}, fmt.Errorf("parse state: %w", err)
	}
	keyBytes, err := base64.StdEncoding.DecodeString(j.PrivKey)
	if err != nil {
		return LocalState{}, fmt.Errorf("decode key: %w", err)
	}
	key, err := wgtypes.NewKey(keyBytes)
	if err != nil {
		return LocalState{}, fmt.Errorf("invalid key: %w", err)
	}
	return LocalState{
		PrivKey:  key,
		VPNIP:    j.VPNIP,
		VPNNet:   j.VPNNet,
		WGPort:   j.WGPort,
		IsRelay:  j.IsRelay,
		IsPublic: j.IsPublic,
	}, nil
}

func saveLocalState(path string, s LocalState) error {
	j := localStateJSON{
		PrivKey:  base64.StdEncoding.EncodeToString(s.PrivKey[:]),
		VPNIP:    s.VPNIP,
		VPNNet:   s.VPNNet,
		WGPort:   s.WGPort,
		IsRelay:  s.IsRelay,
		IsPublic: s.IsPublic,
	}
	data, err := json.MarshalIndent(j, "", "  ")
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return err
	}
	return os.WriteFile(path, data, 0600)
}