// This file contains useful functions stolen from lyrebird's cmd/lyrebird/lyrebird.go
// (and slightly modified to work with our freakish logging).

package main

import (
	"io"
	"net"
	"sync"
	"syscall"

	"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/lyrebird/common/socks5"
	"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/lyrebird/transports/base"
)

func clientAcceptLoop(f base.ClientFactory, ln net.Listener) error {
	defer ln.Close()
	for {
		conn, err := ln.Accept()
		if err != nil {
			// NOTE(eta): fixed up e.Temporary() to e.Timeout() here
			if e, ok := err.(net.Error); ok && !e.Timeout() {
				return err
			}
			continue
		}
		go clientHandler(f, conn)
	}
}

// This is a `dialFn` that lets us call `protect()` on the file descriptor before dialing.
func protectedDial(network, addr string) (net.Conn, error) {
	var dialer net.Dialer
	dialer.Control = func(_, _ string, rc syscall.RawConn) error {
		return rc.Control(func(fd uintptr) {
			Log(levelDebug, "calling protect on fd %v", fd)
			protect(int(fd))
		})
	}
	return dialer.Dial(network, addr)
}

func clientHandler(f base.ClientFactory, conn net.Conn) {
	defer conn.Close()

	name := f.Transport().Name()

	// Read the client's SOCKS handshake.
	socksReq, err := socks5.Handshake(conn)
	if err != nil {
		Log(levelError, "%s - client failed socks handshake: %s", name, err)
		return
	}
	addrStr := ElideAddr(socksReq.Target)

	// Deal with arguments.
	args, err := f.ParseArgs(&socksReq.Args)
	if err != nil {
		Log(levelError, "%s(%s) - invalid arguments: %s", name, addrStr, err)
		_ = socksReq.Reply(socks5.ReplyGeneralFailure)
		return
	}

	// Obtain the proxy dialer if any, and create the outgoing TCP connection.
	dialFn := protectedDial
	remote, err := f.Dial("tcp", socksReq.Target, dialFn, args)
	if err != nil {
		Log(levelError, "%s(%s) - outgoing connection failed: %s", name, addrStr, ElideError(err))
		_ = socksReq.Reply(socks5.ErrorToReplyCode(err))
		return
	}
	defer remote.Close()
	err = socksReq.Reply(socks5.ReplySucceeded)
	if err != nil {
		Log(levelError, "%s(%s) - SOCKS reply failed: %s", name, addrStr, ElideError(err))
		return
	}

	if err = copyLoop(conn, remote); err != nil {
		Log(levelWarn, "%s(%s) - closed connection: %s", name, addrStr, ElideError(err))
	} else {
		Log(levelInfo, "%s(%s) - closed connection", name, addrStr)
	}
}

func copyLoop(a net.Conn, b net.Conn) error {
	// Note: b is always the pt connection.  a is the SOCKS/ORPort connection.
	errChan := make(chan error, 2)

	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		defer b.Close()
		defer a.Close()
		_, err := io.Copy(b, a)
		errChan <- err
	}()
	go func() {
		defer wg.Done()
		defer a.Close()
		defer b.Close()
		_, err := io.Copy(a, b)
		errChan <- err
	}()

	// Wait for both upstream and downstream to close.  Since one side
	// terminating closes the other, the second error in the channel will be
	// something like EINVAL (though io.Copy() will swallow EOF), so only the
	// first error is returned.
	wg.Wait()
	if len(errChan) > 0 {
		return <-errChan
	}

	return nil
}
