aiops/backend/controller/rdp.go

package controller

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"image"
	"image/png"
	"log"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"ezops/middleware"
	"ezops/model"
	"ezops/pkg/crypto"
	rdpengine "ezops/pkg/rdp"
	"ezops/pkg/wshub"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"gorm.io/gorm"
)

// RDPHandler 提供后端渲染型 RDP 会话，浏览器只负责 Canvas 与输入事件 // Logan.Li
type RDPHandler struct {
	db      *gorm.DB
	crypt   *crypto.SecretBox
	wsMgr   *wshub.WSManager
	engine  rdpengine.Engine
	clients sync.Map // sessionID -> context.CancelFunc // Logan.Li
	live    sync.Map // sessionID -> *rdpLiveBroadcaster // Logan.Li
}

const rdpAuditMaxText = 500
const rdpMaxWidth = 1024
const rdpMaxHeight = 800

func NewRDPHandler(db *gorm.DB, masterKey string, wsMgr *wshub.WSManager) *RDPHandler {
	return &RDPHandler{
		db:     db,
		crypt:  crypto.NewSecretBox(masterKey),
		wsMgr:  wsMgr,
		engine: rdpengine.NewEngine(),
	}
}

func (h *RDPHandler) KillSession(sessionID uint) bool {
	if val, ok := h.clients.Load(sessionID); ok {
		cancel := val.(context.CancelFunc)
		cancel()
		h.clients.Delete(sessionID)
		return true
	}
	return false
}

func (h *RDPHandler) ServeLive(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("sessionId"), 10, 64)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "无效的 RDP 会话ID"})
		return
	}
	var session model.RDPSession
	if err := h.db.First(&session, id).Error; err != nil {
		c.JSON(http.StatusNotFound, model.Response{Code: 404, Message: "RDP 会话不存在"})
		return
	}
	if session.Status != "active" {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "RDP 会话已结束，请使用历史回放"})
		return
	}
	val, ok := h.live.Load(uint(id))
	if !ok {
		c.JSON(http.StatusNotFound, model.Response{Code: 404, Message: "RDP 实时流不存在"})
		return
	}
	conn, info, err := h.wsMgr.Upgrade(c, wshub.TypeAudit, middleware.GetUserID(c), middleware.GetUsername(c))
	if err != nil {
		return
	}
	defer h.wsMgr.Unregister(info)
	hub := val.(*rdpLiveBroadcaster)
	ch := hub.subscribe()
	defer hub.unsubscribe(ch)

	_ = writeRDPFrame(h.wsMgr, conn, rdpengine.ServerFrame{
		Type:    "session",
		Message: strconv.FormatUint(uint64(session.ID), 10),
		Width:   session.Width,
		Height:  session.Height,
	})
	for {
		select {
		case frame, ok := <-ch:
			if !ok {
				return
			}
			if err := writeRDPFrame(h.wsMgr, conn, frame); err != nil {
				return
			}
		case <-c.Request.Context().Done():
			return
		}
	}
}

