更新管理登录功能
This commit is contained in:
@@ -21,3 +21,5 @@ frontend/node_modules/
|
||||
.env
|
||||
.env.local
|
||||
.env.*.local
|
||||
|
||||
lmvpnweb.sock
|
||||
@@ -0,0 +1,17 @@
|
||||
# lmvpn_server
|
||||
|
||||
## 目录说明
|
||||
|
||||
- `data/` - 程序运行生成的数据（配置文件、数据库文件）
|
||||
- `dist/` - 前端构建产物，由 Go 服务端托管提供静态资源
|
||||
- `frontend/` - 前端工程（Vue 3 + TypeScript + Vite）
|
||||
- `internal/` - Go 后端源码
|
||||
- `config/` - 配置加载
|
||||
- `db/` - 数据库操作
|
||||
- `handler/` - HTTP 请求处理
|
||||
- `middleware/` - 中间件（认证等）
|
||||
- `model/` - 数据模型
|
||||
- `vpn/` - VPN 相关逻辑（认证、隧道等）
|
||||
- `pytest/` - Python 测试脚本
|
||||
- `main.go` - Go 服务端入口文件
|
||||
- `go.mod` / `go.sum` - Go 模块依赖管理
|
||||
+19
-4
@@ -1,6 +1,14 @@
|
||||
<script setup lang="ts">
|
||||
import { RouterLink, RouterView } from 'vue-router'
|
||||
import HelloWorld from './components/HelloWorld.vue'
|
||||
import { RouterLink, RouterView, useRouter } from 'vue-router'
|
||||
import { useAuthStore } from '@/stores/auth'
|
||||
|
||||
const router = useRouter()
|
||||
const authStore = useAuthStore()
|
||||
|
||||
function handleLogout() {
|
||||
authStore.logout()
|
||||
router.push('/')
|
||||
}
|
||||
</script>
|
||||
|
||||
<template>
|
||||
@@ -8,11 +16,16 @@ import HelloWorld from './components/HelloWorld.vue'
|
||||
<img alt="Vue logo" class="logo" src="@/assets/logo.svg" width="125" height="125" />
|
||||
|
||||
<div class="wrapper">
|
||||
<HelloWorld msg="You did it!" />
|
||||
|
||||
<nav>
|
||||
<RouterLink to="/">Home</RouterLink>
|
||||
<RouterLink to="/about">About</RouterLink>
|
||||
<template v-if="authStore.isLoggedIn">
|
||||
<RouterLink to="/admin">管理后台</RouterLink>
|
||||
<a href="#" @click.prevent="handleLogout">退出</a>
|
||||
</template>
|
||||
<template v-else>
|
||||
<RouterLink to="/login">登录</RouterLink>
|
||||
</template>
|
||||
</nav>
|
||||
</div>
|
||||
</header>
|
||||
@@ -50,6 +63,8 @@ nav a {
|
||||
display: inline-block;
|
||||
padding: 0 1rem;
|
||||
border-left: 1px solid var(--color-border);
|
||||
color: var(--color-text);
|
||||
text-decoration: none;
|
||||
}
|
||||
|
||||
nav a:first-of-type {
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import { createRouter, createWebHistory } from 'vue-router'
|
||||
import HomeView from '../views/HomeView.vue'
|
||||
import { useAuthStore } from '@/stores/auth'
|
||||
|
||||
const router = createRouter({
|
||||
history: createWebHistory(import.meta.env.BASE_URL),
|
||||
@@ -12,12 +13,32 @@ const router = createRouter({
|
||||
{
|
||||
path: '/about',
|
||||
name: 'about',
|
||||
// route level code-splitting
|
||||
// this generates a separate chunk (About.[hash].js) for this route
|
||||
// which is lazy-loaded when the route is visited.
|
||||
component: () => import('../views/AboutView.vue'),
|
||||
},
|
||||
{
|
||||
path: '/login',
|
||||
name: 'login',
|
||||
component: () => import('../views/LoginView.vue'),
|
||||
},
|
||||
{
|
||||
path: '/admin',
|
||||
name: 'admin',
|
||||
component: () => import('../views/AdminView.vue'),
|
||||
meta: { requiresAuth: true },
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
router.beforeEach((to, from, next) => {
|
||||
const authStore = useAuthStore()
|
||||
|
||||
if (to.meta.requiresAuth && !authStore.isLoggedIn) {
|
||||
next({ name: 'login' })
|
||||
} else if (to.name === 'login' && authStore.isLoggedIn) {
|
||||
next({ name: 'admin' })
|
||||
} else {
|
||||
next()
|
||||
}
|
||||
})
|
||||
|
||||
export default router
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
import { ref, computed } from 'vue'
|
||||
import { defineStore } from 'pinia'
|
||||
|
||||
interface UserInfo {
|
||||
id: number
|
||||
username: string
|
||||
role: string
|
||||
}
|
||||
|
||||
export const useAuthStore = defineStore('auth', () => {
|
||||
const token = ref<string>(localStorage.getItem('token') || '')
|
||||
const user = ref<UserInfo | null>(null)
|
||||
|
||||
const isLoggedIn = computed(() => !!token.value)
|
||||
|
||||
function setAuth(t: string, u: UserInfo) {
|
||||
token.value = t
|
||||
user.value = u
|
||||
localStorage.setItem('token', t)
|
||||
}
|
||||
|
||||
function clearAuth() {
|
||||
token.value = ''
|
||||
user.value = null
|
||||
localStorage.removeItem('token')
|
||||
}
|
||||
|
||||
async function login(username: string, password: string) {
|
||||
const res = await fetch('/api/login', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ username, password }),
|
||||
})
|
||||
if (!res.ok) {
|
||||
const data = await res.json()
|
||||
throw new Error(data.error || '登录失败')
|
||||
}
|
||||
const data = await res.json()
|
||||
setAuth(data.token, data.user)
|
||||
}
|
||||
|
||||
function logout() {
|
||||
clearAuth()
|
||||
}
|
||||
|
||||
async function fetchUser() {
|
||||
if (!token.value) return
|
||||
try {
|
||||
const res = await fetch('/api/me', {
|
||||
headers: { Authorization: `Bearer ${token.value}` },
|
||||
})
|
||||
if (!res.ok) {
|
||||
clearAuth()
|
||||
return
|
||||
}
|
||||
user.value = await res.json()
|
||||
} catch {
|
||||
clearAuth()
|
||||
}
|
||||
}
|
||||
|
||||
return { token, user, isLoggedIn, login, logout, fetchUser }
|
||||
})
|
||||
@@ -0,0 +1,67 @@
|
||||
<script setup lang="ts">
|
||||
import { onMounted } from 'vue'
|
||||
import { useRouter } from 'vue-router'
|
||||
import { useAuthStore } from '@/stores/auth'
|
||||
|
||||
const router = useRouter()
|
||||
const authStore = useAuthStore()
|
||||
|
||||
onMounted(async () => {
|
||||
await authStore.fetchUser()
|
||||
})
|
||||
|
||||
function handleLogout() {
|
||||
authStore.logout()
|
||||
router.push('/')
|
||||
}
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<div class="admin-container">
|
||||
<h2>管理后台</h2>
|
||||
<div class="user-card">
|
||||
<p><strong>用户名：</strong>{{ authStore.user?.username }}</p>
|
||||
<p><strong>角色：</strong>{{ authStore.user?.role === 'admin' ? '管理员' : '普通用户' }}</p>
|
||||
</div>
|
||||
<button class="logout-btn" @click="handleLogout">退出登录</button>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<style scoped>
|
||||
.admin-container {
|
||||
max-width: 600px;
|
||||
margin: 2rem auto;
|
||||
padding: 2rem;
|
||||
}
|
||||
|
||||
.admin-container h2 {
|
||||
margin-bottom: 1.5rem;
|
||||
color: var(--color-heading);
|
||||
}
|
||||
|
||||
.user-card {
|
||||
padding: 1rem;
|
||||
border: 1px solid var(--color-border);
|
||||
border-radius: 8px;
|
||||
background: var(--color-background-soft);
|
||||
}
|
||||
|
||||
.user-card p {
|
||||
margin: 0.5rem 0;
|
||||
}
|
||||
|
||||
.logout-btn {
|
||||
margin-top: 1.5rem;
|
||||
padding: 0.6rem 1.5rem;
|
||||
border: none;
|
||||
border-radius: 4px;
|
||||
background: #e74c3c;
|
||||
color: #fff;
|
||||
font-size: 0.95rem;
|
||||
cursor: pointer;
|
||||
}
|
||||
|
||||
.logout-btn:hover {
|
||||
opacity: 0.9;
|
||||
}
|
||||
</style>
|
||||
@@ -0,0 +1,141 @@
|
||||
<script setup lang="ts">
|
||||
import { ref } from 'vue'
|
||||
import { useRouter } from 'vue-router'
|
||||
import { useAuthStore } from '@/stores/auth'
|
||||
|
||||
const router = useRouter()
|
||||
const authStore = useAuthStore()
|
||||
|
||||
const username = ref('')
|
||||
const password = ref('')
|
||||
const error = ref('')
|
||||
const loading = ref(false)
|
||||
|
||||
async function handleLogin() {
|
||||
error.value = ''
|
||||
if (!username.value || !password.value) {
|
||||
error.value = '请输入用户名和密码'
|
||||
return
|
||||
}
|
||||
loading.value = true
|
||||
try {
|
||||
await authStore.login(username.value, password.value)
|
||||
router.push('/admin')
|
||||
} catch (e: any) {
|
||||
error.value = e.message || '登录失败'
|
||||
} finally {
|
||||
loading.value = false
|
||||
}
|
||||
}
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<div class="login-container">
|
||||
<div class="login-card">
|
||||
<h2>LmVPN 登录</h2>
|
||||
<form @submit.prevent="handleLogin">
|
||||
<div class="form-group">
|
||||
<label for="username">用户名</label>
|
||||
<input
|
||||
id="username"
|
||||
v-model="username"
|
||||
type="text"
|
||||
placeholder="请输入用户名"
|
||||
autocomplete="username"
|
||||
/>
|
||||
</div>
|
||||
<div class="form-group">
|
||||
<label for="password">密码</label>
|
||||
<input
|
||||
id="password"
|
||||
v-model="password"
|
||||
type="password"
|
||||
placeholder="请输入密码"
|
||||
autocomplete="current-password"
|
||||
/>
|
||||
</div>
|
||||
<p v-if="error" class="error-msg">{{ error }}</p>
|
||||
<button type="submit" :disabled="loading">
|
||||
{{ loading ? '登录中...' : '登录' }}
|
||||
</button>
|
||||
</form>
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<style scoped>
|
||||
.login-container {
|
||||
display: flex;
|
||||
justify-content: center;
|
||||
align-items: center;
|
||||
min-height: calc(100vh - 120px);
|
||||
}
|
||||
|
||||
.login-card {
|
||||
width: 100%;
|
||||
max-width: 380px;
|
||||
padding: 2rem;
|
||||
border-radius: 8px;
|
||||
box-shadow: 0 2px 12px rgba(0, 0, 0, 0.1);
|
||||
background: var(--color-background);
|
||||
}
|
||||
|
||||
.login-card h2 {
|
||||
text-align: center;
|
||||
margin-bottom: 1.5rem;
|
||||
color: var(--color-heading);
|
||||
}
|
||||
|
||||
.form-group {
|
||||
margin-bottom: 1rem;
|
||||
}
|
||||
|
||||
.form-group label {
|
||||
display: block;
|
||||
margin-bottom: 0.4rem;
|
||||
font-size: 0.9rem;
|
||||
color: var(--color-text);
|
||||
}
|
||||
|
||||
.form-group input {
|
||||
width: 100%;
|
||||
padding: 0.6rem 0.8rem;
|
||||
border: 1px solid var(--color-border);
|
||||
border-radius: 4px;
|
||||
font-size: 1rem;
|
||||
background: var(--color-background);
|
||||
color: var(--color-text);
|
||||
box-sizing: border-box;
|
||||
}
|
||||
|
||||
.form-group input:focus {
|
||||
outline: none;
|
||||
border-color: var(--vt-c-indigo);
|
||||
}
|
||||
|
||||
.error-msg {
|
||||
color: #e74c3c;
|
||||
font-size: 0.85rem;
|
||||
margin-bottom: 0.5rem;
|
||||
}
|
||||
|
||||
button {
|
||||
width: 100%;
|
||||
padding: 0.65rem;
|
||||
border: none;
|
||||
border-radius: 4px;
|
||||
background: var(--vt-c-indigo);
|
||||
color: #fff;
|
||||
font-size: 1rem;
|
||||
cursor: pointer;
|
||||
}
|
||||
|
||||
button:hover {
|
||||
opacity: 0.9;
|
||||
}
|
||||
|
||||
button:disabled {
|
||||
opacity: 0.6;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
</style>
|
||||
@@ -23,6 +23,7 @@ require (
|
||||
github.com/go-playground/validator/v10 v10.30.1 // indirect
|
||||
github.com/go-sql-driver/mysql v1.8.1 // indirect
|
||||
github.com/goccy/go-json v0.10.5 // indirect
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1 // indirect
|
||||
github.com/google/uuid v1.3.0 // indirect
|
||||
github.com/gorilla/websocket v1.5.3 // indirect
|
||||
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||
|
||||
@@ -37,6 +37,8 @@ github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4=
|
||||
github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
|
||||
github.com/goccy/go-yaml v1.19.2 h1:PmFC1S6h8ljIz6gMRBopkjP1TVT7xuwrButHID66PoM=
|
||||
github.com/goccy/go-yaml v1.19.2/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7LkFRi1kA=
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY=
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"lmvpn/internal/db"
|
||||
"lmvpn/internal/middleware"
|
||||
"lmvpn/internal/model"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
type loginRequest struct {
|
||||
Username string `json:"username" binding:"required"`
|
||||
Password string `json:"password" binding:"required"`
|
||||
}
|
||||
|
||||
type loginResponse struct {
|
||||
Token string `json:"token"`
|
||||
User userInfo `json:"user"`
|
||||
}
|
||||
|
||||
type userInfo struct {
|
||||
ID uint `json:"id"`
|
||||
Username string `json:"username"`
|
||||
Role string `json:"role"`
|
||||
}
|
||||
|
||||
func Login(c *gin.Context) {
|
||||
var req loginRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "请输入用户名和密码"})
|
||||
return
|
||||
}
|
||||
|
||||
var user model.User
|
||||
if err := db.DB.Where("username = ?", req.Username).First(&user).Error; err != nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "用户名或密码错误"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(req.Password)); err != nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "用户名或密码错误"})
|
||||
return
|
||||
}
|
||||
|
||||
if user.Status != 1 {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "账号已被禁用"})
|
||||
return
|
||||
}
|
||||
|
||||
token, err := middleware.GenerateToken(user.ID, user.Username, user.Role)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "生成令牌失败"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, loginResponse{
|
||||
Token: token,
|
||||
User: userInfo{
|
||||
ID: user.ID,
|
||||
Username: user.Username,
|
||||
Role: user.Role,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func Me(c *gin.Context) {
|
||||
userID, _ := c.Get("user_id")
|
||||
username, _ := c.Get("username")
|
||||
role, _ := c.Get("role")
|
||||
|
||||
c.JSON(http.StatusOK, userInfo{
|
||||
ID: userID.(uint),
|
||||
Username: username.(string),
|
||||
Role: role.(string),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
const (
|
||||
jwtSecret = "lmvpn-jwt-secret-key-2024"
|
||||
tokenExpire = 24 * time.Hour
|
||||
)
|
||||
|
||||
type Claims struct {
|
||||
UserID uint `json:"user_id"`
|
||||
Username string `json:"username"`
|
||||
Role string `json:"role"`
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
func GenerateToken(userID uint, username, role string) (string, error) {
|
||||
claims := Claims{
|
||||
UserID: userID,
|
||||
Username: username,
|
||||
Role: role,
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(tokenExpire)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
},
|
||||
}
|
||||
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
return token.SignedString([]byte(jwtSecret))
|
||||
}
|
||||
|
||||
func ParseToken(tokenStr string) (*Claims, error) {
|
||||
token, err := jwt.ParseWithClaims(tokenStr, &Claims{}, func(token *jwt.Token) (interface{}, error) {
|
||||
return []byte(jwtSecret), nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if claims, ok := token.Claims.(*Claims); ok && token.Valid {
|
||||
return claims, nil
|
||||
}
|
||||
|
||||
return nil, jwt.ErrSignatureInvalid
|
||||
}
|
||||
|
||||
func AuthMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
authHeader := c.GetHeader("Authorization")
|
||||
if authHeader == "" {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "未提供认证令牌"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
parts := strings.SplitN(authHeader, " ", 2)
|
||||
if len(parts) != 2 || parts[0] != "Bearer" {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "认证格式错误"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
claims, err := ParseToken(parts[1])
|
||||
if err != nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "令牌无效或已过期"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
c.Set("user_id", claims.UserID)
|
||||
c.Set("username", claims.Username)
|
||||
c.Set("role", claims.Role)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,8 @@ import (
|
||||
|
||||
"lmvpn/internal/config"
|
||||
"lmvpn/internal/db"
|
||||
"lmvpn/internal/handler"
|
||||
"lmvpn/internal/middleware"
|
||||
"lmvpn/internal/vpn"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -30,6 +32,14 @@ func main() {
|
||||
|
||||
r.GET("/ws", vpn.HandleWS)
|
||||
|
||||
r.POST("/api/login", handler.Login)
|
||||
|
||||
auth := r.Group("/api")
|
||||
auth.Use(middleware.AuthMiddleware())
|
||||
{
|
||||
auth.GET("/me", handler.Me)
|
||||
}
|
||||
|
||||
fs := http.FileServer(http.Dir("./dist"))
|
||||
r.Use(func(c *gin.Context) {
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/ws") {
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
import os
|
||||
import json
|
||||
import pytest
|
||||
import logging
|
||||
from websockets.asyncio.client import connect
|
||||
|
||||
logging.basicConfig(level=logging.DEBUG)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
WS_URL = os.environ.get("LMVPN_WS_URL", "ws://localhost:8080/ws")
|
||||
AUTH_USERNAME = os.environ.get("LMVPN_AUTH_USERNAME", "admin")
|
||||
AUTH_PASSWORD = os.environ.get("LMVPN_AUTH_PASSWORD", "admin123")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def raw_ws():
|
||||
"""仅建立 WebSocket 连接，不认证"""
|
||||
logger.info("connecting to %s", WS_URL)
|
||||
async with connect(WS_URL, ping_interval=None) as ws:
|
||||
yield ws
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def authenticated_ws():
|
||||
"""建立 WebSocket 连接并完成认证"""
|
||||
logger.info("connecting to %s", WS_URL)
|
||||
async with connect(WS_URL, ping_interval=None) as ws:
|
||||
auth_msg = json.dumps({"type": "auth", "username": AUTH_USERNAME, "password": AUTH_PASSWORD})
|
||||
await ws.send(auth_msg)
|
||||
resp = await ws.recv()
|
||||
resp_data = json.loads(resp)
|
||||
assert resp_data.get("type") == "auth_ok", f"认证失败: {resp_data}"
|
||||
logger.info("authenticated as %s", AUTH_USERNAME)
|
||||
yield ws
|
||||
@@ -0,0 +1,3 @@
|
||||
[pytest]
|
||||
asyncio_mode = auto
|
||||
testpaths = pytest
|
||||
@@ -0,0 +1,3 @@
|
||||
pytest>=8.0
|
||||
pytest-asyncio>=0.24
|
||||
websockets>=13.0
|
||||
@@ -0,0 +1,142 @@
|
||||
import json
|
||||
import asyncio
|
||||
import pytest
|
||||
from websockets import ConnectionClosed
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 认证测试
|
||||
# ============================================================
|
||||
|
||||
async def test_auth_success(raw_ws):
|
||||
"""正确凭据应返回 auth_ok"""
|
||||
msg = json.dumps({"type": "auth", "username": "admin", "password": "admin123"})
|
||||
await raw_ws.send(msg)
|
||||
resp = await raw_ws.recv()
|
||||
data = json.loads(resp)
|
||||
assert data["type"] == "auth_ok"
|
||||
|
||||
|
||||
async def test_auth_wrong_password(raw_ws):
|
||||
"""错误密码应返回 auth_err"""
|
||||
msg = json.dumps({"type": "auth", "username": "admin", "password": "wrong"})
|
||||
await raw_ws.send(msg)
|
||||
resp = await raw_ws.recv()
|
||||
data = json.loads(resp)
|
||||
assert data["type"] == "auth_err"
|
||||
|
||||
|
||||
async def test_auth_nonexistent_user(raw_ws):
|
||||
"""不存在的用户名应返回 auth_err"""
|
||||
msg = json.dumps({"type": "auth", "username": "no_such_user", "password": "123456"})
|
||||
await raw_ws.send(msg)
|
||||
resp = await raw_ws.recv()
|
||||
data = json.loads(resp)
|
||||
assert data["type"] == "auth_err"
|
||||
|
||||
|
||||
async def test_auth_non_json(raw_ws):
|
||||
"""首条消息不是 JSON，服务端应关闭连接"""
|
||||
await raw_ws.send("not json")
|
||||
with pytest.raises(ConnectionClosed):
|
||||
await raw_ws.recv()
|
||||
|
||||
|
||||
async def test_auth_missing_type(raw_ws):
|
||||
"""缺少 type 字段，服务端应关闭连接"""
|
||||
msg = json.dumps({"username": "admin", "password": "admin123"})
|
||||
await raw_ws.send(msg)
|
||||
with pytest.raises(ConnectionClosed):
|
||||
await raw_ws.recv()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 消息回显测试
|
||||
# ============================================================
|
||||
|
||||
async def test_echo_text(authenticated_ws):
|
||||
"""发送文本消息，验证回显一致"""
|
||||
await authenticated_ws.send("hello world")
|
||||
resp = await authenticated_ws.recv()
|
||||
assert resp == "hello world"
|
||||
|
||||
|
||||
async def test_echo_binary(authenticated_ws):
|
||||
"""发送二进制消息，验证回显一致"""
|
||||
data = b"\x00\x01\x02\xff\xfe\xfd"
|
||||
await authenticated_ws.send(data)
|
||||
resp = await authenticated_ws.recv()
|
||||
assert resp == data
|
||||
|
||||
|
||||
async def test_echo_multiple_messages(authenticated_ws):
|
||||
"""连续发送多条消息，逐一验证回显"""
|
||||
messages = ["msg1", "msg2", "msg3", "hello", "世界"]
|
||||
for m in messages:
|
||||
await authenticated_ws.send(m)
|
||||
resp = await authenticated_ws.recv()
|
||||
assert resp == m
|
||||
|
||||
|
||||
async def test_echo_json_message(authenticated_ws):
|
||||
"""发送 JSON 消息，验证回显一致"""
|
||||
payload = json.dumps({"cmd": "ping", "seq": 1, "data": "test"})
|
||||
await authenticated_ws.send(payload)
|
||||
resp = await authenticated_ws.recv()
|
||||
assert resp == payload
|
||||
|
||||
|
||||
async def test_echo_large_message(authenticated_ws):
|
||||
"""发送大消息（64KB），验证回显一致"""
|
||||
large_msg = "A" * 65536
|
||||
await authenticated_ws.send(large_msg)
|
||||
resp = await authenticated_ws.recv()
|
||||
assert resp == large_msg
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 并发测试
|
||||
# ============================================================
|
||||
|
||||
async def test_concurrent_connections():
|
||||
"""同时建立 5 个认证连接，各自发送消息验证回显"""
|
||||
from websockets.asyncio.client import connect
|
||||
import os
|
||||
|
||||
ws_url = os.environ.get("LMVPN_WS_URL", "ws://localhost:8080/ws")
|
||||
|
||||
async def one_session(uid):
|
||||
async with connect(ws_url, ping_interval=None) as ws:
|
||||
auth = json.dumps({"type": "auth", "username": "admin", "password": "admin123"})
|
||||
await ws.send(auth)
|
||||
resp = await ws.recv()
|
||||
assert json.loads(resp)["type"] == "auth_ok"
|
||||
|
||||
msg = f"msg_from_{uid}"
|
||||
await ws.send(msg)
|
||||
resp = await ws.recv()
|
||||
assert resp == msg
|
||||
|
||||
tasks = [one_session(i) for i in range(5)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 协议测试
|
||||
# ============================================================
|
||||
|
||||
async def test_connection_close_after_auth_failure(raw_ws):
|
||||
"""认证失败后服务端应关闭连接"""
|
||||
msg = json.dumps({"type": "auth", "username": "admin", "password": "wrong"})
|
||||
await raw_ws.send(msg)
|
||||
resp = await raw_ws.recv()
|
||||
assert json.loads(resp)["type"] == "auth_err"
|
||||
with pytest.raises(ConnectionClosed):
|
||||
await raw_ws.recv()
|
||||
|
||||
|
||||
async def test_ping_pong(authenticated_ws):
|
||||
"""发送 ping，服务端应在超时前回应 pong"""
|
||||
await authenticated_ws.ping()
|
||||
Reference in New Issue
Block a user