package vpn

import (

"context"

"fmt"

"io"

"strings"

"sync"

"golang.org/x/xerrors"

"google.golang.org/protobuf/proto"

"cdr.dev/slog"

)

type SpeakerRole string

type rpcMessage interface {

proto.Message

GetRpc() *RPC

// EnsureRPC isn't autogenerated, but we'll manually add it for RPC types so that the speaker

// can allocate the RPC.

EnsureRPC() *RPC

}

func (t *TunnelMessage) EnsureRPC() *RPC {

if t.Rpc == nil {

t.Rpc = &RPC{}

}

return t.Rpc

}

func (m *ManagerMessage) EnsureRPC() *RPC {

if m.Rpc == nil {

m.Rpc = &RPC{}

}

return m.Rpc

}

// receivableRPCMessage is an rpcMessage that we can receive, and unmarshal, using generics, from a

// byte stream. proto.Unmarshal requires us to have already allocated the memory for the message

// type we are unmarshalling. All our message types are pointers like *TunnelMessage, so to

// allocate, the compiler needs to know:

//

// a) that the type is a pointer type

// b) what type it is pointing to

//

// So, this generic interface requires that the message is a pointer to the type RR. Then, we pass

// both the receivableRPCMessage and RR as type constraints, so that we'll have access to the

// underlying type when it comes time to allocate it. It's a bit messy, but the alternative is

// reflection, which has its own challenges in understandability.

type receivableRPCMessage[RR any] interface {

rpcMessage

*RR

}

const (

SpeakerRoleManager SpeakerRole = "manager"

SpeakerRoleTunnel SpeakerRole = "tunnel"

)

// speaker is an implementation of the CoderVPN protocol. It handles unary RPCs and their responses,

// as well as the low-level serialization & deserialization to the ReadWriteCloser (rwc).

//

// ┌────────┐ sendCh

// ◄─────│ ◄────────────────────────────────────────────────────────────────── ◄┐

// │ │ ▲ rpc requests

// rwc │ serdes │ │ │ sendReply()

// │ │ ┌───────────────────┐ ┌──────┼──────┐

// ──────► ┼────────► recvFromSerdes() │ rpc │rpc handling │ │

// └────────┘ recvCh │ ┼────────────► ◄──── unaryRPC()

// │ │ responses │ │ │

// │ │ │ │

// │ │ └─────────────┘ ┌ ─ ─│─ ─ ─ ─ ─ ─ ─ ┐

// │ ┼──────────────────────────────────────────► request handling

// └───────────────────┘ requests (outside speaker)

// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘

//

// speaker is implemented as a generic type that accepts the type of message we send (S), the type we receive (R), and

// the underlying type that R points to (RR). The speaker is intended to be wrapped by another, non-generic type for

// the role (manager or tunnel). E.g. Tunnel from this package.

//

// The serdes handles SERialiazation and DESerialization of the low level message types. The wrapping type may send

// non-RPC messages (that is messages that don't expect an explicit reply) by sending on the sendCh.

//

// Unary RPCs are handled by the unaryRPC() function, which handles sending the message and waiting for the response.

//

// recvFromSerdes() reads all incoming messages from the serdes. If they are RPC responses, it dispatches them to the

// waiting unaryRPC() function call, if any. If they are RPC requests or non-RPC messages, it wraps them in a request

// struct and sends them over the requests chan. The manager/tunnel role type must read from this chan and handle

// the requests. If they are RPC types, it should call sendReply() on the request with the reply message.

type speaker[S rpcMessage, R receivableRPCMessage[RR], RR any] struct {

serdes *serdes[S, R, RR]

requests chan *request[S, R]

logger slog.Logger

nextMsgID uint64

ctx context.Context

cancel context.CancelFunc

sendCh chan<- S

recvCh <-chan R

recvLoopDone chan struct{}

mu sync.Mutex

responseChans map[uint64]chan R

}

// newSpeaker creates a new protocol speaker.

