// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package revdial implements a Dialer and Listener which work together
// to turn an accepted connection (for instance, a Hijacked HTTP request) into
// a Dialer which can then create net.Conns connecting back to the original
// dialer, which then gets a net.Listener accepting those conns.
//
// This is basically a very minimal SOCKS5 client & server.
//
// The motivation is that sometimes you want to run a server on a
// machine deep inside a NAT. Rather than connecting to the machine
// directly (which you can't, because of the NAT), you have the
// sequestered machine connect out to a public machine. Both sides
// then use revdial and the public machine can become a client for the
// NATed machine.
package revdial

import (
	"bufio"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/shellhub-io/shellhub/pkg/clock"
	"github.com/shellhub-io/shellhub/pkg/uuid"
	"github.com/shellhub-io/shellhub/pkg/wsconnadapter"
	log "github.com/sirupsen/logrus"
)

// ErrDialerClosed is returned once a dialer can no longer serve connections. It is terminal for the
// dialer, and the caller reconnects.
var ErrDialerClosed = errors.New("revdial.Dialer closed")

const (
	dialerUniqParam = "revdial.dialer"
	dialUUIDParam   = "uuid"
)

const dialerKeepAliveTimeout = 35 * time.Second

// The Dialer can create new connections.
type Dialer struct {
	conn       net.Conn // hijacked client conn
	path       string   // e.g. "/revdial"
	uniqID     string
	pickupPath string // path + uniqID: "/revdial?revdial.dialer="+uniqID

	pendingMu sync.Mutex
	pending   map[string]*pendingDial
	connReady chan string
	donec     chan struct{}
	closeOnce sync.Once
	logger    *log.Entry
	clock     clock.Clock
}

type pendingDial struct {
	conn   chan net.Conn
	failed chan error
}

var dialers = sync.Map{}

// NewDialer returns the side of the connection which will initiate
// new connections. This will typically be the side which did the HTTP
// Hijack. The connection is (typically) the hijacked HTTP client
// connection. The connPath is the HTTP path and optional query (but
// without scheme or host) on the dialer where the ConnHandler is
// mounted. The dialer reads clk to set its write deadlines.
func NewDialer(logger *log.Entry, c net.Conn, connPath string, clk clock.Clock) *Dialer {
	d := &Dialer{
		path:      connPath,
		uniqID:    newUniqID(),
		conn:      c,
		donec:     make(chan struct{}),
		connReady: make(chan string, 8),
		pending:   make(map[string]*pendingDial),
		logger:    logger,
		clock:     clk,
	}

	join := "?"
	if strings.Contains(connPath, "?") {
		join = "&"
	}
	d.pickupPath = connPath + join + dialerUniqParam + "=" + d.uniqID
	d.register()
	go d.serve() //nolint:errcheck

	d.logger.Debug("new dialer connection")

	return d
}

func newUniqID() string {
	buf := make([]byte, 16)
	rand.Read(buf)

	return hex.EncodeToString(buf)
}

func (d *Dialer) register() {
	dialers.Store(d.uniqID, d)
}

func (d *Dialer) unregister() {
	dialers.Delete(d.uniqID)
}

// Done returns a channel which is closed when d is closed (either by
// this process on purpose, by a local error, or close or error from
// the peer).
func (d *Dialer) Done() <-chan struct{} { return d.donec }

// Close closes the Dialer. It never blocks, returns nil, and every later call is a no-op.
func (d *Dialer) Close() error {
	d.closeOnce.Do(d.close)

	return nil
}

func (d *Dialer) close() {
	d.logger.Debug("dialer connection closed")

	d.unregister()
	_ = d.conn.Close()
	close(d.donec)
}

