实现登录功能

2025-11-30 11:44:26 +08:00
parent 4f0815a5f9
commit 72aa2846e5
7 changed files with 441 additions and 38 deletions
--- a/http/server.go
+++ b/http/server.go
@@ -26,6 +26,11 @@ type Server struct {
 	shieldManager *shield.ShieldManager
 	server        *http.Server

+	// 会话管理相关字段
+	sessions      map[string]time.Time // 会话ID到过期时间的映射
+	sessionsMutex sync.Mutex           // 会话映射的互斥锁
+	sessionTTL    time.Duration        // 会话过期时间
+
 	// WebSocket相关字段
 	upgrader      websocket.Upgrader
 	clients       map[*websocket.Conn]bool
@@ -50,10 +55,15 @@ func NewServer(globalConfig *config.Config, dnsServer *dns.Server, shieldManager
 		},
 		clients:       make(map[*websocket.Conn]bool),
 		broadcastChan: make(chan []byte, 100),
+		// 会话管理初始化
+		sessions:   make(map[string]time.Time),
+		sessionTTL: 24 * time.Hour, // 会话有效期24小时
 	}

 	// 启动广播协程
 	go server.startBroadcastLoop()
+	// 启动会话清理协程
+	go server.cleanupSessionsLoop()

 	return server
 }
