endpoint.go

package connet

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"sync/atomic"

	"github.com/quic-go/quic-go"

	"github.com/connet-dev/connet/pkg/iterc"
	"github.com/connet-dev/connet/pkg/proto"
	"github.com/connet-dev/connet/pkg/proto/pbcontrol"
	"github.com/connet-dev/connet/pkg/quicc"
	"github.com/connet-dev/connet/pkg/reliable"
	"github.com/connet-dev/connet/pkg/slogc"
	"github.com/connet-dev/connet/pkg/statusc"
)

var errEndpointClosed = fmt.Errorf("endpoint closed")

type endpointStatus struct {
	Status statusc.Status
	StatusPeer
}

type endpoint struct {
	client *Client
	cfg    endpointConfig
	peer   *peer

	ctx    context.Context
	cancel context.CancelCauseFunc
	closer chan struct{}

	onlineReport func(err error)
	connStatus   atomic.Value

	logger *slog.Logger
}

// a client endpoint could close when:
//   - the user cancels the incoming context. This could happen while setting up the endpoint too.
//   - the user calls Close explicitly.
//   - the parent client is closing, so it calls close on the endpoint too. Session might be closing at the same time.
//   - an error happens in runPeer
//   - a terminal error happens in runAnnounce
func newEndpoint(ctx context.Context, cl *Client, cfg endpointConfig, logger *slog.Logger) (*endpoint, error) {
	p, err := newPeer(cl.directServer, cl.addrs, cl.metadata, logger)
	if err != nil {
		return nil, err
	}

	ctx, cancel := context.WithCancelCause(ctx)
	ep := &endpoint{
		client: cl,
		cfg:    cfg,
		peer:   p,

		ctx:    ctx,
		cancel: cancel,
		closer: make(chan struct{}),

		logger: logger,
	}
	ep.connStatus.Store(statusc.NotConnected)
	context.AfterFunc(ctx, ep.cleanup)

	errCh := make(chan error, 1)
	var reportOnce sync.Once
	ep.onlineReport = func(err error) {
		if err == nil {
			if !ep.connStatus.CompareAndSwap(statusc.Reconnecting, statusc.Connected) {
				ep.connStatus.CompareAndSwap(statusc.NotConnected, statusc.Connected)
			}
		}
		reportOnce.Do(func() {
			if err != nil {
				errCh <- err
			}
			close(errCh)
		})
	}

	go ep.runPeer(ctx)
	go ep.runSession(ctx)

	select {
	case <-ctx.Done():
		ep.cancel(ctx.Err())
		return nil, ctx.Err()
	case err := <-errCh:
		if err != nil {
			ep.cancel(err)
			return nil, err
		}
	}

	return ep, nil
}

func (ep *endpoint) status() (endpointStatus, error) {
	peerStatus, err := ep.peer.status()
	if err != nil {
		return endpointStatus{}, err
	}
	return endpointStatus{
		Status:     ep.connStatus.Load().(statusc.Status),
		StatusPeer: peerStatus,
	}, nil
}

func (ep *endpoint) close(err error) error {
	ep.cancel(err)
	<-ep.closer
	return nil
}

func (ep *endpoint) runPeer(ctx context.Context) {
	if err := ep.peer.run(ctx); err != nil {
		ep.logger.Debug("peer stopped", "err", err)
		ep.cancel(err)
	}
}

func (ep *endpoint) runSession(ctx context.Context) {
	err := ep.client.currentSession.Listen(ctx, func(sess *session) error {
		if sess != nil {
			go ep.runSessionAnnounce(ctx, sess)
		}
		return nil
	})
	if err != nil {
		ep.cancel(err)
	}
}

func (ep *endpoint) runSessionAnnounce(ctx context.Context, sess *session) {
	for {
		err := ep.runSessionAnnounceErr(ctx, sess)
		ep.connStatus.CompareAndSwap(statusc.Connected, statusc.Reconnecting)

		switch {
		case err == nil:
		case errors.Is(err, context.Canceled):
			return
		case sess.conn.Context().Err() != nil:
			return
		default:
			ep.logger.Debug("announce stopped", "err", err)
		}
	}
}

func (ep *endpoint) runSessionAnnounceErr(ctx context.Context, sess *session) error {
	if ep.cfg.route.AllowRelay() {
		g := reliable.NewGroup(ctx)
		g.Go(reliable.Bind2(sess.conn, sess.wv, ep.runAnnounce))
		g.Go(reliable.Bind2(sess.conn, sess.wv, ep.runRelay))
		return g.Wait()
	}

	return ep.runAnnounce(ctx, sess.conn, sess.wv)
}