// Dial creates a new connection back to the Listener. It returns the connection the device
// dialled back for this call, the error the device reported when it could not, ErrDialerClosed
// once the dialer is closed, or ctx's error when ctx ends first. A connection that arrives
// after Dial returned is closed.
func (d *Dialer) Dial(ctx context.Context) (net.Conn, error) {
	id := uuid.Generate()

	pending := &pendingDial{
		conn:   make(chan net.Conn, 1),
		failed: make(chan error, 1),
	}

	d.addPending(id, pending)
	defer d.removePending(id, pending)

	select {
	case d.connReady <- id:
		d.logger.Debug("message true to conn ready channel")
	case <-d.donec:
		d.logger.Debug("dial done")

		return nil, ErrDialerClosed
	case <-ctx.Done():
		d.logger.Debug("dial done due context cancellation")

		return nil, ctx.Err()
	}

	select {
	case c := <-pending.conn:
		d.logger.Debug("new incoming connection")

		return c, nil
	case err := <-pending.failed:
		d.logger.Debug("failed to pick-up connection")

		return nil, err
	case <-d.donec:
		d.logger.Debug("dial done on pick-up")

		return nil, ErrDialerClosed
	case <-ctx.Done():
		d.logger.Debug("dial done on pick-up due context cancellation")

		return nil, ctx.Err()
	}
}

func (d *Dialer) addPending(id string, pending *pendingDial) {
	d.pendingMu.Lock()
	defer d.pendingMu.Unlock()

	d.pending[id] = pending
}

func (d *Dialer) removePending(id string, pending *pendingDial) {
	d.pendingMu.Lock()
	delete(d.pending, id)
	d.pendingMu.Unlock()

	pending.abandon()
}

func (d *Dialer) matchConn(id string, c net.Conn) {
	d.pendingMu.Lock()
	defer d.pendingMu.Unlock()

	pending, ok := d.pending[id]
	if !ok {
		d.logger.WithField("uuid", id).Debug("closing a connection no dial is waiting for")

		_ = c.Close()

		return
	}

	if !pending.deliver(c) {
		d.logger.WithField("uuid", id).Debug("closing a second connection for the same dial")

		_ = c.Close()
	}
}

func (d *Dialer) failPickup(connPath string, err error) {
	parsed, parseErr := url.ParseRequestURI(connPath)
	if parseErr != nil {
		d.logger.WithError(parseErr).WithField("connPath", connPath).Debug("pick-up failed for a path that names no dial")

		return
	}

	id := parsed.Query().Get(dialUUIDParam)

	d.pendingMu.Lock()
	defer d.pendingMu.Unlock()

	pending, ok := d.pending[id]
	if !ok {
		d.logger.WithError(err).WithField("uuid", id).Debug("pick-up failed for a dial that is no longer waiting")

		return
	}

	pending.fail(err)
}

func (p *pendingDial) deliver(c net.Conn) bool {
	select {
	case p.conn <- c:
		return true
	default:
		return false
	}
}

func (p *pendingDial) fail(err error) {
	select {
	case p.failed <- err:
	default:
	}
}

func (p *pendingDial) abandon() {
	select {
	case c := <-p.conn:
		_ = c.Close()
	default:
	}
}

func (d *Dialer) serve() error {
	defer d.Close() //nolint:errcheck

	go func() {
		defer d.Close() //nolint:errcheck
		defer d.logger.Debug("dialer serve done")

		br := bufio.NewReader(d.conn)
		for {
			line, err := br.ReadSlice('\n')
			if err != nil {
				d.logger.WithError(err).Trace("failed to read the agent's command")

				unexpectedError := websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure)
				if !errors.Is(err, net.ErrClosed) && unexpectedError {
					d.logger.WithError(err).Error("revdial.Dialer failed to read")
				}

				return
			}
			var msg controlMsg
			if err := json.Unmarshal(line, &msg); err != nil {
				d.logger.WithError(err).WithField("line", line).Printf("revdial.Dialer read invalid JSON")

				return
			}

			switch msg.Command {
			case "pickup-failed":
				d.failPickup(msg.ConnPath, fmt.Errorf("revdial listener failed to pick up connection: %v", msg.Err))
			case "keep-alive":
			default:
				log.WithField("message", msg.Command).Debug("unknown message received")
			}
		}
	}()
	for {
		if err := d.sendMessage(controlMsg{Command: "keep-alive"}); err != nil {
			d.logger.WithError(err).Debug("failed to send keep-alive message to device")

			return err
		}

		t := time.NewTimer(30 * time.Second)
		select {
		case <-t.C:
			continue
		case uuid := <-d.connReady:
			t.Stop()
			if err := d.sendMessage(controlMsg{
				Command:  "conn-ready",
				ConnPath: d.pickupPath + "&" + dialUUIDParam + "=" + uuid,
			}); err != nil {
				d.logger.WithError(err).Debug("failed to send conn-ready message to device")

				return err
			}
		case <-d.donec:
			t.Stop()

			return ErrDialerClosed
		}
	}
}

