diff --git a/code/client/main.go b/code/client/main.go index cfb2c92..c99c3e3 100644 --- a/code/client/main.go +++ b/code/client/main.go @@ -146,6 +146,11 @@ func main() { os.Exit(1) } + // for test + // work := worker.NewWorker() + // work.TestCapture() + // return + dir, err := filepath.Abs(*conf) runtime.Assert(err) diff --git a/code/client/pool/send_vnc.go b/code/client/pool/send_vnc.go index a6dcd12..5cf715d 100644 --- a/code/client/pool/send_vnc.go +++ b/code/client/pool/send_vnc.go @@ -9,6 +9,11 @@ import ( // SendVNCImage send vnc image data func (conn *Conn) SendVNCImage(to string, toIdx uint32, id string, screen, rect image.Rectangle, encode network.VncImageEncoding, data []byte) { + dup := func(data []byte) []byte { + ret := make([]byte, len(data)) + copy(ret, data) + return ret + } var msg network.Msg msg.To = to msg.ToIdx = toIdx @@ -25,7 +30,7 @@ func (conn *Conn) SendVNCImage(to string, toIdx uint32, id string, screen, rect RectHeight: uint32(rect.Dy()), }, Encode: encode, - Data: data, + Data: dup(data), }, } select { diff --git a/code/client/tunnel/vnc/process/process_windows.go b/code/client/tunnel/vnc/process/process_windows.go index 699da82..ba584df 100644 --- a/code/client/tunnel/vnc/process/process_windows.go +++ b/code/client/tunnel/vnc/process/process_windows.go @@ -4,6 +4,7 @@ import ( "errors" "fmt" "image" + "image/jpeg" "natpass/code/client/tunnel/vnc/define" "natpass/code/client/tunnel/vnc/vncnetwork" "os" @@ -13,6 +14,7 @@ import ( "unicode/utf16" "unsafe" + "github.com/lwch/logging" "golang.org/x/sys/windows" ) @@ -137,14 +139,8 @@ func (p *Process) Capture(timeout time.Duration) (*image.RGBA, error) { p.chWrite <- &msg trans := func(data *vncnetwork.ImageData) *image.RGBA { img := image.NewRGBA(image.Rect(0, 0, int(data.GetWidth()), int(data.GetHeight()))) - switch data.GetBits() { - case 8: - // TODO - case 24: - // TODO - case 32: - copy(img.Pix, data.GetData()) - } + copy(img.Pix, data.GetData()) + // dumpImage(img) return img } if timeout > 0 { @@ -159,3 +155,17 @@ func (p *Process) Capture(timeout time.Duration) (*image.RGBA, error) { return trans(data), nil } } + +func dumpImage(img image.Image) { + f, err := os.Create(`C:\Users\lwch\Pictures\debug.jpeg`) + if err != nil { + logging.Error("debug: %v", err) + return + } + defer f.Close() + err = jpeg.Encode(f, img, nil) + if err != nil { + logging.Error("encode: %v", err) + return + } +} diff --git a/code/client/tunnel/vnc/worker/capture.go b/code/client/tunnel/vnc/worker/capture.go deleted file mode 100644 index d2d2c4f..0000000 --- a/code/client/tunnel/vnc/worker/capture.go +++ /dev/null @@ -1,37 +0,0 @@ -package worker - -import ( - "natpass/code/client/tunnel/vnc/vncnetwork" - - "github.com/gorilla/websocket" - "google.golang.org/protobuf/proto" -) - -func captureError(conn *websocket.Conn, msg string) { - var m vncnetwork.VncMsg - m.XType = vncnetwork.VncMsg_capture_data - m.Payload = &vncnetwork.VncMsg_Data{ - Data: &vncnetwork.ImageData{ - Ok: false, - Msg: msg, - }, - } - data, _ := proto.Marshal(&m) - conn.WriteMessage(websocket.BinaryMessage, data) -} - -func captureOK(conn *websocket.Conn, bits, width, height int, data []byte) { - var msg vncnetwork.VncMsg - msg.XType = vncnetwork.VncMsg_capture_data - msg.Payload = &vncnetwork.VncMsg_Data{ - Data: &vncnetwork.ImageData{ - Ok: true, - Bits: uint32(bits), - Width: uint32(width), - Height: uint32(height), - Data: data, - }, - } - data, _ = proto.Marshal(&msg) - conn.WriteMessage(websocket.BinaryMessage, data) -} diff --git a/code/client/tunnel/vnc/worker/capture_windows.go b/code/client/tunnel/vnc/worker/capture_windows.go index 1c025e7..3865e7f 100644 --- a/code/client/tunnel/vnc/worker/capture_windows.go +++ b/code/client/tunnel/vnc/worker/capture_windows.go @@ -3,25 +3,36 @@ package worker import ( "errors" "natpass/code/client/tunnel/vnc/define" + "natpass/code/client/tunnel/vnc/vncnetwork" "syscall" "unsafe" - "github.com/gorilla/websocket" "github.com/lwch/logging" ) -func (worker *Worker) runCapture(conn *websocket.Conn) { +func (worker *Worker) runCapture() vncnetwork.ImageData { err := worker.capture() if err != nil { - logging.Error("capture: %v", err.Error()) - captureError(conn, err.Error()) - return + return vncnetwork.ImageData{ + Ok: false, + Msg: err.Error(), + } } data := make([]byte, worker.info.width*worker.info.height*worker.info.bits/8) for i := 0; i < len(data); i++ { data[i] = *(*uint8)(unsafe.Pointer(worker.buffer + uintptr(i))) } - captureOK(conn, worker.info.bits, worker.info.width, worker.info.height, data) + // BGR => RGB + for i := 0; i < len(data); i += (worker.info.bits / 8) { + data[i], data[i+2] = data[i+2], data[i] + } + return vncnetwork.ImageData{ + Ok: true, + Bits: uint32(worker.info.bits), + Width: uint32(worker.info.width), + Height: uint32(worker.info.height), + Data: data, + } } func (worker *Worker) capture() error { @@ -43,6 +54,7 @@ func (worker *Worker) capture() error { return errors.New("update buffer: " + err.Error()) } } + logging.Info("width=%d, height=%d, bits=%d", info.width, info.height, info.bits) memDC, _, err := syscall.Syscall(define.FuncCreateCompatibleDC, 1, worker.hdc, 0, 0) if memDC == 0 { return errors.New("create dc: " + err.Error()) @@ -59,7 +71,7 @@ func (worker *Worker) capture() error { return errors.New("select object: " + err.Error()) } defer syscall.Syscall(define.FuncSelectObject, 2, memDC, oldDC, 0) - ok, _, err := syscall.Syscall9(define.FuncBitBlt, 0, memDC, 0, 0, + ok, _, err := syscall.Syscall9(define.FuncBitBlt, 9, memDC, 0, 0, uintptr(worker.info.width), uintptr(worker.info.height), worker.hdc, 0, 0, define.SRCCOPY) if ok == 0 { return errors.New("bitblt: " + err.Error()) @@ -92,6 +104,9 @@ func (worker *Worker) copyImageData(bitmap uintptr) { hdr.BiHeight = int32(-worker.info.height) hdr.BiCompression = define.BI_RGB hdr.BiSizeImage = 0 - syscall.Syscall9(define.FuncGetDIBits, 7, worker.hdc, bitmap, 0, uintptr(worker.info.height), + lines, _, err := syscall.Syscall9(define.FuncGetDIBits, 7, worker.hdc, bitmap, 0, uintptr(worker.info.height), worker.buffer, uintptr(unsafe.Pointer(&hdr)), define.DIB_RGB_COLORS, 0, 0) + if lines == 0 { + logging.Error("get bits: %v", err) + } } diff --git a/code/client/tunnel/vnc/worker/capture_xx.go b/code/client/tunnel/vnc/worker/capture_xx.go index 09c604d..bc817ee 100644 --- a/code/client/tunnel/vnc/worker/capture_xx.go +++ b/code/client/tunnel/vnc/worker/capture_xx.go @@ -5,5 +5,6 @@ package worker import "github.com/gorilla/websocket" -func (worker *Worker) runCapture(conn *websocket.Conn) { +func (worker *Worker) runCapture() vncnetwork.ImageData { + return vncnetwork.ImageData{} } diff --git a/code/client/tunnel/vnc/worker/worker.go b/code/client/tunnel/vnc/worker/worker.go index c277ce0..74527cc 100644 --- a/code/client/tunnel/vnc/worker/worker.go +++ b/code/client/tunnel/vnc/worker/worker.go @@ -1,7 +1,10 @@ package worker import ( + "image" + "image/jpeg" "natpass/code/client/tunnel/vnc/vncnetwork" + "os" "github.com/gorilla/websocket" "github.com/lwch/logging" @@ -42,7 +45,40 @@ func (worker *Worker) Do(conn *websocket.Conn) { } switch msg.GetXType() { case vncnetwork.VncMsg_capture_req: - worker.runCapture(conn) + data := worker.runCapture() + // if data.Ok { + // dumpImage(data.Data, int(data.Width), int(data.Height)) + // } + if !data.Ok { + logging.Error("capture: %s", data.Msg) + } + msg.XType = vncnetwork.VncMsg_capture_data + msg.Payload = &vncnetwork.VncMsg_Data{ + Data: &data, + } + enc, _ := proto.Marshal(&msg) + conn.WriteMessage(websocket.BinaryMessage, enc) } } } + +func (worker *Worker) TestCapture() { + msg := worker.runCapture() + dumpImage(msg.Data, int(msg.Width), int(msg.Height)) +} + +func dumpImage(data []byte, width, height int) { + f, err := os.Create(`C:\Users\lwch\Pictures\debug.jpeg`) + if err != nil { + logging.Error("debug: %v", err) + return + } + defer f.Close() + img := image.NewRGBA(image.Rect(0, 0, width, height)) + copy(img.Pix, data) + err = jpeg.Encode(f, img, nil) + if err != nil { + logging.Error("encode: %v", err) + return + } +} diff --git a/code/client/tunnel/vnc/worker/worker_windows.go b/code/client/tunnel/vnc/worker/worker_windows.go index f89e937..39d4aa5 100644 --- a/code/client/tunnel/vnc/worker/worker_windows.go +++ b/code/client/tunnel/vnc/worker/worker_windows.go @@ -92,6 +92,10 @@ func (worker *Worker) updateInfo() error { worker.info.bits = int(bits) worker.info.width = int(width) worker.info.height = int(height) + if worker.info.bits != 32 { + worker.info.bits = 32 + logging.Info("reset bits to 32") + } return nil }