func (h *RDPHandler) Serve(c *gin.Context) {
	serverID := c.Query("id")
	if serverID == "" {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "缺少服务器ID"})
		return
	}

	width := clampPositiveInt(c.Query("width"), 1280, rdpMaxWidth)
	height := clampPositiveInt(c.Query("height"), 800, rdpMaxHeight)

	var server model.Server
	if err := h.db.First(&server, serverID).Error; err != nil {
		c.JSON(http.StatusNotFound, model.Response{Code: 404, Message: "服务器不存在"})
		return
	}
	normalizeServerRemoteAccess(&server)
	if !isRDPServer(server) {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "该资产未配置为 RDP/Windows"})
		return
	}

	rdpCred, err := h.loadRDPCredential(server)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: err.Error()})
		return
	}

	userID := middleware.GetUserID(c)
	username := middleware.GetUsername(c)
	conn, info, err := h.wsMgr.Upgrade(c, "rdp", userID, username)
	if err != nil {
		return
	}
	log.Printf("[RDP] ws connected session pending server=%s ip=%s user=%s size=%dx%d", server.Hostname, server.IP, username, width, height)
	defer h.wsMgr.Unregister(info)

	var userIDPtr *uint
	if userID > 0 {
		userIDPtr = &userID
	}
	port := server.RDPPort
	if port == 0 {
		port = 3389
	}
	session := model.RDPSession{
		UserID:     userIDPtr,
		Username:   username,
		ServerID:   server.ID,
		ServerName: server.Hostname,
		ServerIP:   server.IP,
		ServerPort: port,
		ClientIP:   c.ClientIP(),
		Width:      width,
		Height:     height,
		Status:     "active",
		StartTime:  time.Now(),
	}
	if err := h.db.Create(&session).Error; err != nil {
		_ = conn.WriteJSON(rdpengine.ServerFrame{Type: "error", Message: "创建 RDP 会话失败"})
		return
	}

	ctx, cancel := context.WithCancel(c.Request.Context())
	h.clients.Store(session.ID, cancel)
	liveHub := newRDPLiveBroadcaster()
	h.live.Store(session.ID, liveHub)
	defer func() {
		cancel()
		h.clients.Delete(session.ID)
		h.live.Delete(session.ID)
		liveHub.close()
	}()

	events := make(chan rdpengine.ClientEvent, 128)
	sink := &wsFrameSink{conn: conn, mgr: h.wsMgr, recorder: rdpengine.NewMemoryRecorder(500), live: liveHub, sessionID: session.ID}
	_ = sink.SendFrame(rdpengine.ServerFrame{Type: "session", Message: strconv.Itoa(int(session.ID)), Width: width, Height: height})
	engineDone := make(chan error, 1)
	go func() {
		log.Printf("[RDP] engine start session=%d server=%s:%d size=%dx%d", session.ID, server.IP, port, width, height)
		engineDone <- h.engine.Run(ctx, rdpengine.Config{
			Host:     server.IP,
			Port:     port,
			Username: rdpCred.Username,
			Password: rdpCred.Password,
			Domain:   rdpCred.Domain,
			Security: rdpCred.Security,
			Width:    width,
			Height:   height,
		}, sink, events)
	}()

	status := "closed"
	errMsg := ""
	defer func() {
		cancel()
		end := time.Now()
		duration := int(end.Sub(session.StartTime).Seconds())
		recordData := sink.recorder.JSON()
		updates := map[string]interface{}{
			"status":      status,
			"duration":    duration,
			"end_time":    end,
			"record_data": recordData,
			"record_size": len(recordData),
		}
		if errMsg != "" {
			updates["error"] = errMsg
		}
		_ = h.db.Model(&model.RDPSession{}).Where("id = ?", session.ID).Updates(updates).Error
	}()

	audit := newRDPOperationRecorder(h.db, session, userIDPtr, username)
	defer audit.Flush()

	readDone := make(chan struct{})
	go func() {
		defer close(readDone)
		for {
			var event rdpengine.ClientEvent
			if err := conn.ReadJSON(&event); err != nil {
				return
			}
			audit.Record(event)
			enqueueRDPEvent(events, event)
		}
	}()

	select {
	case err := <-engineDone:
		if err != nil && !errors.Is(err, context.Canceled) {
			status = "failed"
			errMsg = err.Error()
			log.Printf("[RDP] engine failed session=%d err=%v", session.ID, err)
		} else {
			log.Printf("[RDP] engine closed session=%d err=%v", session.ID, err)
		}
	case <-readDone:
		log.Printf("[RDP] ws read closed session=%d", session.ID)
	}
}

func enqueueRDPEvent(events chan rdpengine.ClientEvent, event rdpengine.ClientEvent) {
	select {
	case events <- event:
		return
	default:
	}
	if event.Type == "mousemove" || event.Type == "mouse_move" {
		// 队列已满时丢弃当前移动事件，避免旧坐标堆积拖慢点击/拖拽响应。 // Logan.Li
		return
	}
	// 关键输入不能被旧 mousemove 堵住，优先清理一批移动事件后再投递。 // Logan.Li
	for i := 0; i < cap(events)/2; i++ {
		select {
		case queued := <-events:
			if queued.Type != "mousemove" && queued.Type != "mouse_move" {
				select {
				case events <- queued:
				default:
				}
			}
		default:
			i = cap(events)
		}
		select {
		case events <- event:
			return
		default:
		}
	}
	select {
	case events <- event:
	case <-time.After(200 * time.Millisecond):
		log.Printf("[RDP] input queue congested, drop critical event type=%s", event.Type)
	}
}

