rdplib/plugin/cliprdr/cliprdr_image_test.go

// cliprdr_image_test.go — 剪贴板图片链路的协议层单元测试（native 可跑）
package cliprdr_test

import (
	"bytes"
	"encoding/binary"
	"image/png"
	"sync"
	"testing"

	"git.zeroonesoft.cn/golib/rdplib/core"
	"git.zeroonesoft.cn/golib/rdplib/plugin/cliprdr"
)

// captureSender 捕获 handler 发出的 PDU，供断言
type captureSender struct {
	mu   sync.Mutex
	pdus [][]byte
}

func (s *captureSender) GetType() (string, uint32) { return "cliprdr", 0 }
func (s *captureSender) Sender(f core.ChannelSender) {}
func (s *captureSender) Process(data []byte)         {}

func (s *captureSender) SendToChannel(name string, data []byte) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	pdu := make([]byte, len(data))
	copy(pdu, data)
	s.pdus = append(s.pdus, pdu)
	return len(data), nil
}

func (s *captureSender) findByType(msgType uint16) [][]byte {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out [][]byte
	for _, p := range s.pdus {
		if len(p) >= 8 && binary.LittleEndian.Uint16(p[0:]) == msgType {
			out = append(out, p)
		}
	}
	return out
}

// fakeDIB 构造一张 2x2 的 32bpp BI_RGB 底行优先 DIB
func fakeDIB() []byte {
	const w, h = 2, 2
	dib := make([]byte, 40+w*4*h)
	binary.LittleEndian.PutUint32(dib[0:4], 40)   // biSize
	binary.LittleEndian.PutUint32(dib[4:8], w)    // width
	binary.LittleEndian.PutUint32(dib[8:12], h)   // height（正 = 底行优先）
	binary.LittleEndian.PutUint16(dib[12:14], 1)  // planes
	binary.LittleEndian.PutUint16(dib[14:16], 32) // bpp
	pix := dib[40:]
	// 视觉顶部 = DIB 最后一行：红、绿
	off := (h - 1) * w * 4
	binary.LittleEndian.PutUint32(pix[off:off+4], 0x00FF0000)   // BGRA → 红
	binary.LittleEndian.PutUint32(pix[off+4:off+8], 0x0000FF00) // 绿
	// 视觉底部 = DIB 首行：蓝、白
	binary.LittleEndian.PutUint32(pix[0:4], 0x000000FF)
	binary.LittleEndian.PutUint32(pix[4:8], 0x00FFFFFF)
	return dib
}

// capsPDU 构造 CB_CLIP_CAPS（声明长格式名），handler 需先收到它才会按长名解析
func capsPDU() []byte {
	body := make([]byte, 0, 20)
	body = binary.LittleEndian.AppendUint16(body, 1) // cCapSets
	body = binary.LittleEndian.AppendUint16(body, 0) // pad
	body = binary.LittleEndian.AppendUint16(body, cliprdr.CB_CAPSTYPE_GENERAL)
	body = binary.LittleEndian.AppendUint16(body, cliprdr.CB_CAPSTYPE_GENERAL_LEN)
	body = binary.LittleEndian.AppendUint32(body, cliprdr.CB_CAPS_VERSION_2)
	body = binary.LittleEndian.AppendUint32(body, cliprdr.CB_USE_LONG_FORMAT_NAMES)
	p := pduHeader(cliprdr.CB_CLIP_CAPS, 0, len(body))
	return append(p, body...)
}

func TestDIBToPNG(t *testing.T) {
	pngBytes, err := cliprdr.DIBToPNG(fakeDIB())
	if err != nil {
		t.Fatalf("DIBToPNG: %v", err)
	}
	img, err := png.Decode(bytes.NewReader(pngBytes))
	if err != nil {
		t.Fatalf("decoded PNG invalid: %v", err)
	}
	if img.Bounds().Dx() != 2 || img.Bounds().Dy() != 2 {
		t.Fatalf("unexpected size %v", img.Bounds())
	}
	r, g, b, _ := img.At(0, 0).RGBA() // RGBA() 分量为 16 位
	if r != 0xFFFF || g != 0 || b != 0 {
		t.Fatalf("pixel(0,0) want red, got r=%x g=%x b=%x", r, g, b)
	}
}

