pkg/beam/unix.go

package beam

import (
	"bufio"
	"fmt"
	"net"
	"os"
	"syscall"
)

func debugCheckpoint(msg string, args ...interface{}) {
	if os.Getenv("DEBUG") == "" {
		return
	}
	os.Stdout.Sync()
	tty, _ := os.OpenFile("/dev/tty", os.O_RDWR, 0700)
	fmt.Fprintf(tty, msg, args...)
	bufio.NewScanner(tty).Scan()
	tty.Close()
}

type UnixConn struct {
	*net.UnixConn
	fds []*os.File
}

// Framing:
// In order to handle framing in Send/Recieve, as these give frame
// boundaries we use a very simple 4 bytes header. It is a big endiand
// uint32 where the high bit is set if the message includes a file
// descriptor. The rest of the uint32 is the length of the next frame.
// We need the bit in order to be able to assign recieved fds to
// the right message, as multiple messages may be coalesced into
// a single recieve operation.
func makeHeader(data []byte, fds []int) ([]byte, error) {
	header := make([]byte, 4)

	length := uint32(len(data))

	if length > 0x7fffffff {
		return nil, fmt.Errorf("Data to large")
	}

	if len(fds) != 0 {
		length = length | 0x80000000
	}
	header[0] = byte((length >> 24) & 0xff)
	header[1] = byte((length >> 16) & 0xff)
	header[2] = byte((length >> 8) & 0xff)
	header[3] = byte((length >> 0) & 0xff)

	return header, nil
}

func parseHeader(header []byte) (uint32, bool) {
	length := uint32(header[0])<<24 | uint32(header[1])<<16 | uint32(header[2])<<8 | uint32(header[3])
	hasFd := length&0x80000000 != 0
	length = length & ^uint32(0x80000000)

	return length, hasFd
}

func FileConn(f *os.File) (*UnixConn, error) {
	conn, err := net.FileConn(f)
	if err != nil {
		return nil, err
	}
	uconn, ok := conn.(*net.UnixConn)
	if !ok {
		conn.Close()
		return nil, fmt.Errorf("%d: not a unix connection", f.Fd())
	}
	return &UnixConn{UnixConn: uconn}, nil

}

// Send sends a new message on conn with data and f as payload and
// attachment, respectively.
// On success, f is closed
func (conn *UnixConn) Send(data []byte, f *os.File) error {
	{
		var fd int = -1
		if f != nil {
			fd = int(f.Fd())
		}
		debugCheckpoint("===DEBUG=== about to send '%s'[%d]. Hit enter to confirm: ", data, fd)
	}
	var fds []int
	if f != nil {
		fds = append(fds, int(f.Fd()))
	}
	if err := conn.sendUnix(data, fds...); err != nil {
		return err
	}

	if f != nil {
		f.Close()
	}
	return nil
}

// Receive waits for a new message on conn, and receives its payload
// and attachment, or an error if any.
//
// If more than 1 file descriptor is sent in the message, they are all
// closed except for the first, which is the attachment.
// It is legal for a message to have no attachment or an empty payload.
func (conn *UnixConn) Receive() (rdata []byte, rf *os.File, rerr error) {
	defer func() {
		var fd int = -1
		if rf != nil {
			fd = int(rf.Fd())
		}
		debugCheckpoint("===DEBUG=== Receive() -> '%s'[%d]. Hit enter to continue.\n", rdata, fd)
	}()

	// Read header
	header := make([]byte, 4)
	nRead := uint32(0)

	for nRead < 4 {
		n, err := conn.receiveUnix(header[nRead:])
		if err != nil {
			return nil, nil, err
		}
		nRead = nRead + uint32(n)
	}

	length, hasFd := parseHeader(header)

	if hasFd {
		if len(conn.fds) == 0 {
			return nil, nil, fmt.Errorf("No expected file descriptor in message")
		}

		rf = conn.fds[0]
		conn.fds = conn.fds[1:]
	}

	rdata = make([]byte, length)

	nRead = 0
	for nRead < length {
		n, err := conn.receiveUnix(rdata[nRead:])
		if err != nil {
			return nil, nil, err
		}
		nRead = nRead + uint32(n)
	}

	return
}

func (conn *UnixConn) receiveUnix(buf []byte) (int, error) {
	oob := make([]byte, syscall.CmsgSpace(4))
	bufn, oobn, _, _, err := conn.ReadMsgUnix(buf, oob)
	if err != nil {
		return 0, err
	}
	fd := extractFd(oob[:oobn])
	if fd != -1 {
		f := os.NewFile(uintptr(fd), "")
		conn.fds = append(conn.fds, f)
	}

	return bufn, nil
}

