aiops/backend/pkg/rdp/session_engine.go

package rdp

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	"ezops/pkg/rdp/display"
	"ezops/pkg/rdp/trace"
	"ezops/pkg/rdp/transport"
)

// SessionEngine 是自研完整 RDP 图形会话引擎 // Logan.Li
type SessionEngine struct{}

func (e *SessionEngine) Run(ctx context.Context, cfg Config, sink FrameSink, events <-chan ClientEvent) error {
	if cfg.Port == 0 {
		cfg.Port = 3389
	}
	trace.Logf("engine=rdp connect target=%s:%d size=%dx%d security=%s", cfg.Host, cfg.Port, cfg.Width, cfg.Height, cfg.Security)
	_ = sink.SendFrame(ServerFrame{
		Type:    "status",
		Message: fmt.Sprintf("正在连接 RDP %s:%d", cfg.Host, cfg.Port),
	})

	sess, err := OpenSession(ctx, cfg)
	if err != nil {
		trace.Logf("engine=rdp open failed target=%s:%d err=%v", cfg.Host, cfg.Port, err)
		_ = sink.SendFrame(ServerFrame{Type: "error", Message: "RDP 会话建立失败：" + err.Error()})
		return err
	}
	defer func() {
		if c, ok := sess.Conn.(io.Closer); ok {
			_ = c.Close()
		}
	}()

	_ = sink.SendFrame(ServerFrame{Type: "status", Message: "RDP 图形会话已建立，等待桌面画面..."})
	trace.Logf("engine=rdp open ok target=%s:%d", cfg.Host, cfg.Port)
	writeMu := &sync.Mutex{}
	input := NewInputSenderWithLock(sess.Conn, writeMu)
	receiver := NewFastPathReceiver()

	readErr := make(chan error, 1)
	go func() {
		readErr <- e.readLoop(ctx, sess, receiver, sink, writeMu)
	}()
	go e.inputLoop(ctx, input, sink, events)

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case err := <-readErr:
			if err != nil {
				_ = sink.SendFrame(ServerFrame{Type: "error", Message: err.Error()})
			}
			return err
		}
	}
}

func (e *SessionEngine) inputLoop(ctx context.Context, input *InputSender, sink FrameSink, events <-chan ClientEvent) {
	for {
		select {
		case <-ctx.Done():
			return
		case event, ok := <-events:
			if !ok {
				return
			}
			if event.Type != "mousemove" && event.Type != "mouse_move" {
				trace.Logf("engine=rdp input type=%s x=%d y=%d button=%d code=%s down=%t", event.Type, event.X, event.Y, event.Button, event.Code, event.Down)
			}
			if err := e.handleEvent(input, sink, event); err != nil {
				_ = sink.SendFrame(ServerFrame{Type: "status", Message: "输入事件失败：" + err.Error()})
			}
		}
	}
}

