capture.go

capture.go

package capture

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"syscall"
	"testing"
)

type Capture struct {
	done    chan struct{}
	fd      int
	oldfile *os.File
	io.Reader
}

func (c *Capture) Close() error {
	// fd -> closefd
	closefd, err := syscall.Dup(c.fd)
	if err != nil {
		return err
	}

	// oldfile.fd -> fd
	err = syscall.Dup2(int(c.oldfile.Fd()), c.fd)
	if err != nil {
		return err
	}

	err = syscall.Close(closefd)
	if err != nil {
		return err
	}

	<-c.done

	return nil
}

func NewCapture(fd int) (*Capture, error) {
	c := &Capture{
		done: make(chan struct{}),
		fd:   fd,
	}

	// fd -> oldfd
	oldfd, err := syscall.Dup(c.fd)
	if err != nil {
		return nil, err
	}

	c.oldfile = os.NewFile(uintptr(oldfd), "")

	ospr, ospw, err := os.Pipe()
	if err != nil {
		c.oldfile.Close()
		return nil, err
	}
	defer ospw.Close()

	// ospw -> fd
	err = syscall.Dup2(int(ospw.Fd()), c.fd)
	if err != nil {
		c.oldfile.Close()
		ospr.Close()
		return nil, err
	}

	iopr, iopw := io.Pipe()

	go func() {
		// copy from ospr -> c.oldfile and iopw (read by c.Reader)
		_, err = io.Copy(io.MultiWriter(c.oldfile, iopw), ospr)
		iopw.CloseWithError(err)
		close(c.done)
	}()

	c.Reader = iopr

	return c, nil
}

func TestCapture(t *testing.T) {
	c, err := NewCapture(1)
	if err != nil {
		t.Fatal(err)
	}

	done := make(chan struct{})
	go func() {
		br := bufio.NewReader(c)
		for {
			s, err := br.ReadString('\n')
			if err != nil {
				break
			}
			fmt.Fprint(os.Stderr, s) // TODO: write to insights
		}
		close(done)
	}()

	fmt.Println("Hello!")

	err = c.Close()
	if err != nil {
		t.Fatal(err)
	}

	<-done

	fmt.Println("Message 2")
}