type rdpOperationRecorder struct {
	db          *gorm.DB
	session     model.RDPSession
	userID      *uint
	username    string
	inputBuffer string
	lastInput   time.Time
}

func newRDPOperationRecorder(db *gorm.DB, session model.RDPSession, userID *uint, username string) *rdpOperationRecorder {
	return &rdpOperationRecorder{db: db, session: session, userID: userID, username: username}
}

func (r *rdpOperationRecorder) Record(event rdpengine.ClientEvent) {
	now := time.Now()
	if r.inputBuffer != "" && now.Sub(r.lastInput) > 2*time.Second {
		r.Flush()
	}
	switch event.Type {
	case "key":
		r.recordKey(event, now)
	case "mousemove", "mouse_move", "mousedown", "mouse_down", "mouseup", "mouse_up", "wheel":
		return // 鼠标轨迹只进图形录屏，不写命令审计，避免 RDP 命令表噪声。 // Logan.Li
	case "clipboard":
		r.Flush()
		r.create("[RDP粘贴] " + truncateAuditText(event.Text, rdpAuditMaxText))
	case "file_upload":
		r.Flush()
		r.create("[RDP上传] " + truncateAuditText(event.Text, rdpAuditMaxText))
	case "file_download":
		r.Flush()
		r.create("[RDP下载] " + truncateAuditText(event.Text, rdpAuditMaxText))
	}
}

func (r *rdpOperationRecorder) recordKey(event rdpengine.ClientEvent, now time.Time) {
	if !event.Down {
		return
	}
	if event.Key == "Control" || event.Key == "Shift" || event.Key == "Alt" || event.Key == "Meta" {
		return
	}
	if event.Code == "Enter" || event.Code == "NumpadEnter" {
		if strings.TrimSpace(r.inputBuffer) != "" {
			r.Flush()
			return
		}
		r.create("[RDP按键] Enter")
		return
	}
	if event.Code == "Backspace" {
		if len(r.inputBuffer) > 0 {
			r.inputBuffer = r.inputBuffer[:len(r.inputBuffer)-1]
			r.lastInput = now
		}
		return
	}
	if len(event.Key) == 1 {
		r.inputBuffer += event.Key
		r.lastInput = now
		if len(r.inputBuffer) >= rdpAuditMaxText {
			r.Flush()
		}
		return
	}
	if isAuditableRDPKey(event.Key) {
		r.Flush()
		r.create("[RDP按键] " + event.Key)
	}
}

// isAuditableRDPKey 只保留对 Windows 操作有审计意义的非文本按键。 // Logan.Li
func isAuditableRDPKey(key string) bool {
	switch key {
	case "Tab", "Escape", "Delete", "Home", "End", "PageUp", "PageDown":
		return true
	}
	if strings.HasPrefix(key, "F") {
		n, err := strconv.Atoi(strings.TrimPrefix(key, "F"))
		return err == nil && n >= 1 && n <= 24
	}
	return false
}

func (r *rdpOperationRecorder) Flush() {
	if strings.TrimSpace(r.inputBuffer) == "" {
		r.inputBuffer = ""
		return
	}
	r.create("[RDP输入] " + truncateAuditText(r.inputBuffer, rdpAuditMaxText))
	r.inputBuffer = ""
}

func (r *rdpOperationRecorder) create(command string) {
	cmd := model.SessionCommand{
		SessionID:   r.session.ID,
		SessionType: "rdp",
		UserID:      r.userID,
		Username:    r.username,
		ServerID:    r.session.ServerID,
		ServerName:  r.session.ServerName,
		Command:     command,
		Timestamp:   time.Now(),
	}
	_ = r.db.Create(&cmd).Error
}

