natpass/code/client/rule/vnc/process/process.go

package process

import (
	"errors"
	"fmt"
	"image"
	"image/jpeg"
	"net"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/lwch/logging"
	"github.com/lwch/natpass/code/client/rule/vnc/vncnetwork"
	"github.com/lwch/natpass/code/utils"
	"google.golang.org/protobuf/proto"
)

const (
	listenBegin = 6155
	listenEnd   = 6955
)

// Process process
type Process struct {
	pid         int
	srv         *http.Server
	chWrite     chan *vncnetwork.VncMsg
	chImage     chan *vncnetwork.ImageData
	chClipboard chan *vncnetwork.ClipboardData
}

func (p *Process) listenAndServe() (uint16, error) {
	mux := http.NewServeMux()
	mux.HandleFunc("/", p.ws)
	port := uint16(listenBegin)
	for {
		if port > listenEnd {
			return 0, errors.New("no port available")
		}
		p.srv = &http.Server{
			Addr:    fmt.Sprintf("127.0.0.1:%d", port),
			Handler: mux,
		}
		ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
		if err != nil {
			port++
			continue
		}
		go p.srv.Serve(ln)
		return port, nil
	}
}

var upgrader = websocket.Upgrader{EnableCompression: true}

func (p *Process) ws(w http.ResponseWriter, r *http.Request) {
	logging.Info("child process connected")
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer conn.Close()
	defer p.Close()
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer utils.Recover("ws read")
		defer wg.Done()
		for {
			_, data, err := conn.ReadMessage()
			if err != nil {
				logging.Error("read message: %v", err)
				return
			}
			var msg vncnetwork.VncMsg
			err = proto.Unmarshal(data, &msg)
			if err != nil {
				continue
			}
			switch msg.GetXType() {
			case vncnetwork.VncMsg_capture_data:
				p.chImage <- msg.GetData()
			case vncnetwork.VncMsg_clipboard_event:
				p.chClipboard <- msg.GetClipboard()
			default:
			}
		}
	}()
	go func() {
		defer utils.Recover("ws write")
		defer wg.Done()
		for {
			msg := <-p.chWrite
			data, err := proto.Marshal(msg)
			if err != nil {
				continue
			}
			err = conn.WriteMessage(websocket.BinaryMessage, data)
			if err != nil {
				logging.Error("write message: %v", err)
				return
			}
		}
	}()
	wg.Wait()
}

func (p *Process) kill() {
	ps, _ := os.FindProcess(p.pid)
	if ps != nil {
		ps.Kill()
	}
}

// Close close process
func (p *Process) Close() {
	if p.srv != nil {
		p.srv.Close()
	}
	if p.chImage != nil {
		close(p.chImage)
		p.chImage = nil
	}
	if p.chClipboard != nil {
		close(p.chClipboard)
		p.chClipboard = nil
	}
	if p.chWrite != nil {
		close(p.chWrite)
		p.chWrite = nil
	}
	p.kill()
}

// Capture capture desktop image
func (p *Process) Capture(timeout time.Duration) (*image.RGBA, error) {
	var msg vncnetwork.VncMsg
	msg.XType = vncnetwork.VncMsg_capture_req
	p.chWrite <- &msg
	trans := func(data *vncnetwork.ImageData) *image.RGBA {
		img := image.NewRGBA(image.Rect(0, 0, int(data.GetWidth()), int(data.GetHeight())))
		copy(img.Pix, data.GetData())
		// dumpImage(img)
		return img
	}
	if timeout > 0 {
		select {
		case data := <-p.chImage:
			return trans(data), nil
		case <-time.After(timeout):
			return nil, errors.New("timeout")
		}
	} else {
		data := <-p.chImage
		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
	}
}