func TestFormatDataResponseDIBConvertedToPNG(t *testing.T) {
	h := cliprdr.NewHandler(nil, nil)
	sender := &captureSender{}
	h.Sender(sender)

	var gotPNG []byte
	h.SetImageCallbacks(func(p []byte) { gotPNG = p }, nil)

	// 先声明能力（长格式名），再让服务端提供仅含 CF_DIB 的 FormatList
	h.Process(capsPDU())
	h.Process(mkFormatListLong([]cliprdr.CliprdrFormat{{FormatId: cliprdr.CF_DIB}}))
	h.Process(responsePDU(cliprdr.CF_DIB, fakeDIB()))

	if len(gotPNG) == 0 {
		t.Fatal("expected PNG bytes via onRemoteClipboardImage")
	}
	if _, err := png.Decode(bytes.NewReader(gotPNG)); err != nil {
		t.Fatalf("converted PNG invalid: %v", err)
	}
}

func TestFormatListPNGTriggersRequest(t *testing.T) {
	h := cliprdr.NewHandler(nil, nil)
	sender := &captureSender{}
	h.Sender(sender)

	h.Process(capsPDU())
	h.Process(mkFormatListLong([]cliprdr.CliprdrFormat{
		{FormatId: 0xC3F0, FormatName: "PNG"},
		{FormatId: cliprdr.CF_UNICODETEXT, FormatName: ""},
	}))

	requests := sender.findByType(0x0004) // CB_FORMAT_DATA_REQUEST
	for _, p := range requests {
		id := binary.LittleEndian.Uint32(p[8:12])
		if id == 0xC3F0 { // 应按服务端给出的 PNG 注册 ID 发起请求
			return
		}
	}
	t.Fatal("expected PNG format data request for server-offered PNG format")
}

func TestFormatDataRequestPNGServedFromProvider(t *testing.T) {
	h := cliprdr.NewHandler(nil, nil)
	sender := &captureSender{}
	h.Sender(sender)
	h.SetImageCallbacks(nil, func() []byte { return []byte("FAKEPNG") })

	h.Process(requestPDU(cliprdr.CF_PNG))

	responses := sender.findByType(0x0005) // CB_FORMAT_DATA_RESPONSE
	for _, p := range responses {
		if binary.LittleEndian.Uint16(p[2:]) == 1 && bytes.Equal(p[8:15], []byte("FAKEPNG")) {
			return
		}
	}
	t.Fatalf("expected PNG data served from provider, got %d responses", len(responses))
}

// --- helpers: 构造标准 CLIPRDR PDU ---

func pduHeader(msgType, msgFlags uint16, bodyLen int) []byte {
	h := make([]byte, 8)
	binary.LittleEndian.PutUint16(h[0:], msgType)
	binary.LittleEndian.PutUint16(h[2:], msgFlags)
	binary.LittleEndian.PutUint32(h[4:], uint32(bodyLen))
	return h
}

func requestPDU(formatId uint32) []byte {
	p := pduHeader(0x0004, 0, 4)
	return binary.LittleEndian.AppendUint32(p, formatId)
}

func responsePDU(formatId uint32, data []byte) []byte {
	_ = formatId // 响应 PDU 本身不含格式 ID，格式由我方请求记录
	p := pduHeader(0x0005, 1, len(data))
	return append(p, data...)
}

func mkFormatListLong(formats []cliprdr.CliprdrFormat) []byte {
	var body bytes.Buffer
	for _, f := range formats {
		binary.Write(&body, binary.LittleEndian, f.FormatId)
		body.Write(encodeUTF16(f.FormatName))
		body.Write([]byte{0, 0})
	}
	p := pduHeader(0x0002, 0, body.Len())
	return append(p, body.Bytes()...)
}

func encodeUTF16(s string) []byte {
	u := []uint16{}
	for _, r := range s {
		u = append(u, uint16(r))
	}
	out := make([]byte, len(u)*2)
	for i, v := range u {
		binary.LittleEndian.PutUint16(out[i*2:], v)
	}
	return out
}