func (h *RDPHandler) loadRDPCredential(server model.Server) (model.RDPSecretContent, error) {
	return loadRDPCredential(h.crypt, server)
}

func parsePositiveInt(raw string, fallback int) int {
	v, err := strconv.Atoi(raw)
	if err != nil || v <= 0 {
		return fallback
	}
	return v
}

func clampPositiveInt(raw string, fallback int, max int) int {
	v := parsePositiveInt(raw, fallback)
	if max > 0 && v > max {
		return max
	}
	return v
}

type wsFrameSink struct {
	conn       *websocket.Conn
	mgr        *wshub.WSManager
	mu         sync.Mutex
	recorder   *rdpengine.MemoryRecorder
	live       *rdpLiveBroadcaster
	sessionID  uint
	bitmaps    int
	lastLog    time.Time
	lastRecord time.Time
}

func (s *wsFrameSink) SendFrame(frame rdpengine.ServerFrame) error {
	s.recorder.Append(frame)
	if s.live != nil {
		s.live.publish(frame)
	}
	if snap, ok := s.snapshotFrame(frame); ok {
		s.recorder.AppendSnapshot(snap)
	}
	if frame.Type == "bitmap" {
		s.bitmaps++
		if s.bitmaps <= 3 || time.Since(s.lastLog) > 5*time.Second {
			log.Printf("[RDP] bitmap session=%d count=%d rect=%dx%d@%d,%d image=%d rgba=%d", s.sessionID, s.bitmaps, frame.Width, frame.Height, frame.X, frame.Y, len(frame.Image), len(frame.RGBA))
			s.lastLog = time.Now()
		}
	} else if frame.Type == "status" || frame.Type == "error" || frame.Type == "session" {
		log.Printf("[RDP] frame session=%d type=%s msg=%q size=%dx%d", s.sessionID, frame.Type, frame.Message, frame.Width, frame.Height)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	return writeRDPFrame(s.mgr, s.conn, frame)
}

func (s *wsFrameSink) snapshotFrame(frame rdpengine.ServerFrame) (rdpengine.ServerFrame, bool) {
	if frame.Type != "bitmap" || len(frame.RGBA) == 0 || frame.Width <= 0 || frame.Height <= 0 {
		return rdpengine.ServerFrame{}, false
	}
	now := time.Now()
	if !s.lastRecord.IsZero() && now.Sub(s.lastRecord) < 500*time.Millisecond {
		return rdpengine.ServerFrame{}, false
	}
	s.lastRecord = now
	img := image.NewRGBA(image.Rect(0, 0, frame.Width, frame.Height))
	copy(img.Pix, frame.RGBA)
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		return rdpengine.ServerFrame{}, false
	}
	frame.RGBA = nil
	frame.Image = "data:image/png;base64," + base64.StdEncoding.EncodeToString(buf.Bytes())
	return frame, true
}

func writeRDPFrame(mgr *wshub.WSManager, conn *websocket.Conn, frame rdpengine.ServerFrame) error {
	mgr.SetWriteDeadline(conn)
	if frame.Type == "bitmap" && len(frame.RGBA) > 0 {
		return conn.WriteMessage(websocket.BinaryMessage, encodeRDPBitmapBinary(frame))
	}
	return conn.WriteJSON(frame)
}

func encodeRDPBitmapBinary(frame rdpengine.ServerFrame) []byte {
	headerLen := 20
	out := make([]byte, headerLen+len(frame.RGBA))
	copy(out[0:4], []byte{'R', 'D', 'P', 'R'})
	binary.LittleEndian.PutUint32(out[4:8], uint32(frame.X))
	binary.LittleEndian.PutUint32(out[8:12], uint32(frame.Y))
	binary.LittleEndian.PutUint32(out[12:16], uint32(frame.Width))
	binary.LittleEndian.PutUint32(out[16:20], uint32(frame.Height))
	copy(out[headerLen:], frame.RGBA)
	return out
}