func (ep *endpoint) runAnnounce(ctx context.Context, conn *quic.Conn, wv proto.WireVersion) error {
	stream, err := conn.OpenStreamSync(ctx)
	if err != nil {
		return fmt.Errorf("announce open stream: %w", err)
	}
	defer func() {
		if err := stream.Close(); err != nil {
			slogc.Fine(ep.logger, "error closing announce stream", "err", err)
		}
	}()

	g := reliable.NewGroup(ctx)
	g.Go(quicc.CancelStream(stream))

	g.Go(func(ctx context.Context) error {
		defer ep.logger.Debug("completed announce notify")
		return ep.peer.selfListen(ctx, func(rawPeer *pbcontrol.Peer) error {
			peer := &pbcontrol.Peer{
				ServerCertificate: rawPeer.ServerCertificate,
				ClientCertificate: rawPeer.ClientCertificate,
			}
			if ep.cfg.route.AllowDirect() {
				peer.Directs = rawPeer.Directs
			}
			if ep.cfg.route.AllowRelay() {
				peer.RelayIds = rawPeer.RelayIds
			}

			ep.logger.Debug("updated announce", "direct", len(peer.Directs), "relays", len(peer.RelayIds))
			return proto.Write(stream, &pbcontrol.Request{
				Announce: &pbcontrol.Request_Announce{
					Endpoint: ep.cfg.endpoint.PB(),
					Role:     ep.cfg.role.PB(),
					Peer:     peer,
				},
			}, wv)
		})
	})

	g.Go(func(ctx context.Context) error {
		for {
			resp, err := pbcontrol.ReadResponse(stream, wv)
			ep.onlineReport(err)
			if err != nil {
				return err
			}
			if resp.Announce == nil {
				return fmt.Errorf("announce unexpected response")
			}

			// TODO on server restart peers is reset and client loses active peers
			// only for them to come back at the next tick, with different ID
			ep.peer.setPeers(iterc.MapSlice(resp.Announce.Peers, func(rawPeer *pbcontrol.RemotePeer) *pbcontrol.RemotePeer {
				peer := &pbcontrol.RemotePeer{
					Id:       rawPeer.Id,
					Metadata: rawPeer.Metadata,
					Peer: &pbcontrol.Peer{
						ServerCertificate: rawPeer.Peer.ServerCertificate,
						ClientCertificate: rawPeer.Peer.ClientCertificate,
					},
				}

				if ep.cfg.route.AllowDirect() {
					peer.Peer.Directs = rawPeer.Peer.Directs
				}
				if ep.cfg.route.AllowRelay() {
					peer.Peer.RelayIds = rawPeer.Peer.RelayIds
				}

				return peer
			}))
		}
	})

	return g.Wait()
}

func (ep *endpoint) runRelay(ctx context.Context, conn *quic.Conn, wv proto.WireVersion) error {
	stream, err := conn.OpenStreamSync(ctx)
	if err != nil {
		return fmt.Errorf("relay open stream: %w", err)
	}
	defer func() {
		if err := stream.Close(); err != nil {
			slogc.Fine(ep.logger, "error closing relay stream", "err", err)
		}
	}()

	if err := proto.Write(stream, &pbcontrol.Request{
		Relay: &pbcontrol.Request_Relay{
			Endpoint:          ep.cfg.endpoint.PB(),
			Role:              ep.cfg.role.PB(),
			ClientCertificate: ep.peer.clientCert.Leaf.Raw,
		},
	}, wv); err != nil {
		return err
	}

	g := reliable.NewGroup(ctx)
	g.Go(quicc.CancelStream(stream))

	g.Go(func(ctx context.Context) error {
		for {
			resp, err := pbcontrol.ReadResponse(stream, wv)
			if err != nil {
				ep.onlineReport(err)
				return err
			}
			if resp.Relay == nil {
				return fmt.Errorf("relay unexpected response")
			}

			ep.peer.setRelays(resp.Relay.Relays)
		}
	})

	return g.Wait()
}

func (ep *endpoint) cleanup() {
	defer close(ep.closer)
	defer ep.connStatus.Store(statusc.Disconnected)

	switch ep.cfg.role {
	case RoleDestination:
		ep.client.removeDestination(ep.cfg.endpoint)
	case RoleSource:
		ep.client.removeSource(ep.cfg.endpoint)
	}
}