// Copyright (C) 2016-2021  Nexedi SA and Contributors.
//                          Kirill Smelkov <kirr@nexedi.com>
//
// This program is free software: you can Use, Study, Modify and Redistribute
// it under the terms of the GNU General Public License version 3, or (at your
// option) any later version, as published by the Free Software Foundation.
//
// You can also Link and Combine this program with other software covered by
// the terms of any of the Free Software licenses or any of the Open Source
// Initiative approved licenses and Convey the resulting work. Corresponding
// source of such a combination shall include the source code for all other
// software used.
//
// This program is distributed WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
//
// See COPYING file for full licensing terms.
// See https://www.nexedi.com/licensing for rationale and options.

package neonet
// link establishment

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net"

	"github.com/tinylib/msgp/msgp"

	"lab.nexedi.com/kirr/go123/xerr"
	"lab.nexedi.com/kirr/go123/xnet"
	"lab.nexedi.com/kirr/neo/go/internal/xcontext"
	"lab.nexedi.com/kirr/neo/go/internal/xio"
	"lab.nexedi.com/kirr/neo/go/neo/internal/tneonet"
	"lab.nexedi.com/kirr/neo/go/neo/proto"
)

// ---- Handshake ----

// XXX _Handshake{Client,Server} may be needed to become public in case when we have already
// established raw connection and want to hand-over it to NEO. But currently we
// do not have such uses.

func _HandshakeClient(ctx context.Context, conn net.Conn) (*NodeLink, error) {
	return handshakeClient(ctx, conn, proto.Version, dialEncTryOrder[0])
}

func _HandshakeServer(ctx context.Context, conn net.Conn) (*NodeLink, error) {
	return handshakeServer(ctx, conn, proto.Version)
}

// _HandshakeError is returned when there is an error while performing handshake.
type _HandshakeError struct {
	LocalRole  _LinkRole
	LocalAddr  net.Addr
	RemoteAddr net.Addr
	Err        error
}

func (e *_HandshakeError) Error() string {
	role := ""
	switch e.LocalRole {
	case _LinkServer: role = "server"
	case _LinkClient: role = "client"
	default:          panic("bug")
	}
	return fmt.Sprintf("%s - %s: handshake (%s): %s", e.LocalAddr, e.RemoteAddr, role, e.Err.Error())
}

func (e *_HandshakeError) Cause()  error { return e.Err }
func (e *_HandshakeError) Unwrap() error { return e.Err }

// handshakeClient implements client-side NEO protocol handshake just after raw
// connection between 2 nodes was established.
//
// Client indicates its version and preferred encoding, but accepts any
// encoding chosen to use by server.
//
// On success raw connection is returned wrapped into NodeLink.
// On error raw connection is closed.
func handshakeClient(ctx context.Context, conn net.Conn, version uint32, encPrefer proto.Encoding) (*NodeLink, error) {
	enc, rxbuf, err := _handshakeClient(ctx, conn, version, encPrefer)
	if err != nil {
		return nil, err
	}
	return newNodeLink(conn, enc, _LinkClient, rxbuf), nil
}

// handshakeServer implements server-side NEO protocol handshake just after raw
// connection between 2 nodes was established.
//
// Server verifies that its version matches Client and accepts client preferred encoding.
//
// On success raw connection is returned wrapped into NodeLink.
// On error raw connection is closed.
func handshakeServer(ctx context.Context, conn net.Conn, version uint32) (*NodeLink, error) {
	enc, rxbuf, err := _handshakeServer(ctx, conn, version)
	if err != nil {
		return nil, err
	}
	return newNodeLink(conn, enc, _LinkServer, rxbuf), nil
}

