弃用延时决定截图时机，使用协程检查图像情况

client/rdpClient.go

@@ -39,6 +39,12 @@ type Position struct {
 }
 
 var (
+	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)
 )
@@ -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, ":", "_")
-
+		var bitmapList = make([]*common.Bitmap, 0)
-		//for common.Opened(ctx) {}
 
 		// TODO: 寻找合适的截图时机
-		time.Sleep(5 * time.Second)
+		<-bitmapRevSig
-		modifyMux.Lock()
 
-		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()
 	})
 }
 

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)
 	}