func (d *Dialer) sendMessage(m controlMsg) error {
	if err := d.conn.SetWriteDeadline(d.clock.Now().Add(10 * time.Second)); err != nil {
		d.logger.WithError(err).Debug("failed to set the write dead line to device")

		return err
	}

	j, err := json.Marshal(m)
	if err != nil {
		d.logger.WithError(err).Error("failed to encode the control message")

		return err
	}

	j = append(j, '\n')

	if _, err := d.conn.Write(j); err != nil {
		d.logger.WithError(err).Debug("failed to write on the connection")

		return err
	}

	return d.conn.SetWriteDeadline(time.Time{})
}

// NewListener returns a new Listener, accepting connections which
// arrive from the provided server connection, which should be after
// any necessary authentication (usually after an HTTP exchange).
//
// The provided dialServer func is responsible for connecting back to
// the server and doing TLS setup.
func NewListener(serverConn net.Conn, dialServer func(context.Context, string) (*websocket.Conn, *http.Response, error)) *Listener {
	ln := &Listener{
		sc:    serverConn,
		dial:  dialServer,
		connc: make(chan net.Conn, acceptQueueSize),
		donec: make(chan struct{}),
	}
	go ln.run()

	return ln
}

var _ net.Listener = (*Listener)(nil)

// Listener is a net.Listener, returning new connections which arrive
// from a corresponding Dialer.
type Listener struct {
	sc     net.Conn
	connc  chan net.Conn
	donec  chan struct{}
	dial   func(context.Context, string) (*websocket.Conn, *http.Response, error)
	writec chan<- []byte

	mu     sync.Mutex
	closed bool
}

type controlMsg struct {
	Command  string `json:"command,omitempty"`  // "keep-alive", "conn-ready", "pickup-failed"
	ConnPath string `json:"connPath,omitempty"` // conn pick-up URL path for "conn-url", "pickup-failed"
	Err      string `json:"err,omitempty"`
}

func (ln *Listener) run() {
	done := func() {
		_ = ln.Close()
	}

	var onceDefer sync.Once
	defer onceDefer.Do(done)

	closeTimer := time.AfterFunc(dialerKeepAliveTimeout, done)

	writec := make(chan []byte, controlWriteQueueSize)
	ln.writec = writec
	go func() {
		defer onceDefer.Do(done)

		for {
			select {
			case <-ln.donec:
				return
			case msg := <-writec:
				if _, err := ln.sc.Write(msg); err != nil {
					log.Printf("revdial.Listener: error writing message to server: %v", err)

					return
				}
			}
		}
	}()

	go func() {
		defer onceDefer.Do(done)

		br := bufio.NewReader(ln.sc)
		for {
			line, err := br.ReadSlice('\n')
			if err != nil {
				return
			}
			var msg controlMsg
			if err := json.Unmarshal(line, &msg); err != nil {
				log.Printf("revdial.Listener read invalid JSON: %q: %v", line, err)

				return
			}
			switch msg.Command {
			case "keep-alive":
				closeTimer.Reset(dialerKeepAliveTimeout)
			case "conn-ready":
				go ln.grabConn(msg.ConnPath)
			default:
			}
		}
	}()

	for {
		ln.sendMessage(controlMsg{Command: "keep-alive"})

		t := time.NewTimer(30 * time.Second)
		select {
		case <-t.C:
			continue
		case <-ln.donec:
			t.Stop()

			return
		}
	}
}

