Include encoders

Signed-off-by: Daniel Widerin <daniel@widerin.net>

encoder.go

@@ -0,0 +1,180 @@
+package main
+
+/**
+* XXX: Ugly workaround for https://github.com/amitbet/vnc2video/issues/10. I've copied the file and build a
+* X264ImageCustomEncoder. Once this is merged, we can drop the encoder.go file again.
+*/
+
+import (
+	"errors"
+	"fmt"
+	vnc "github.com/amitbet/vnc2video"
+	"github.com/amitbet/vnc2video/encoders"
+	log "github.com/sirupsen/logrus"
+	"image"
+	"image/color"
+	"io"
+	"os"
+	"os/exec"
+	"strconv"
+)
+
+func encodePPMforRGBA(w io.Writer, img *image.RGBA) error {
+	maxvalue := 255
+	size := img.Bounds()
+	// write ppm header
+	_, err := fmt.Fprintf(w, "P6\n%d %d\n%d\n", size.Dx(), size.Dy(), maxvalue)
+	if err != nil {
+		return err
+	}
+
+	if convImage == nil {
+		convImage = make([]uint8, size.Dy()*size.Dx()*3)
+	}
+
+	rowCount := 0
+	for i := 0; i < len(img.Pix); i++ {
+		if (i % 4) != 3 {
+			convImage[rowCount] = img.Pix[i]
+			rowCount++
+		}
+	}
+
+	if _, err := w.Write(convImage); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func encodePPMGeneric(w io.Writer, img image.Image) error {
+	maxvalue := 255
+	size := img.Bounds()
+	// write ppm header
+	_, err := fmt.Fprintf(w, "P6\n%d %d\n%d\n", size.Dx(), size.Dy(), maxvalue)
+	if err != nil {
+		return err
+	}
+
+	// write the bitmap
+	colModel := color.RGBAModel
+	row := make([]uint8, size.Dx()*3)
+	for y := size.Min.Y; y < size.Max.Y; y++ {
+		i := 0
+		for x := size.Min.X; x < size.Max.X; x++ {
+			color := colModel.Convert(img.At(x, y)).(color.RGBA)
+			row[i] = color.R
+			row[i+1] = color.G
+			row[i+2] = color.B
+			i += 3
+		}
+		if _, err := w.Write(row); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+var convImage []uint8
+
+func encodePPM(w io.Writer, img image.Image) error {
+	if img == nil {
+		return errors.New("nil image")
+	}
+	img1, isRGBImage := img.(*vnc.RGBImage)
+	img2, isRGBA := img.(*image.RGBA)
+	if isRGBImage {
+		return encodePPMforRGBImage(w, img1)
+	} else if isRGBA {
+		return encodePPMforRGBA(w, img2)
+	}
+	return encodePPMGeneric(w, img)
+}
+func encodePPMforRGBImage(w io.Writer, img *vnc.RGBImage) error {
+	maxvalue := 255
+	size := img.Bounds()
+	// write ppm header
+	_, err := fmt.Fprintf(w, "P6\n%d %d\n%d\n", size.Dx(), size.Dy(), maxvalue)
+	if err != nil {
+		return err
+	}
+
+	if _, err := w.Write(img.Pix); err != nil {
+		return err
+	}
+	return nil
+}
+
+type X264ImageCustomEncoder struct {
+	encoders.X264ImageEncoder
+	FFMpegBinPath string
+	cmd           *exec.Cmd
+	input         io.WriteCloser
+	closed        bool
+	Framerate     int
+}
+
+func (enc *X264ImageCustomEncoder) Init(videoFileName string) {
+	if enc.Framerate == 0 {
+		enc.Framerate = 12
+	}
+	cmd := exec.Command(enc.FFMpegBinPath,
+		"-f", "image2pipe",
+		"-vcodec", "ppm",
+		"-r", strconv.Itoa(enc.Framerate),
+		"-an", //no audio
+		"-y",
+		"-i", "-",
+		"-vcodec", "libx264",
+		"-threads", "8",
+		"-preset", "veryfast",
+		"-g", "250",
+		"-crf", "37",
+		videoFileName,
+	)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	encInput, err := cmd.StdinPipe()
+	enc.input = encInput
+	if err != nil {
+		log.Error("can't get ffmpeg input pipe")
+	}
+	enc.cmd = cmd
+}
+func (enc *X264ImageCustomEncoder) Run(videoFileName string) error {
+	if _, err := os.Stat(enc.FFMpegBinPath); os.IsNotExist(err) {
+		if _, err := os.Stat(enc.FFMpegBinPath + ".exe"); os.IsNotExist(err) {
+			log.Error("encoder file doesn't exist in path:", enc.FFMpegBinPath)
+			return errors.New("encoder file doesn't exist in path" + videoFileName)
+		} else {
+			enc.FFMpegBinPath = enc.FFMpegBinPath + ".exe"
+		}
+	}
+
+	enc.Init(videoFileName)
+	log.Infof("launching binary: %v", enc.cmd)
+	err := enc.cmd.Run()
+	if err != nil {
+		log.Errorf("error while launching ffmpeg: %v\n err: %v", enc.cmd.Args, err)
+		return err
+	}
+	return nil
+}
+func (enc *X264ImageCustomEncoder) Encode(img image.Image) {
+	if enc.input == nil || enc.closed {
+		return
+	}
+
+	err := encodePPM(enc.input, img)
+	if err != nil {
+		log.Error("error while encoding image:", err)
+	}
+}
+
+func (enc *X264ImageCustomEncoder) Close() {
+	if enc.closed {
+		return
+	}
+	enc.closed = true
+	enc.input.Close()
+}

main.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	vnc "github.com/amitbet/vnc2video"
-	"github.com/amitbet/vnc2video/encoders"
 	log "github.com/sirupsen/logrus"
 	cli "github.com/urfave/cli/v2"
 	"net"
@@ -141,7 +140,7 @@ func recorder(c *cli.Context) error {
 		panic(err)
 	}
 	log.Infof("Using %s for encoding", ffmpeg_path)
-	vcodec := &encoders.X264ImageEncoder{
+	vcodec := &X264ImageCustomEncoder{
 		FFMpegBinPath: ffmpeg_path,
 		Framerate:     c.Int("framerate"),
 	}
@@ -182,12 +181,16 @@ func recorder(c *cli.Context) error {
 		}
 	}()
 
-	sigc := make(chan os.Signal, 1)
+	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigc,
+	signal.Notify(sigCh,
+		os.Interrupt,
 		syscall.SIGHUP,
 		syscall.SIGINT,
 		syscall.SIGTERM,
-		syscall.SIGQUIT)
+		syscall.SIGQUIT,
+		syscall.SIGKILL,
+	)
+
 	frameBufferReq := 0
 	timeStart := time.Now()
 
@@ -202,22 +205,18 @@ func recorder(c *cli.Context) error {
 				secsPassed := time.Now().Sub(timeStart).Seconds()
 				frameBufferReq++
 				reqPerSec := float64(frameBufferReq) / secsPassed
-				//counter++
-				//jpeg.Encode(out, screenImage, nil)
-				///vcodec.Encode(screenImage)
 				log.Debugf("reqs=%d, seconds=%f, Req Per second= %f", frameBufferReq, secsPassed, reqPerSec)
 
 				reqMsg := vnc.FramebufferUpdateRequest{Inc: 1, X: 0, Y: 0, Width: cc.Width(), Height: cc.Height()}
-				//cc.ResetAllEncodings()
 				reqMsg.Write(cc)
 			}
-		case signal := <-sigc:
+		case signal := <-sigCh:
 			if signal != nil {
 				log.Info(signal, " received, exit.")
 				vcodec.Close()
 				// give some time to write the file
-				time.Sleep(time.Second * 5)
+				time.Sleep(time.Second * 1)
-				return nil
+				os.Exit(0)
 			}
 		}
 	}