aiops/backend/pkg/rdp/session_connect.go

package rdp

import (
	"context"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"ezops/pkg/rdp/credssp"
	"ezops/pkg/rdp/transport"
)

// Session 保存已建立的 RDP 图形会话状态 // Logan.Li
type Session struct {
	Conn     io.ReadWriter
	UserID   uint16
	SharedID uint32
	Width    int
	Height   int
	// ChannelNameByID 记录服务端分配的静态虚拟通道 ID→名称(cliprdr 等)，
	// 供 readLoop 按通道分发数据给对应 Handler。空表示服务端未授权任何虚拟通道。// Logan.Li
	ChannelNameByID map[uint16]string
	// PendingChannelPDUs 缓存握手收尾阶段提前到达的虚拟通道包，避免 cliprdr 初始化被吞。// Logan.Li
	PendingChannelPDUs []mcsDataIndication
}

// OpenSession 完成 CredSSP + MCS/GCC 握手并进入可收帧状态 // Logan.Li
func OpenSession(ctx context.Context, cfg Config) (*Session, error) {
	port := cfg.Port
	if port == 0 {
		port = 3389
	}
	if cfg.MCSMode == "" {
		cfg.MCSMode = strings.TrimSpace(os.Getenv("RDP_MCS_MODE"))
	}
	if cfg.MCSMode == "" {
		cfg.MCSMode = "modern"
	}
	address := fmt.Sprintf("%s:%d", cfg.Host, port)
	strategies := []credssp.Options{
		{Version: 6, PubKeyStyle: "binding"},
		{Version: 6, PubKeyStyle: "binding", UseCBT: true},
		{Version: 6, PubKeyStyle: "binding", ForceMIC: true},
		{Version: 2, PubKeyStyle: "raw"},
	}
	var lastErr error
	for _, opts := range strategies {
		sess, err := openSessionOnce(ctx, cfg, address, opts)
		if err == nil {
			return sess, nil
		}
		lastErr = err
	}
	if lastErr == nil {
		lastErr = fmt.Errorf("RDP 会话建立失败")
	}
	return nil, lastErr
}

func openSessionOnce(ctx context.Context, cfg Config, address string, opts credssp.Options) (*Session, error) {
	tlsConn, _, err := dialCredSSPOnce(ctx, cfg, address, opts)
	if err != nil {
		return nil, err
	}
	_ = tlsConn.SetDeadline(time.Now().Add(45 * time.Second))
	if err := transport.WriteX224Data(tlsConn, buildMCSConnectInitial(cfg)); err != nil {
		_ = tlsConn.Close()
		return nil, err
	}
	payload, err := transport.ReadX224Data(tlsConn)
	if err != nil {
		_ = tlsConn.Close()
		return nil, err
	}
	if _, err := parseMCSConnectResponse(payload); err != nil {
		_ = tlsConn.Close()
		return nil, err
	}
	userID, err := attachMCSUser(tlsConn)
	if err != nil {
		_ = tlsConn.Close()
		return nil, err
	}
	mcsInfo, _ := parseMCSConnectResponse(payload)
	joinIDs := append([]uint16{mcsGlobalChannelID, userID}, mcsInfo.ServerChannelIDs...)
	if mcsInfo.MessageChannelID != 0 {
		joinIDs = append(joinIDs, mcsInfo.MessageChannelID)
	}
	for _, channelID := range joinIDs {
		if _, err := joinMCSChannel(tlsConn, userID, channelID); err != nil {
			_ = tlsConn.Close()
			return nil, err
		}
	}
	_, demand, err := sendClientInfoAndReadDemand(tlsConn, cfg, userID)
	if err != nil {
		_ = tlsConn.Close()
		return nil, err
	}
	_, pending, err := sendConfirmActiveAndFinalize(tlsConn, cfg, userID, demand)
	if err != nil {
		_ = tlsConn.Close()
		return nil, err
	}
	morePending, err := requestDisplayUpdates(tlsConn, cfg, userID, demand.sharedID)
	if err != nil {
		_ = tlsConn.Close()
		return nil, err
	}
	pending = append(pending, morePending...)
	width := cfg.Width
	if width <= 0 {
		width = 1280
	}
	height := cfg.Height
	if height <= 0 {
		height = 800
	}
	_ = tlsConn.SetDeadline(time.Time{})
	return &Session{
		Conn:               tlsConn,
		UserID:             userID,
		SharedID:           demand.sharedID,
		Width:              width,
		Height:             height,
		ChannelNameByID:    mcsInfo.ChannelNameByID,
		PendingChannelPDUs: pending,
	}, nil
}

func requestDisplayUpdates(rw io.ReadWriter, cfg Config, userID uint16, sharedID uint32) ([]mcsDataIndication, error) {
	var pending []mcsDataIndication
	if conn, ok := rw.(interface{ SetReadDeadline(time.Time) error }); ok {
		_ = conn.SetReadDeadline(time.Now().Add(3 * time.Second))
		pkt, err := ReadWirePacket(rw)
		_ = conn.SetReadDeadline(time.Time{})
		if err == nil && !pkt.FastPath {
			if ind, perr := parseMCSDataIndication(pkt.Payload); perr == nil && ind.channelID != mcsGlobalChannelID {
				pending = append(pending, ind)
			}
		}
	}
	body := buildSuppressOutputBody(cfg, true)
	pdu := buildDataPDU(userID, sharedID, pduType2SuppressOutput, body)
	if err := transport.WriteX224Data(rw, buildMCSSendDataRequest(userID, mcsGlobalChannelID, pdu)); err != nil {
		return pending, fmt.Errorf("发送 SuppressOutput 失败: %w", err)
	}
	return pending, nil
}

func buildSuppressOutputBody(cfg Config, allow bool) []byte {
	width := cfg.Width
	if width <= 0 {
		width = 1280
	}
	height := cfg.Height
	if height <= 0 {
		height = 800
	}
	out := make([]byte, 0, 12)
	if allow {
		out = append(out, 1, 0, 0, 0)
	} else {
		out = append(out, 0, 0, 0, 0)
	}
	out = append(out, 0, 0, 0, 0, byte(width-1), byte((width-1)>>8), byte(height-1), byte((height-1)>>8))
	return out
}

func buildRefreshRectBody(cfg Config) []byte {
	width := cfg.Width
	if width <= 0 {
		width = 1280
	}
	height := cfg.Height
	if height <= 0 {
		height = 800
	}
	out := make([]byte, 0, 10)
	out = append(out, 1, 0)
	out = append(out, 0, 0, 0, 0, byte(width-1), byte((width-1)>>8), byte(height-1), byte((height-1)>>8))
	return out
}