func (e *SessionEngine) readLoop(ctx context.Context, sess *Session, receiver *FastPathReceiver, sink FrameSink, writeMu *sync.Mutex) error {
	framebuffer := display.NewFramebuffer(sess.Width, sess.Height)
	ticker := time.NewTicker(16 * time.Millisecond)
	defer ticker.Stop()
	refreshRequested := false
	flushFrame := func() {
		frame, ok := framebuffer.FlushRGBA()
		if !ok {
			return
		}
		_ = sink.SendFrame(ServerFrame{
			Type:   "bitmap",
			X:      frame.X,
			Y:      frame.Y,
			Width:  frame.Width,
			Height: frame.Height,
			RGBA:   frame.RGBA,
		})
	}
	sendBitmaps := func(rects []BitmapRect) {
		for _, rect := range rects {
			if rect.Width <= 0 || rect.Height <= 0 {
				continue
			}
			if !refreshRequested && isBlackFullScreen(rect, sess.Width, sess.Height) {
				refreshRequested = true
				go func() {
					time.Sleep(400 * time.Millisecond)
					if err := e.requestRefresh(sess, writeMu); err != nil {
						trace.Logf("engine=rdp refresh request failed err=%v", err)
					}
				}()
			}
			framebuffer.Draw(display.Bitmap{
				X:            rect.DestLeft,
				Y:            rect.DestTop,
				Width:        rect.Width,
				Height:       rect.Height,
				BitsPerPixel: rect.BitsPerPixel,
				Data:         rect.Data,
			})
		}
	}

	type packetResult struct {
		packet WirePacket
		err    error
	}
	packetCh := make(chan packetResult, 1)
	go func() {
		for {
			pkt, err := ReadWirePacket(sess.Conn)
			select {
			case packetCh <- packetResult{packet: pkt, err: err}:
			case <-ctx.Done():
				return
			}
			if err != nil {
				return
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			flushFrame()
			return ctx.Err()
		case <-ticker.C:
			flushFrame()
		case got := <-packetCh:
			if got.err != nil {
				flushFrame()
				return got.err
			}
			var rects []BitmapRect
			var err error
			if got.packet.FastPath {
				rects, err = receiver.Handle(got.packet.SecFlag, got.packet.Payload)
			} else {
				rects, err = receiver.HandleSlowPath(got.packet.Payload)
			}
			if err != nil {
				continue
			}
			if err := e.ackSurfaceFrames(sess, receiver.TakeFrameIDs(), writeMu); err != nil {
				trace.Logf("engine=rdp frame ack failed err=%v", err)
			}
			if len(rects) > 0 {
				trace.Logf("engine=rdp bitmap rects=%d fastpath=%t", len(rects), got.packet.FastPath)
				sendBitmaps(rects)
			}
		}
	}
}

func isBlackFullScreen(rect BitmapRect, width, height int) bool {
	if rect.DestLeft != 0 || rect.DestTop != 0 || rect.Width < width || rect.Height < height {
		return false
	}
	limit := len(rect.Data)
	if limit > 256*1024 {
		limit = 256 * 1024
	}
	for i := 0; i < limit; i++ {
		if rect.Data[i] > 4 {
			return false
		}
	}
	return limit > 0
}

func (e *SessionEngine) requestRefresh(sess *Session, writeMu *sync.Mutex) error {
	body := buildRefreshRectBody(Config{Width: sess.Width, Height: sess.Height})
	pdu := buildDataPDU(sess.UserID, sess.SharedID, pduType2RefreshRect, body)
	trace.Logf("engine=rdp refresh rect %dx%d", sess.Width, sess.Height)
	writeMu.Lock()
	err := transport.WriteX224Data(sess.Conn, buildMCSSendDataRequest(sess.UserID, mcsGlobalChannelID, pdu))
	writeMu.Unlock()
	return err
}

func (e *SessionEngine) ackSurfaceFrames(sess *Session, frameIDs []uint32, writeMu *sync.Mutex) error {
	for _, frameID := range frameIDs {
		body := []byte{byte(frameID), byte(frameID >> 8), byte(frameID >> 16), byte(frameID >> 24)}
		pdu := buildDataPDU(sess.UserID, sess.SharedID, pduType2FrameAcknowledge, body)
		trace.Logf("engine=rdp frame ack frameID=%d", frameID)
		writeMu.Lock()
		err := transport.WriteX224Data(sess.Conn, buildMCSSendDataRequest(sess.UserID, mcsGlobalChannelID, pdu))
		writeMu.Unlock()
		if err != nil {
			return err
		}
	}
	return nil
}

func (e *SessionEngine) handleEvent(input *InputSender, sink FrameSink, event ClientEvent) error {
	switch event.Type {
	case "mousemove", "mouse_move":
		return input.MouseMove(event.X, event.Y)
	case "mousedown", "mouse_down":
		return input.MouseButton(event.X, event.Y, event.Button, true)
	case "mouseup", "mouse_up":
		return input.MouseButton(event.X, event.Y, event.Button, false)
	case "wheel":
		return input.MouseWheel(event.X, event.Y, event.DeltaY)
	case "key":
		sc, ok := ScanCodeFromJS(event.Code)
		if !ok {
			return nil
		}
		return input.Key(sc, event.Down)
	case "clipboard":
		if event.Text != "" {
			return input.TypeText(event.Text)
		}
		return nil
	case "auto_exec":
		if event.Text == "" {
			return nil
		}
		_ = sink.SendFrame(ServerFrame{Type: "status", Message: "正在远端执行：" + event.Text})
		return input.AutoExec(event.Text)
	case "resize":
		_ = sink.SendFrame(ServerFrame{Type: "status", Message: "分辨率变更请求已记录，重连后生效"})
		return nil
	default:
		return nil
	}
}

// ProbeSessionFrames 用于诊断：建立会话后读取若干秒位图帧数量 // Logan.Li
func ProbeSessionFrames(ctx context.Context, cfg Config, seconds int) (int, error) {
	sess, err := OpenSession(ctx, cfg)
	if err != nil {
		return 0, err
	}
	defer func() {
		if c, ok := sess.Conn.(io.Closer); ok {
			_ = c.Close()
		}
	}()
	receiver := NewFastPathReceiver()
	deadline := time.Now().Add(time.Duration(seconds) * time.Second)
	count := 0
	for time.Now().Before(deadline) {
		if c, ok := sess.Conn.(interface{ SetReadDeadline(time.Time) error }); ok {
			_ = c.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
		}
		pkt, err := ReadWirePacket(sess.Conn)
		if err != nil {
			if time.Now().Before(deadline) {
				continue
			}
			return count, nil
		}
		var rects []BitmapRect
		if pkt.FastPath {
			rects, err = receiver.Handle(pkt.SecFlag, pkt.Payload)
		} else {
			rects, err = receiver.HandleSlowPath(pkt.Payload)
		}
		if err == nil {
			count += len(rects)
		}
	}
	if c, ok := sess.Conn.(interface{ SetReadDeadline(time.Time) error }); ok {
		_ = c.SetReadDeadline(time.Time{})
	}
	return count, nil
}