func _handshakeClient(ctx context.Context, conn net.Conn, version uint32, encPrefer proto.Encoding) (enc proto.Encoding, rxbuf *xbufReader, err error) {
	defer func() {
		if err != nil {
			err = &_HandshakeError{_LinkClient, conn.LocalAddr(), conn.RemoteAddr(), err}
		}
	}()

	rxbuf = newXBufReader(conn, /*any non-small limit*/1024)

	var peerEnc proto.Encoding
	err = xcontext.WithCloseOnErrCancel(ctx, conn, func() error {
		// tx client hello
		err := txHello("tx hello", conn, version, encPrefer)
		if err != nil {
			return err
		}

		// rx server hello reply
		var peerVer uint32
		peerEnc, peerVer, err = rxHello("rx hello reply", rxbuf)
		if err != nil {
			return err
		}

		// verify version
		if peerVer != version {
			return fmt.Errorf("protocol version mismatch: peer = %08x  ; our side = %08x", peerVer, version)
		}

		return nil
	})
	if err != nil {
		return 0, nil, err
	}

	// use peer encoding (server should return the same, but we are ok if
	// it asks to switch to different)
	return peerEnc, rxbuf, nil
}

func _handshakeServer(ctx context.Context, conn net.Conn, version uint32) (enc proto.Encoding, rxbuf *xbufReader, err error) {
	defer func() {
		if err != nil {
			err = &_HandshakeError{_LinkServer, conn.LocalAddr(), conn.RemoteAddr(), err}
		}
	}()

	rxbuf = newXBufReader(conn, /*any non-small limit*/1024)

	var peerEnc proto.Encoding
	err = xcontext.WithCloseOnErrCancel(ctx, conn, func() error {
		// rx client hello
		var peerVer uint32
		var err     error
		peerEnc, peerVer, err = rxHello("rx hello", rxbuf)
		if err != nil {
			return err
		}

		// tx server reply
		//
		// do it before version check so that client can also detect "version
		// mismatch" instead of just getting "disconnect".
		err = txHello("tx hello reply", conn, version, peerEnc)
		if err != nil {
			return err
		}

		// verify version
		if peerVer != version {
			return fmt.Errorf("protocol version mismatch: peer = %08x  ; our side = %08x", peerVer, version)
		}

		return nil
	})
	if err != nil {
		return 0, nil, err
	}

	return peerEnc, rxbuf, nil
}

// handshake hello:
//
// - 00 00 00 <ver>	for 'N' encoding, and
// - 92 c4 03 NEO ...	for 'M' encoding (= msgpack of (b"NEO", <ver>))
//
// the first byte is different from TLS handshake (0x16).

func txHello(errctx string, conn net.Conn, version uint32, enc proto.Encoding) (err error) {
	defer xerr.Context(&err, errctx)

	var b []byte
	switch enc {
	case 'N':
		// 00 00 00 <v>
		b = make([]byte, 4)
		if version > 0xff {
			panic("encoding N supports versions only in range [0, 0xff]")
		}
		b[3] = uint8(version)

	case 'M':
		// (b"NEO", <V>) encoded as msgpack (= 92 c4 03 NEO int(<V>))
		b = msgp.AppendArrayHeader(b, 2)       // 92
		b = msgp.AppendBytes(b, []byte("NEO")) // c4 03 NEO
		b = msgp.AppendUint32(b, version)      // u?intX version

	default:
		panic("bug")
	}

	_, err = conn.Write(b)
	if err != nil {
		return err
	}

	return nil
}