func newSpeaker[S rpcMessage, R receivableRPCMessage[RR], RR any](

ctx context.Context, logger slog.Logger, conn io.ReadWriteCloser,

me, them SpeakerRole,

) (

*speaker[S, R, RR], error,

) {

ctx, cancel := context.WithCancel(ctx)

if err := handshake(ctx, conn, logger, me, them); err != nil {

cancel()

return nil, xerrors.Errorf("handshake failed: %w", err)

}

sendCh := make(chan S)

recvCh := make(chan R)

s := &speaker[S, R, RR]{

serdes: newSerdes(ctx, logger, conn, sendCh, recvCh),

logger: logger,

requests: make(chan *request[S, R]),

responseChans: make(map[uint64]chan R),

nextMsgID: 1,

ctx: ctx,

cancel: cancel,

sendCh: sendCh,

recvCh: recvCh,

recvLoopDone: make(chan struct{}),

}

return s, nil

}

// start starts the serialzation/deserialization. It's important this happens

// after any assignments of the speaker to its owning Tunnel or Manager, since

// the mutex is copied and that is not threadsafe.

// nolint: revive

func (s *speaker[_, _, _]) start() {

s.serdes.start()

go s.recvFromSerdes()

}

func (s *speaker[S, R, _]) recvFromSerdes() {

defer close(s.recvLoopDone)

defer close(s.requests)

for {

select {

case <-s.ctx.Done():

s.logger.Debug(s.ctx, "recvFromSerdes context done while waiting for proto", slog.Error(s.ctx.Err()))

return

case msg, ok := <-s.recvCh:

if !ok {

s.logger.Debug(s.ctx, "recvCh is closed")

return

}

rpc := msg.GetRpc()

if rpc != nil && rpc.ResponseTo != 0 {

// this is a unary response

s.tryToDeliverResponse(msg)

continue

}

req := &request[S, R]{

ctx: s.ctx,

msg: msg,

replyCh: s.sendCh,

}

select {

case <-s.ctx.Done():

s.logger.Debug(s.ctx, "recvFromSerdes context done while waiting for request handler", slog.Error(s.ctx.Err()))

return

case s.requests <- req:

}

}

}

}

// Close closes the speaker

// nolint: revive

func (s *speaker[_, _, _]) Close() error {

s.cancel()

err := s.serdes.Close()

return err

}

// unaryRPC sends a request/response style RPC over the protocol, waits for the response, then

// returns the response

func (s *speaker[S, R, _]) unaryRPC(ctx context.Context, req S) (resp R, err error) {

rpc := req.EnsureRPC()

msgID, respCh := s.newRPC()

rpc.MsgId = msgID

logger := s.logger.With(slog.F("msg_id", msgID))

select {

case <-ctx.Done():

return resp, ctx.Err()

case <-s.ctx.Done():

return resp, xerrors.Errorf("vpn protocol closed: %w", s.ctx.Err())

case <-s.recvLoopDone:

logger.Debug(s.ctx, "recvLoopDone while sending request")

return resp, io.ErrUnexpectedEOF

case s.sendCh <- req:

logger.Debug(s.ctx, "sent rpc request", slog.F("req", req))

}

select {

case <-ctx.Done():

s.rmResponseChan(msgID)

return resp, ctx.Err()

case <-s.ctx.Done():

s.rmResponseChan(msgID)

return resp, xerrors.Errorf("vpn protocol closed: %w", s.ctx.Err())

case <-s.recvLoopDone:

logger.Debug(s.ctx, "recvLoopDone while waiting for response")

return resp, io.ErrUnexpectedEOF

case resp = <-respCh:

logger.Debug(s.ctx, "got response", slog.F("resp", resp))

return resp, nil

}

}

func (s *speaker[_, R, _]) newRPC() (uint64, chan R) {

s.mu.Lock()

defer s.mu.Unlock()

msgID := s.nextMsgID

s.nextMsgID++

c := make(chan R)

s.responseChans[msgID] = c

return msgID, c

}

func (s *speaker[_, _, _]) rmResponseChan(msgID uint64) {

s.mu.Lock()

defer s.mu.Unlock()

delete(s.responseChans, msgID)

}

func (s *speaker[_, R, _]) tryToDeliverResponse(resp R) {

msgID := resp.GetRpc().GetResponseTo()

s.mu.Lock()

defer s.mu.Unlock()

c, ok := s.responseChans[msgID]

if ok {

c <- resp

// Remove the channel since we delivered a response. This ensures that each response channel

// gets _at most_ one response. Since the channels are buffered with size 1, send will

// never block.

delete(s.responseChans, msgID)

}

}

