package main

import (
	"fmt"
	"flag"
	"net"
	"os"
	"sync"
	"syscall"
	"strconv"
	"math/rand"
	"time"
)

func unixPair(typ int) (*net.UnixConn, *net.UnixConn) {
	fds, err := syscall.Socketpair(syscall.AF_UNIX, typ, 0)
	if err != nil {
		panic(os.NewSyscallError("socketpair", err))
	}
	file := os.NewFile(uintptr(fds[0]), "")
	c0, err := net.FileConn(file)
	if err != nil {
		panic(err)
	}
	if err := file.Close(); err != nil {
		panic(err)
	}
	file = os.NewFile(uintptr(fds[1]), "")
	c1, err := net.FileConn(file)
	if err != nil {
		panic(err)
	}
	if err := file.Close(); err != nil {
		panic(err)
	}
	return c0.(*net.UnixConn), c1.(*net.UnixConn)
}

func UnixPair() (*net.UnixConn, *net.UnixConn) {
	return unixPair(syscall.SOCK_STREAM)
}

func UnixgramPair() (*net.UnixConn, *net.UnixConn) {
	return unixPair(syscall.SOCK_DGRAM)
}

func UnixpacketPair() (*net.UnixConn, *net.UnixConn) {
	return unixPair(syscall.SOCK_SEQPACKET)
}

func UnixConnectedPair() (*net.UnixConn, *net.UnixConn) {
	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
	addr := "@foobarbaz-" + strconv.Itoa(rng.Int())
	l, err := net.Listen("unix", addr)
	if err != nil {
		panic(err)
	}
	ch := make(chan net.Conn)
	go func() {
		c1, err := l.Accept()
		if err != nil {
			panic(err)
		}
		ch <- c1
	}()
	c2, err := net.Dial("unix", addr)
	if err != nil {
		panic(err)
	}
	c1 := <-ch
	return c1.(*net.UnixConn), c2.(*net.UnixConn)
}

func sameProcClose(wg *sync.WaitGroup, c net.Conn, closeWrite bool) {
	defer wg.Done()
	fmt.Printf("sameProcClose closeWrite=%v\n", closeWrite)
	fmt.Printf("closing\n")
	if closeWrite {
		c.(*net.UnixConn).CloseWrite()
	}
	c.Close()
	fmt.Printf("closed\n")
}

func childProcClose(wg *sync.WaitGroup, c net.Conn, flag string) {
	defer wg.Done()
	fmt.Printf("childProcClose flag=%v\n", flag)
	argv := []string{os.Args[0], "-child=" + flag}
	file, err := c.(*net.UnixConn).File()
	if err != nil {
		panic(err)
	}
	attr := os.ProcAttr{Files: []*os.File{nil, nil, os.Stderr, file}}
	attr.Sys = &syscall.SysProcAttr{Pdeathsig: syscall.SIGKILL}
	p, err := os.StartProcess(os.Args[0], argv, &attr)
	if err != nil {
		panic(err)
	}
	file.Close()
	c.Close()
	p.Wait()
}

func childClose() {
	file := os.NewFile(3, "")
	conn, err := net.FileConn(file)
	if err != nil {
		panic(err)
	}
	if err := file.Close(); err != nil {
		panic(err)
	}
	fmt.Printf("closing\n")
	conn.(*net.UnixConn).Close()
	fmt.Printf("closed\n")
}

func childCloseWrite() {
	file := os.NewFile(3, "")
	conn, err := net.FileConn(file)
	if err != nil {
		panic(err)
	}
	if err := file.Close(); err != nil {
		panic(err)
	}
	fmt.Printf("closing with closeWrite\n")
	conn.(*net.UnixConn).CloseWrite()
	fmt.Printf("closed\n")
}

type newPairFunc func() (*net.UnixConn, *net.UnixConn)
type closerFunc func(*sync.WaitGroup, net.Conn) ()

func test(newPair newPairFunc, closer closerFunc) {
	c1, c2 := newPair()
	var wg sync.WaitGroup
	wg.Add(2)
	go closer(&wg, c2)
	go func() {
		defer wg.Done()
		buf := make([]byte, 1024)
		for {
			fmt.Printf("reading\n")
			if _, err := c1.Read(buf); err != nil {
				fmt.Printf("read err=%v\n", err)
				return
			}
		}
	}()
	wg.Wait()
}

func main() {
	var child string
	flag.StringVar(&child, "child", "", "panic|close|closewrite")
	flag.Parse()
	if child != "" {
		switch child {
		case "panic":
			panic("boom")
		case "close":
			childClose()
		case "closewrite":
			childCloseWrite()
		default:
			panic("invalid child value: " + child)
		}
		os.Exit(0)
	}

	test(UnixPair, func(wg *sync.WaitGroup, c net.Conn) { sameProcClose(wg, c, false) })
	test(UnixPair, func(wg *sync.WaitGroup, c net.Conn) { sameProcClose(wg, c, true) })
	test(UnixPair, func(wg *sync.WaitGroup, c net.Conn) { childProcClose(wg, c, "panic") })
	test(UnixPair, func(wg *sync.WaitGroup, c net.Conn) { childProcClose(wg, c, "close") })
	test(UnixPair, func(wg *sync.WaitGroup, c net.Conn) { childProcClose(wg, c, "closewrite") })

	test(UnixConnectedPair, func(wg *sync.WaitGroup, c net.Conn) { sameProcClose(wg, c, false) })
	test(UnixConnectedPair, func(wg *sync.WaitGroup, c net.Conn) { sameProcClose(wg, c, true) })
	test(UnixConnectedPair, func(wg *sync.WaitGroup, c net.Conn) { childProcClose(wg, c, "panic") })
	test(UnixConnectedPair, func(wg *sync.WaitGroup, c net.Conn) { childProcClose(wg, c, "close") })
	test(UnixConnectedPair, func(wg *sync.WaitGroup, c net.Conn) { childProcClose(wg, c, "closewrite") })
}