@@ -62,17 +72,26 @@ func NewServer(globalConfig *config.Config, dnsServer *dns.Server, shieldManager
 func (s *Server) Start() error {
 	mux := http.NewServeMux()

+	// 登录路由，不需要认证
+	mux.HandleFunc("/login", func(w http.ResponseWriter, r *http.Request) {
+		// 重定向到登录页面HTML
+		http.Redirect(w, r, "/login.html", http.StatusFound)
+	})
+
 	// API路由
 	if s.config.EnableAPI {
-		// 重定向/api到Swagger UI页面
-		mux.HandleFunc("/api", func(w http.ResponseWriter, r *http.Request) {
-			http.Redirect(w, r, "/api/index.html", http.StatusMovedPermanently)
-		})
+		// 登录API端点，不需要认证
+		mux.HandleFunc("/api/login", s.handleLogin)

-		// 注册所有API端点
-		mux.HandleFunc("/api/stats", s.handleStats)
-		mux.HandleFunc("/api/shield", s.handleShield)
-		mux.HandleFunc("/api/shield/localrules", func(w http.ResponseWriter, r *http.Request) {
+		// 重定向/api到Swagger UI页面
+		mux.HandleFunc("/api", s.loginRequired(func(w http.ResponseWriter, r *http.Request) {
+			http.Redirect(w, r, "/api/index.html", http.StatusMovedPermanently)
+		}))
+
+		// 注册所有API端点，应用登录中间件
+		mux.HandleFunc("/api/stats", s.loginRequired(s.handleStats))
+		mux.HandleFunc("/api/shield", s.loginRequired(s.handleShield))
+		mux.HandleFunc("/api/shield/localrules", s.loginRequired(func(w http.ResponseWriter, r *http.Request) {
 			w.Header().Set("Content-Type", "application/json")
 			if r.Method == http.MethodGet {
 				localRules := s.shieldManager.GetLocalRules()
@@ -80,8 +99,8 @@ func (s *Server) Start() error {
 				return
 			}
 			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
-		})
-		mux.HandleFunc("/api/shield/remoterules", func(w http.ResponseWriter, r *http.Request) {
+		}))
+		mux.HandleFunc("/api/shield/remoterules", s.loginRequired(func(w http.ResponseWriter, r *http.Request) {
 			w.Header().Set("Content-Type", "application/json")
 			if r.Method == http.MethodGet {
 				remoteRules := s.shieldManager.GetRemoteRules()
@@ -89,45 +108,57 @@ func (s *Server) Start() error {
 				return
 			}
 			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
-		})
-		mux.HandleFunc("/api/shield/hosts", s.handleShieldHosts)
-		mux.HandleFunc("/api/shield/blacklists", s.handleShieldBlacklists)
-		mux.HandleFunc("/api/query", s.handleQuery)
-		mux.HandleFunc("/api/status", s.handleStatus)
-		mux.HandleFunc("/api/config", s.handleConfig)
-		mux.HandleFunc("/api/config/restart", s.handleRestart)
+		}))
+		mux.HandleFunc("/api/shield/hosts", s.loginRequired(s.handleShieldHosts))
+		mux.HandleFunc("/api/shield/blacklists", s.loginRequired(s.handleShieldBlacklists))
+		mux.HandleFunc("/api/query", s.loginRequired(s.handleQuery))
+		mux.HandleFunc("/api/status", s.loginRequired(s.handleStatus))
+		mux.HandleFunc("/api/config", s.loginRequired(s.handleConfig))
+		mux.HandleFunc("/api/config/restart", s.loginRequired(s.handleRestart))
 		// 添加统计相关接口
-		mux.HandleFunc("/api/top-blocked", s.handleTopBlockedDomains)
-		mux.HandleFunc("/api/top-resolved", s.handleTopResolvedDomains)
-		mux.HandleFunc("/api/top-clients", s.handleTopClients)
-		mux.HandleFunc("/api/top-domains", s.handleTopDomains)
-		mux.HandleFunc("/api/recent-blocked", s.handleRecentBlockedDomains)
-		mux.HandleFunc("/api/hourly-stats", s.handleHourlyStats)
-		mux.HandleFunc("/api/daily-stats", s.handleDailyStats)
-		mux.HandleFunc("/api/monthly-stats", s.handleMonthlyStats)
-		mux.HandleFunc("/api/query/type", s.handleQueryTypeStats)
+		mux.HandleFunc("/api/top-blocked", s.loginRequired(s.handleTopBlockedDomains))
+		mux.HandleFunc("/api/top-resolved", s.loginRequired(s.handleTopResolvedDomains))
+		mux.HandleFunc("/api/top-clients", s.loginRequired(s.handleTopClients))
+		mux.HandleFunc("/api/top-domains", s.loginRequired(s.handleTopDomains))
+		mux.HandleFunc("/api/recent-blocked", s.loginRequired(s.handleRecentBlockedDomains))
+		mux.HandleFunc("/api/hourly-stats", s.loginRequired(s.handleHourlyStats))
+		mux.HandleFunc("/api/daily-stats", s.loginRequired(s.handleDailyStats))
+		mux.HandleFunc("/api/monthly-stats", s.loginRequired(s.handleMonthlyStats))
+		mux.HandleFunc("/api/query/type", s.loginRequired(s.handleQueryTypeStats))
 		// 日志统计相关接口
-		mux.HandleFunc("/api/logs/stats", s.handleLogsStats)
-		mux.HandleFunc("/api/logs/query", s.handleLogsQuery)
-		mux.HandleFunc("/api/logs/count", s.handleLogsCount)
+		mux.HandleFunc("/api/logs/stats", s.loginRequired(s.handleLogsStats))
+		mux.HandleFunc("/api/logs/query", s.loginRequired(s.handleLogsQuery))
+		mux.HandleFunc("/api/logs/count", s.loginRequired(s.handleLogsCount))
 		// WebSocket端点
-		mux.HandleFunc("/ws/stats", s.handleWebSocketStats)
+		mux.HandleFunc("/ws/stats", s.loginRequired(s.handleWebSocketStats))

 		// 将/api/下的静态文件服务指向static/api目录，放在最后以避免覆盖API端点
 		apiFileServer := http.FileServer(http.Dir("./static/api"))
-		mux.Handle("/api/", http.StripPrefix("/api", apiFileServer))
+		mux.Handle("/api/", s.loginRequired(http.StripPrefix("/api", apiFileServer).ServeHTTP))
 	}

 	// 自定义静态文件服务处理器，用于禁用浏览器缓存，放在API路由之后
 	fileServer := http.FileServer(http.Dir("./static"))
-	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+
+	// 单独处理login.html，不需要登录
+	mux.HandleFunc("/login.html", func(w http.ResponseWriter, r *http.Request) {
+		// 添加Cache-Control头，禁用浏览器缓存
+		w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
+		w.Header().Set("Pragma", "no-cache")
+		w.Header().Set("Expires", "Thu, 01 Jan 1970 00:00:00 GMT")
+		// 直接提供login.html文件
+		http.ServeFile(w, r, "./static/login.html")
+	})
+
+	// 其他静态文件需要登录
+	mux.HandleFunc("/", s.loginRequired(func(w http.ResponseWriter, r *http.Request) {
 		// 添加Cache-Control头，禁用浏览器缓存
 		w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
 		w.Header().Set("Pragma", "no-cache")
 		w.Header().Set("Expires", "Thu, 01 Jan 1970 00:00:00 GMT")
 		// 使用StripPrefix处理路径
 		http.StripPrefix("/", fileServer).ServeHTTP(w, r)
-	})
+	}))

 	s.server = &http.Server{
 		Addr:         fmt.Sprintf("%s:%d", s.config.Host, s.config.Port),
@@ -378,6 +409,78 @@ func (s *Server) startBroadcastLoop() {
 	}
 }

+// cleanupSessionsLoop 定期清理过期会话
+func (s *Server) cleanupSessionsLoop() {
+	for {
+		time.Sleep(1 * time.Hour) // 每小时清理一次
+		s.sessionsMutex.Lock()
+		now := time.Now()
+		for sessionID, expiryTime := range s.sessions {
+			if now.After(expiryTime) {
+				delete(s.sessions, sessionID)
+			}
+		}
+		s.sessionsMutex.Unlock()
+	}
+}
+
+// isAuthenticated 检查用户是否已认证
+func (s *Server) isAuthenticated(r *http.Request) bool {
+	// 从Cookie中获取会话ID
+	cookie, err := r.Cookie("session_id")
+	if err != nil {
+		return false
+	}
+
+	sessionID := cookie.Value
+	s.sessionsMutex.Lock()
+	defer s.sessionsMutex.Unlock()
+
+	// 检查会话是否存在且未过期
+	expiryTime, exists := s.sessions[sessionID]
+	if !exists {
+		return false
+	}
+
+	if time.Now().After(expiryTime) {
+		// 会话已过期，删除它
+		delete(s.sessions, sessionID)
+		return false
+	}
+
+	// 延长会话有效期
+	s.sessions[sessionID] = time.Now().Add(s.sessionTTL)
+	return true
+}
+
+// loginRequired 登录中间件
+func (s *Server) loginRequired(next http.HandlerFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		// 检查是否为登录页面或登录API，允许直接访问
+		if r.URL.Path == "/login" || r.URL.Path == "/api/login" {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		// 检查是否已认证
+		if !s.isAuthenticated(r) {
+			// 如果是API请求，返回401错误
+			if strings.HasPrefix(r.URL.Path, "/api/") || strings.HasPrefix(r.URL.Path, "/ws/") {
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusUnauthorized)
+				json.NewEncoder(w).Encode(map[string]string{"error": "未授权访问"})
+				return
+			}
+			// 否则重定向到登录页面
+			http.Redirect(w, r, "/login", http.StatusFound)
+			return
+		}
+
+		// 已认证，继续处理请求
+		next.ServeHTTP(w, r)
+	}
+}
+
 // handleTopBlockedDomains 处理TOP屏蔽域名请求
 func (s *Server) handleTopBlockedDomains(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodGet {
@@ -1312,3 +1415,56 @@ func (s *Server) handleRestart(w http.ResponseWriter, r *http.Request) {
 	json.NewEncoder(w).Encode(map[string]string{"status": "success", "message": "服务已重启"})
 	logger.Info("服务重启成功")
 }
+
+// handleLogin 处理登录请求
+func (s *Server) handleLogin(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodPost {
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	// 解析请求体
+	var loginData struct {
+		Username string `json:"username"`
+		Password string `json:"password"`
+	}
+
+	if err := json.NewDecoder(r.Body).Decode(&loginData); err != nil {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusBadRequest)
+		json.NewEncoder(w).Encode(map[string]string{"error": "无效的请求体"})
+		return
+	}
+
+	// 验证用户名和密码
+	if loginData.Username != s.config.Username || loginData.Password != s.config.Password {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusUnauthorized)
+		json.NewEncoder(w).Encode(map[string]string{"error": "用户名或密码错误"})
+		return
+	}
+
+	// 生成会话ID
+	sessionID := fmt.Sprintf("%d_%d", time.Now().UnixNano(), len(s.sessions))
+
+	// 保存会话
+	s.sessionsMutex.Lock()
+	s.sessions[sessionID] = time.Now().Add(s.sessionTTL)
+	s.sessionsMutex.Unlock()
+
+	// 设置Cookie
+	cookie := &http.Cookie{
+		Name:     "session_id",
+		Value:    sessionID,
+		Path:     "/",
+		Expires:  time.Now().Add(s.sessionTTL),
+		HttpOnly: true,
+		Secure:   false, // 开发环境下使用false，生产环境应使用true
+	}
+	http.SetCookie(w, cookie)
+
+	// 返回成功响应
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(map[string]string{"status": "success", "message": "登录成功"})
+	logger.Info(fmt.Sprintf("用户 %s 登录成功", loginData.Username))
+}