func (ln *Listener) sendMessage(m controlMsg) {
	j, err := json.Marshal(m)
	if err != nil {
		log.WithError(err).Error("failed to encode the control message")

		return
	}

	j = append(j, '\n')

	select {
	case ln.writec <- j:
	case <-ln.donec:
	}
}

func (ln *Listener) grabConn(path string) {
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()

	wsConn, resp, err := ln.dial(ctx, path) //nolint:bodyclose // gorilla/websocket documents that the handshake response body need not be closed.
	if err != nil {
		ln.sendMessage(controlMsg{Command: "pickup-failed", ConnPath: path, Err: err.Error()})

		return
	}

	failPickup := func(err error) {
		_ = wsConn.Close()
		log.Printf("revdial.Listener: failed to pick up connection to %s: %v", path, err)
		ln.sendMessage(controlMsg{Command: "pickup-failed", ConnPath: path, Err: err.Error()})
	}

	if resp.StatusCode != http.StatusSwitchingProtocols {
		failPickup(fmt.Errorf("non-101 response %v", resp.Status))

		return
	}

	err = ln.enqueue(wsconnadapter.New(wsConn))

	switch {
	case errors.Is(err, ErrListenerClosed):
		_ = wsConn.Close()
	case err != nil:
		failPickup(err)
	}
}

const (
	acceptQueueSize       = 8
	controlWriteQueueSize = 8
)

var errTooManyPickups = errors.New("too many picked-up connections are waiting to be accepted")

func (ln *Listener) enqueue(c net.Conn) error {
	ln.mu.Lock()
	defer ln.mu.Unlock()

	if ln.closed {
		return ErrListenerClosed
	}

	select {
	case ln.connc <- c:
		return nil
	default:
		return errTooManyPickups
	}
}

// Accept blocks and returns a new connection, or ErrListenerClosed once the Listener is closed:
// by Close, or because its control connection to the server failed, carried a malformed message,
// or stopped delivering keep-alives.
func (ln *Listener) Accept() (net.Conn, error) {
	select {
	case c := <-ln.connc:
		return c, nil
	case <-ln.donec:
		return nil, ErrListenerClosed
	}
}

// ErrListenerClosed is returned by Accept once the Listener is closed. Accept lists what closes it.
var ErrListenerClosed = errors.New("revdial: Listener closed")

// Close closes the Listener, making future Accept calls return
// ErrListenerClosed, and closes every picked-up connection that no Accept has
// taken yet. A connection an Accept running concurrently takes belongs to that
// caller. It never blocks and always returns nil.
func (ln *Listener) Close() error {
	ln.mu.Lock()
	defer ln.mu.Unlock()
	if ln.closed {
		return nil
	}
	go func() { _ = ln.sc.Close() }()
	ln.closed = true
	close(ln.donec)

	for {
		select {
		case c := <-ln.connc:
			_ = c.Close()
		default:
			return nil
		}
	}
}

// Addr returns a dummy address. This exists only to conform to the
// net.Listener interface.
func (ln *Listener) Addr() net.Addr { return fakeAddr{} }

type fakeAddr struct{}

func (fakeAddr) Network() string { return "revdial" }
func (fakeAddr) String() string  { return "revdialconn" }

// ConnHandler returns the HTTP handler that needs to be mounted somewhere
// that the Listeners can dial out and get to. A dialer to connect to it
// is given to NewListener and the path to reach it is given to NewDialer
// to use in messages to the listener. It answers an unknown dialer with 400,
// and when the upgrade fails it logs the cause through the dialer's logger and
// writes nothing, because the upgrader has already written the response.
func ConnHandler(upgrader websocket.Upgrader) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		dialerUniq := r.FormValue(dialerUniqParam)
		uuid := r.FormValue(dialUUIDParam)

		d, _ := dialers.Load(dialerUniq)

		dialer, ok := d.(*Dialer)
		if !ok {
			http.Error(w, "unknown dialer", http.StatusBadRequest)

			return
		}

		wsConn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			dialer.logger.WithError(err).WithField("uuid", uuid).Warn("failed to upgrade the pickup connection")

			return
		}

		dialer.matchConn(uuid, wsconnadapter.New(wsConn))
	})
}