TUN-8424: Refactor capnp registration server

Move RegistrationServer and RegistrationClient into tunnelrpc module
to properly abstract out the capnp aspects internal to the module only.

tunnelrpc/pogs/registration_server.go

@@ -105,6 +105,13 @@ type RegistrationServer_PogsClient struct {
 	Conn   *rpc.Conn
 }
 
+func NewRegistrationServer_PogsClient(client capnp.Client, conn *rpc.Conn) RegistrationServer_PogsClient {
+	return RegistrationServer_PogsClient{
+		Client: client,
+		Conn:   conn,
+	}
+}
+
 func (c RegistrationServer_PogsClient) Close() error {
 	c.Client.Close()
 	return c.Conn.Close()

tunnelrpc/registration_client.go

@@ -0,0 +1,76 @@
+package tunnelrpc
+
+import (
+	"context"
+	"io"
+	"net"
+	"time"
+
+	"github.com/google/uuid"
+	"zombiezen.com/go/capnproto2/rpc"
+
+	"github.com/cloudflare/cloudflared/tunnelrpc/pogs"
+)
+
+type RegistrationClient interface {
+	RegisterConnection(
+		ctx context.Context,
+		auth pogs.TunnelAuth,
+		tunnelID uuid.UUID,
+		options *pogs.ConnectionOptions,
+		connIndex uint8,
+		edgeAddress net.IP,
+	) (*pogs.ConnectionDetails, error)
+	SendLocalConfiguration(ctx context.Context, config []byte) error
+	GracefulShutdown(ctx context.Context, gracePeriod time.Duration)
+	Close()
+}
+
+type registrationClient struct {
+	client         pogs.RegistrationServer_PogsClient
+	transport      rpc.Transport
+	requestTimeout time.Duration
+}
+
+func NewRegistrationClient(ctx context.Context, stream io.ReadWriteCloser, requestTimeout time.Duration) RegistrationClient {
+	transport := SafeTransport(stream)
+	conn := rpc.NewConn(transport)
+	client := pogs.NewRegistrationServer_PogsClient(conn.Bootstrap(ctx), conn)
+	return &registrationClient{
+		client:         client,
+		transport:      transport,
+		requestTimeout: requestTimeout,
+	}
+}
+
+func (r *registrationClient) RegisterConnection(
+	ctx context.Context,
+	auth pogs.TunnelAuth,
+	tunnelID uuid.UUID,
+	options *pogs.ConnectionOptions,
+	connIndex uint8,
+	edgeAddress net.IP,
+) (*pogs.ConnectionDetails, error) {
+	ctx, cancel := context.WithTimeout(ctx, r.requestTimeout)
+	defer cancel()
+	return r.client.RegisterConnection(ctx, auth, tunnelID, connIndex, options)
+}
+
+func (r *registrationClient) SendLocalConfiguration(ctx context.Context, config []byte) error {
+	ctx, cancel := context.WithTimeout(ctx, r.requestTimeout)
+	defer cancel()
+	return r.client.SendLocalConfiguration(ctx, config)
+}
+
+func (r *registrationClient) GracefulShutdown(ctx context.Context, gracePeriod time.Duration) {
+	ctx, cancel := context.WithTimeout(ctx, gracePeriod)
+	defer cancel()
+	_ = r.client.UnregisterConnection(ctx)
+}
+
+func (r *registrationClient) Close() {
+	// Closing the client will also close the connection
+	_ = r.client.Close()
+	// Closing the transport also closes the stream
+	_ = r.transport.Close()
+}

tunnelrpc/registration_server.go

@@ -0,0 +1,40 @@
+package tunnelrpc
+
+import (
+	"context"
+	"io"
+
+	"zombiezen.com/go/capnproto2/rpc"
+
+	"github.com/cloudflare/cloudflared/tunnelrpc/pogs"
+)
+
+// RegistrationServer provides a handler interface for a client to provide methods to handle the different types of
+// requests that can be communicated by the stream.
+type RegistrationServer struct {
+	registrationServer pogs.RegistrationServer
+}
+
+func NewRegistrationServer(registrationServer pogs.RegistrationServer) *RegistrationServer {
+	return &RegistrationServer{
+		registrationServer: registrationServer,
+	}
+}
+
+// Serve listens for all RegistrationServer RPCs, including UnregisterConnection until the underlying connection
+// is terminated.
+func (s *RegistrationServer) Serve(ctx context.Context, stream io.ReadWriteCloser) error {
+	transport := SafeTransport(stream)
+	defer transport.Close()
+
+	main := pogs.RegistrationServer_ServerToClient(s.registrationServer)
+	rpcConn := rpc.NewConn(transport, rpc.MainInterface(main.Client))
+	defer rpcConn.Close()
+
+	select {
+	case <-rpcConn.Done():
+		return rpcConn.Wait()
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}