From d8b3c8734790b29bc8edf315264a715c26215eff Mon Sep 17 00:00:00 2001
From: yv1ing <me@yvling.cn>
Date: Sat, 22 Mar 2025 02:58:25 +0800
Subject: [PATCH] =?UTF-8?q?=E5=BC=83=E7=94=A8=E5=BB=B6=E6=97=B6=E5=86=B3?=
 =?UTF-8?q?=E5=AE=9A=E6=88=AA=E5=9B=BE=E6=97=B6=E6=9C=BA=EF=BC=8C=E4=BD=BF?=
 =?UTF-8?q?=E7=94=A8=E5=8D=8F=E7=A8=8B=E6=A3=80=E6=9F=A5=E5=9B=BE=E5=83=8F?=
 =?UTF-8?q?=E6=83=85=E5=86=B5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 client/rdpClient.go      | 55 ++++++++++++++++++++++++++++++++--------
 client/rdpClient_test.go |  2 +-
 2 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/client/rdpClient.go b/client/rdpClient.go
index 7ea3d67..a22d99e 100644
--- a/client/rdpClient.go
+++ b/client/rdpClient.go
@@ -39,8 +39,14 @@ type Position struct {
 }
 
 var (
-	modifyMux   = sync.Mutex{}
-	positionMap = make(map[Position]*common.Bitmap)
+	timerCycle     = 500 * time.Millisecond
+	timer          = time.NewTimer(timerCycle)
+	bitmapRevOne   = sync.Once{}
+	bitmapRevList  = make([]pdu.BitmapData, 0)
+	bitmapRevCount = 0
+	bitmapRevSig   = make(chan struct{})
+	modifyMux      = sync.Mutex{}
+	positionMap    = make(map[Position]*common.Bitmap)
 )
 
 func Newclient(host string, logLevel glog.LEVEL) *RdpClient {
@@ -69,6 +75,21 @@ func (client *RdpClient) Login(domain, username, password string) ([]byte, error
 
 	// 设置事件处理
 	client.setEventHandler()
+	go func() {
+		for {
+			select {
+			case <-timer.C:
+				modifyMux.Lock()
+				// 检查是否有新元素添加
+				if len(bitmapRevList) == bitmapRevCount {
+					bitmapRevOne.Do(func() {
+						close(bitmapRevSig)
+					})
+				}
+				modifyMux.Unlock()
+			}
+		}
+	}()
 
 	// 合并登录界面位图
 	wg := &sync.WaitGroup{}
@@ -78,22 +99,19 @@ func (client *RdpClient) Login(domain, username, password string) ([]byte, error
 	surfaceflingerCtx, _ := context.WithCancel(context.Background())
 	go func(ctx context.Context, c *RdpClient) {
 		outputName := strings.ReplaceAll(client.Host, ":", "_")
-
-		//for common.Opened(ctx) {}
+		var bitmapList = make([]*common.Bitmap, 0)
 
 		// TODO: 寻找合适的截图时机
-		time.Sleep(5 * time.Second)
-		modifyMux.Lock()
+		<-bitmapRevSig
 
-		var bitmapList = make([]*common.Bitmap, 0)
+		modifyMux.Lock()
 		for _, bitmap := range positionMap {
 			//common.DrawOneBitmap(bitmap, fmt.Sprintf("%d-%d-%d-%d.png", pos.Left, pos.Top, pos.Right, pos.Bottom))
 			bitmapList = append(bitmapList, bitmap)
 		}
-
-		buff = common.DrawFullImage(outputName, bitmapList)
 		modifyMux.Unlock()
 
+		buff = common.DrawFullImage(outputName, bitmapList)
 		wg.Done()
 	}(surfaceflingerCtx, client)
 
@@ -112,7 +130,7 @@ func (client *RdpClient) initProtocolStack(conn net.Conn, domain, username, pass
 
 	// 配置桌面信息
 	client.channels = plugin.NewChannels(client.sec)
-	client.mcs.SetClientDesktop(uint16(1920), uint16(1080))
+	client.mcs.SetClientDesktop(uint16(800), uint16(600))
 
 	// 设置认证信息
 	client.sec.SetDomain(domain)
@@ -150,6 +168,22 @@ func (client *RdpClient) setEventHandler() {
 
 	client.pdu.On("bitmap", func(rectangles []pdu.BitmapData) {
 		modifyMux.Lock()
+		defer modifyMux.Unlock()
+
+		// 重置定时器
+		if !timer.Stop() {
+			<-timer.C
+		}
+		timer.Reset(timerCycle)
+
+		oldCount := len(bitmapRevList)
+		bitmapRevList = append(bitmapRevList, rectangles...)
+		newCount := len(bitmapRevList)
+
+		if newCount > oldCount {
+			bitmapRevCount = newCount
+		}
+
 		for _, rectangle := range rectangles {
 			pos := Position{Left: int(rectangle.DestLeft / rectangle.Width), Top: int(rectangle.DestTop / rectangle.Height), Right: int(rectangle.DestRight / rectangle.Width), Bottom: int(rectangle.DestBottom / rectangle.Height)}
 			isCompress := rectangle.IsCompress()
@@ -172,7 +206,6 @@ func (client *RdpClient) setEventHandler() {
 			}
 			positionMap[pos] = bitmap
 		}
-		modifyMux.Unlock()
 	})
 }
 
diff --git a/client/rdpClient_test.go b/client/rdpClient_test.go
index dba74a6..1212d4c 100644
--- a/client/rdpClient_test.go
+++ b/client/rdpClient_test.go
@@ -13,7 +13,7 @@ func TestGetScreen(t *testing.T) {
 	Windows 10 ✔
 	*/
 
-	imgBytes, err := getScreen("192.168.250.128:3389")
+	imgBytes, err := RealGetScreen("127.0.0.1:3389")
 	if err != nil {
 		t.Fatal(err)
 	}