// handshake performs the initial CoderVPN protocol handshake over the given conn

func handshake(

ctx context.Context, conn io.ReadWriteCloser, logger slog.Logger, me, them SpeakerRole,

) error {

// read and write simultaneously to avoid deadlocking if the conn is not buffered

errCh := make(chan error, 2)

go func() {

ours := headerString(me, CurrentSupportedVersions)

_, err := conn.Write([]byte(ours))

logger.Debug(ctx, "wrote out header")

if err != nil {

err = xerrors.Errorf("write header: %w", err)

}

errCh <- err

}()

headerCh := make(chan string, 1)

go func() {

// we can't use bufio.Scanner here because we need to ensure we don't read beyond the

// first newline. So, we'll read one byte at a time. It's inefficient, but the initial

// header is only a few characters, so we'll keep this code simple.

buf := make([]byte, 256)

have := 0

for {

_, err := conn.Read(buf[have : have+1])

if err != nil {

errCh <- xerrors.Errorf("read header: %w", err)

return

}

if buf[have] == '\n' {

logger.Debug(ctx, "got newline header delimiter")

// use have (not have+1) since we don't want the delimiter for verification.

headerCh <- string(buf[:have])

return

}

have++

if have >= len(buf) {

errCh <- xerrors.Errorf("header malformed or too large: %s", string(buf))

return

}

}

}()

writeOK := false

theirHeader := ""

readOK := false

for !(readOK && writeOK) {

select {

case <-ctx.Done():

_ = conn.Close() // ensure our read/write goroutines get a chance to clean up

return ctx.Err()

case err := <-errCh:

if err == nil {

// write goroutine sends nil when completing successfully.

logger.Debug(ctx, "write ok")

writeOK = true

continue

}

_ = conn.Close()

return err

case theirHeader = <-headerCh:

logger.Debug(ctx, "read ok")

readOK = true

}

}

logger.Debug(ctx, "handshake read/write complete", slog.F("their_header", theirHeader))

gotVersion, err := validateHeader(theirHeader, them, CurrentSupportedVersions)

if err != nil {

return xerrors.Errorf("validate header (%s): %w", theirHeader, err)

}

logger.Debug(ctx, "handshake validated", slog.F("common_version", gotVersion))

// TODO: actually use the common version to perform different behavior once

// we have multiple versions

return nil

}

const headerPreamble = "codervpn"

func headerString(role SpeakerRole, versions RPCVersionList) string {

return fmt.Sprintf("%s %s %s\n", headerPreamble, role, versions.String())

}

func validateHeader(header string, expectedRole SpeakerRole, supportedVersions RPCVersionList) (RPCVersion, error) {

parts := strings.Split(header, " ")

if len(parts) != 3 {

return RPCVersion{}, xerrors.New("wrong number of parts")

}

if parts[0] != headerPreamble {

return RPCVersion{}, xerrors.New("invalid preamble")

}

if parts[1] != string(expectedRole) {

return RPCVersion{}, xerrors.New("unexpected role")

}

otherVersions, err := ParseRPCVersionList(parts[2])

if err != nil {

return RPCVersion{}, xerrors.Errorf("parse version list %q: %w", parts[2], err)

}

compatibleVersion, ok := supportedVersions.IsCompatibleWith(otherVersions)

if !ok {

return RPCVersion{},

xerrors.Errorf("current supported versions %q is not compatible with peer versions %q", supportedVersions.String(), otherVersions.String())

}

return compatibleVersion, nil

}

type request[S rpcMessage, R rpcMessage] struct {

ctx context.Context

msg R

replyCh chan<- S

}

func (r *request[S, _]) sendReply(reply S) error {

rrpc := reply.EnsureRPC()

mrpc := r.msg.GetRpc()

if mrpc == nil {

return xerrors.Errorf("message didn't want a reply")

}

rrpc.ResponseTo = mrpc.MsgId

select {

case <-r.ctx.Done():

return r.ctx.Err()

case r.replyCh <- reply:

}

return nil

}