func (conn *UnixConn) sendUnix(data []byte, fds ...int) error {
	header, err := makeHeader(data, fds)
	if err != nil {
		return err
	}

	// There is a bug in conn.WriteMsgUnix where it doesn't correctly return
	// the number of bytes writte (http://code.google.com/p/go/issues/detail?id=7645)
	// So, we can't rely on the return value from it. However, we must use it to
	// send the fds. In order to handle this we only write one byte using WriteMsgUnix
	// (when we have to), as that can only ever block or fully suceed. We then write
	// the rest with conn.Write()
	// The reader side should not rely on this though, as hopefully this gets fixed
	// in go later.
	written := 0
	if len(fds) != 0 {
		oob := syscall.UnixRights(fds...)
		wrote, _, err := conn.WriteMsgUnix(header[0:1], oob, nil)
		if err != nil {
			return err
		}
		written = written + wrote
	}

	for written < len(header) {
		wrote, err := conn.Write(header[written:])
		if err != nil {
			return err
		}
		written = written + wrote
	}

	written = 0
	for written < len(data) {
		wrote, err := conn.Write(data[written:])
		if err != nil {
			return err
		}
		written = written + wrote
	}

	return nil
}

func extractFd(oob []byte) int {
	// Grab forklock to make sure no forks accidentally inherit the new
	// fds before they are made CLOEXEC
	// There is a slight race condition between ReadMsgUnix returns and
	// when we grap the lock, so this is not perfect. Unfortunately
	// There is no way to pass MSG_CMSG_CLOEXEC to recvmsg() nor any
	// way to implement non-blocking i/o in go, so this is hard to fix.
	syscall.ForkLock.Lock()
	defer syscall.ForkLock.Unlock()
	scms, err := syscall.ParseSocketControlMessage(oob)
	if err != nil {
		return -1
	}

	foundFd := -1
	for _, scm := range scms {
		fds, err := syscall.ParseUnixRights(&scm)
		if err != nil {
			continue
		}

		for _, fd := range fds {
			if foundFd == -1 {
				syscall.CloseOnExec(fd)
				foundFd = fd
			} else {
				syscall.Close(fd)
			}
		}
	}

	return foundFd
}

func socketpair() ([2]int, error) {
	return syscall.Socketpair(syscall.AF_LOCAL, syscall.SOCK_STREAM|syscall.FD_CLOEXEC, 0)
}

// SocketPair is a convenience wrapper around the socketpair(2) syscall.
// It returns a unix socket of type SOCK_STREAM in the form of 2 file descriptors
// not bound to the underlying filesystem.
// Messages sent on one end are received on the other, and vice-versa.
// It is the caller's responsibility to close both ends.
func SocketPair() (a *os.File, b *os.File, err error) {
	defer func() {
		var (
			fdA int = -1
			fdB int = -1
		)
		if a != nil {
			fdA = int(a.Fd())
		}
		if b != nil {
			fdB = int(b.Fd())
		}
		debugCheckpoint("===DEBUG=== SocketPair() = [%d-%d]. Hit enter to confirm: ", fdA, fdB)
	}()
	pair, err := socketpair()
	if err != nil {
		return nil, nil, err
	}
	return os.NewFile(uintptr(pair[0]), ""), os.NewFile(uintptr(pair[1]), ""), nil
}

func USocketPair() (*UnixConn, *UnixConn, error) {
	debugCheckpoint("===DEBUG=== USocketPair(). Hit enter to confirm: ")
	defer debugCheckpoint("===DEBUG=== USocketPair() returned. Hit enter to confirm ")
	a, b, err := SocketPair()
	if err != nil {
		return nil, nil, err
	}
	defer a.Close()
	defer b.Close()
	uA, err := FileConn(a)
	if err != nil {
		return nil, nil, err
	}
	uB, err := FileConn(b)
	if err != nil {
		uA.Close()
		return nil, nil, err
	}
	return uA, uB, nil
}

// FdConn wraps a file descriptor in a standard *net.UnixConn object, or
// returns an error if the file descriptor does not point to a unix socket.
// This creates a duplicate file descriptor. It's the caller's responsibility
// to close both.
func FdConn(fd int) (n *net.UnixConn, err error) {
	{
		debugCheckpoint("===DEBUG=== FdConn([%d]) = (unknown fd). Hit enter to confirm: ", fd)
	}
	f := os.NewFile(uintptr(fd), fmt.Sprintf("%d", fd))
	conn, err := net.FileConn(f)
	if err != nil {
		return nil, err
	}
	uconn, ok := conn.(*net.UnixConn)
	if !ok {
		conn.Close()
		return nil, fmt.Errorf("%d: not a unix connection", fd)
	}
	return uconn, nil
}