func rxHello(errctx string, rx *xbufReader) (enc proto.Encoding, version uint32, err error) {
	defer xerr.Context(&err, errctx)

	b := make([]byte, 4)
	_, err = io.ReadFull(rx, b)
	err = xio.NoEOF(err)
	if err != nil {
		return 0, 0, err
	}

	var peerEnc proto.Encoding
	var peerVer uint32
	badMagic := false
	switch {
	case bytes.Equal(b[:3], []byte{0,0,0}):
		peerEnc = 'N'
		peerVer = uint32(b[3])

	case bytes.Equal(b, []byte{0x92, 0xc4, 3, 'N'}): // start of "fixarray<2> bin8 'N | EO' ...
		b = append(b, []byte{0,0}...)
		_, err = io.ReadFull(rx, b[4:])
		err = xio.NoEOF(err)
		if err != nil {
			return 0, 0, err
		}
		if !bytes.Equal(b[4:], []byte{'E','O'}) {
			badMagic = true
			break
		}

		peerEnc = 'M'
		rxM := msgp.Reader{R: rx.Reader}
		peerVer, err = rxM.ReadUint32()
		if err != nil {
			return 0, 0, fmt.Errorf("M: recv peer version: %s", err) // XXX + "read magic" ctx
		}

	default:
		badMagic = true
	}

	if badMagic {
		return 0, 0, fmt.Errorf("invalid magic %x", b)
	}

	return peerEnc, peerVer, nil
}


// ---- Dial & Listen at NodeLink level ----

// dialEncTryOrder is the order of trials for encoding when establishing a NEO link from client side.
//
// NEO/go server autodetects client preferred encoding and adjusts to that, but
// NEO/py rejects connections if client encoding does not exactly match server.
//
// To autodetect what NEO/py server uses DialLink retries dial + handshake with
// client-preferred encodings in the order specified by dialEncTryOrder.
//
// NOTE tests change this to verify that autodetection of NEO/py server
// encoding actually works.
var dialEncTryOrder = []proto.Encoding{'N', 'M'}
func init() {
	tneonet.DialEncTryOrder = &dialEncTryOrder
}

// DialLink connects to address on given network, performs NEO protocol
// handshake and wraps the connection as NodeLink.
func DialLink(ctx context.Context, net xnet.Networker, addr string) (link *NodeLink, err error) {
	for _, enc := range dialEncTryOrder {
		peerConn, err := net.Dial(ctx, addr)
		if err != nil {
			return nil, err
		}

		link, err = handshakeClient(ctx, peerConn, proto.Version, enc)

		// NEO/py closes connection if it sees unexpected magic, version, etc.
		// -> in such case retry with next encoding trying to autodetect and match server.
		// -> stop trying on success, or on any other error.
		if err == nil || !errors.Is(err, io.ErrUnexpectedEOF) {
			break
		}
	}

	// either link is established or the error indicates we should not retry anymore
	return link, err
}

// ListenLink starts listening on laddr for incoming connections and wraps them as NodeLink.
//
// The listener accepts only those connections that pass NEO protocol handshake.
func ListenLink(ctx context.Context, net xnet.Networker, laddr string) (LinkListener, error) {
	rawl, err := net.Listen(ctx, laddr)
	if err != nil {
		return nil, err
	}

	return NewLinkListener(rawl), nil
}

// NewLinkListener creates LinkListener which accepts connections from an inner
// net.Listener and wraps them as NodeLink.
//
// The listener accepts only those connections that pass NEO protocol handshake.
func NewLinkListener(inner xnet.Listener) LinkListener {
	return &linkListener{l: inner}
}

// LinkListener is xnet.Listener adapted to return handshaked NodeLink on Accept.
type LinkListener interface {
	// from xnet.Listener:
	Close() error
	Addr() net.Addr

	// Accept returns new incoming connection wrapped into NodeLink.
	// It accepts only those connections which pass NEO protocol handshake.
	Accept(ctx context.Context) (*NodeLink, error)
}

// linkListener implements LinkListener.
type linkListener struct {
	l xnet.Listener
}

func (l *linkListener) Accept(ctx context.Context) (*NodeLink, error) {
	// XXX err ctx?

	conn, err := l.l.Accept(ctx)
	if err != nil {
		return nil, err
	}

	// NOTE Handshake closes conn in case of failure
	link, err := _HandshakeServer(ctx, conn)
	if err != nil {
		return nil, err
	}

	return link, nil
}

func (l *linkListener) Close() error   { return l.l.Close() }
func (l *linkListener) Addr() net.Addr { return l.l.Addr()  }
