130 lines
3.8 KiB
Go
130 lines
3.8 KiB
Go
package wireguard
|
||
|
||
import (
|
||
"os"
|
||
"strings"
|
||
"testing"
|
||
|
||
"vpn-platform/backend/internal/model"
|
||
"vpn-platform/backend/internal/store"
|
||
)
|
||
|
||
func setupDeviceTestTunnel(t *testing.T) *model.Tunnel {
|
||
t.Helper()
|
||
base := t.TempDir()
|
||
t.Setenv("VPN_DATA_DIR", base)
|
||
if err := store.Init(base); err != nil {
|
||
t.Fatalf("store.Init: %v", err)
|
||
}
|
||
tun := &model.Tunnel{
|
||
Name: "客户端测试专线",
|
||
Protocol: "wireguard",
|
||
State: "stopped",
|
||
Config: `{"address":"10.9.0.1/24","listen_port":53999}`,
|
||
}
|
||
if err := store.DB.Create(tun).Error; err != nil {
|
||
t.Fatalf("create tunnel: %v", err)
|
||
}
|
||
if err := New().Create(tun); err != nil {
|
||
t.Fatalf("Create: %v", err)
|
||
}
|
||
return tun
|
||
}
|
||
|
||
func TestRegisterAndRemoveDevicePeer(t *testing.T) {
|
||
tun := setupDeviceTestTunnel(t)
|
||
svc := New()
|
||
|
||
// 初始 Create 会把无 peers 的旧配置迁移出 client-1 托管 Peer（默认 10.8.0.2/32）
|
||
used, err := svc.UsedClientAddresses(tun)
|
||
if err != nil {
|
||
t.Fatalf("UsedClientAddresses: %v", err)
|
||
}
|
||
if !used["10.9.0.1"] || !used["10.8.0.2"] {
|
||
t.Fatalf("应包含服务端与初始 Peer 地址: %v", used)
|
||
}
|
||
|
||
priv, pub, err := genKey()
|
||
if err != nil || priv == "" {
|
||
t.Fatalf("genKey: %v", err)
|
||
}
|
||
psk, err := svc.RegisterDevicePeer(tun, "dev-1", "测试手机", pub, "10.9.0.7/32")
|
||
if err != nil {
|
||
t.Fatalf("RegisterDevicePeer: %v", err)
|
||
}
|
||
if !isWireGuardKey(psk) {
|
||
t.Fatal("应为设备 Peer 生成合法 PresharedKey")
|
||
}
|
||
|
||
params, err := svc.DeviceConnectionParams(tun, "dev-1")
|
||
if err != nil {
|
||
t.Fatalf("DeviceConnectionParams: %v", err)
|
||
}
|
||
if params["client_address"] != "10.9.0.7/32" || params["preshared_key"] != psk {
|
||
t.Fatalf("连接参数不含分配地址或 PSK 不匹配: %v", params)
|
||
}
|
||
if params["server_public_key"] == "" || params["endpoint"] == "" {
|
||
t.Fatalf("连接参数缺少服务端公钥或端点: %v", params)
|
||
}
|
||
|
||
// 重复公钥或重复地址必须被拒绝
|
||
if _, err := svc.RegisterDevicePeer(tun, "dev-2", "重复密钥", pub, "10.9.0.8/32"); err == nil {
|
||
t.Fatal("重复公钥必须被拒绝")
|
||
}
|
||
_, pub2, _ := genKey()
|
||
if _, err := svc.RegisterDevicePeer(tun, "dev-3", "重复地址", pub2, "10.9.0.7/32"); err == nil {
|
||
t.Fatal("重复地址必须被拒绝")
|
||
}
|
||
|
||
// 服务端 conf 必须包含新 Peer
|
||
conf := readFile(t, InstanceConfPath(tun))
|
||
if !strings.Contains(conf, pub) {
|
||
t.Fatal("服务端配置未包含设备公钥")
|
||
}
|
||
|
||
if err := svc.RemoveDevicePeer(tun, "dev-1"); err != nil {
|
||
t.Fatalf("RemoveDevicePeer: %v", err)
|
||
}
|
||
if _, err := svc.DeviceConnectionParams(tun, "dev-1"); err == nil {
|
||
t.Fatal("撤销后不应再能取到连接参数")
|
||
}
|
||
conf = readFile(t, InstanceConfPath(tun))
|
||
if strings.Contains(conf, pub) {
|
||
t.Fatal("撤销后服务端配置仍包含设备公钥")
|
||
}
|
||
}
|
||
|
||
func TestRemoveDevicePeerKeepsLastPeer(t *testing.T) {
|
||
base := t.TempDir()
|
||
t.Setenv("VPN_DATA_DIR", base)
|
||
if err := store.Init(base); err != nil {
|
||
t.Fatalf("store.Init: %v", err)
|
||
}
|
||
// 没有任何既有 Peer 的隧道：登记一个设备后它就是唯一 Peer
|
||
tun := &model.Tunnel{
|
||
Name: "空专线", Protocol: "wireguard", State: "stopped",
|
||
Config: `{"address":"10.10.0.1/24","listen_port":53998,"peers":[]}`,
|
||
}
|
||
if err := store.DB.Create(tun).Error; err != nil {
|
||
t.Fatalf("create tunnel: %v", err)
|
||
}
|
||
// peers 为空数组时 Create 会拒绝；直接手工初始化密钥再登记设备
|
||
tun.Config = `{"address":"10.10.0.1/24","listen_port":53998,"peers":[{"id":"seed","name":"种子","public_key":"","client_address":"10.10.0.2/32"}]}`
|
||
if err := New().Create(tun); err != nil {
|
||
t.Fatalf("Create: %v", err)
|
||
}
|
||
svc := New()
|
||
if err := svc.RemoveDevicePeer(tun, "seed"); err == nil {
|
||
t.Fatal("不允许移除最后一个 Peer")
|
||
}
|
||
}
|
||
|
||
func readFile(t *testing.T, path string) string {
|
||
t.Helper()
|
||
raw, err := os.ReadFile(path)
|
||
if err != nil {
|
||
t.Fatalf("read %s: %v", path, err)
|
||
}
|
||
return string(raw)
|
||
}
|