TUN-8415: Refactor capnp rpc into a single module

Combines the tunnelrpc and quic/schema capnp files into the same module. To help reduce future issues with capnp id generation, capnpids are provided in the capnp files from the existing capnp struct ids generated in the go files. Reduces the overall interface of the Capnp methods to the rest of the code by providing an interface that will handle the quic protocol selection. Introduces a new `rpc-timeout` config that will allow all of the SessionManager and ConfigurationManager RPC requests to have a timeout. The timeout for these values is set to 5 seconds as non of these operations for the managers should take a long time to complete. Removed the RPC-specific logger as it never provided good debugging value as the RPC method names were not visible in the logs.
2025-07-27 20:59:58 +00:00 · 2024-05-13 21:22:06 -07:00
parent 7d76ce2d24
commit eb2e4349e8
39 changed files with 1121 additions and 1028 deletions
--- a/tunnelrpc/pogs/auth_serialize.go
+++ b/tunnelrpc/pogs/auth_serialize.go
@@ -6,10 +6,10 @@ import (
 	"zombiezen.com/go/capnproto2/pogs"
 	"zombiezen.com/go/capnproto2/server"

-	"github.com/cloudflare/cloudflared/tunnelrpc"
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
 )

-func (i TunnelServer_PogsImpl) Authenticate(p tunnelrpc.TunnelServer_authenticate) error {
+func (i TunnelServer_PogsImpl) Authenticate(p proto.TunnelServer_authenticate) error {
 	originCert, err := p.Params.OriginCert()
 	if err != nil {
 		return err
@@ -39,13 +39,13 @@ func (i TunnelServer_PogsImpl) Authenticate(p tunnelrpc.TunnelServer_authenticat
 	return MarshalAuthenticateResponse(result, resp)
 }

-func MarshalAuthenticateResponse(s tunnelrpc.AuthenticateResponse, p *AuthenticateResponse) error {
-	return pogs.Insert(tunnelrpc.AuthenticateResponse_TypeID, s.Struct, p)
+func MarshalAuthenticateResponse(s proto.AuthenticateResponse, p *AuthenticateResponse) error {
+	return pogs.Insert(proto.AuthenticateResponse_TypeID, s.Struct, p)
 }

 func (c TunnelServer_PogsClient) Authenticate(ctx context.Context, originCert []byte, hostname string, options *RegistrationOptions) (*AuthenticateResponse, error) {
-	client := tunnelrpc.TunnelServer{Client: c.Client}
-	promise := client.Authenticate(ctx, func(p tunnelrpc.TunnelServer_authenticate_Params) error {
+	client := proto.TunnelServer{Client: c.Client}
+	promise := client.Authenticate(ctx, func(p proto.TunnelServer_authenticate_Params) error {
 		err := p.SetOriginCert(originCert)
 		if err != nil {
 			return err
@@ -71,8 +71,8 @@ func (c TunnelServer_PogsClient) Authenticate(ctx context.Context, originCert []
 	return UnmarshalAuthenticateResponse(retval)
 }

-func UnmarshalAuthenticateResponse(s tunnelrpc.AuthenticateResponse) (*AuthenticateResponse, error) {
+func UnmarshalAuthenticateResponse(s proto.AuthenticateResponse) (*AuthenticateResponse, error) {
 	p := new(AuthenticateResponse)
-	err := pogs.Extract(p, tunnelrpc.AuthenticateResponse_TypeID, s.Struct)
+	err := pogs.Extract(p, proto.AuthenticateResponse_TypeID, s.Struct)
 	return p, err
 }
--- a/tunnelrpc/pogs/auth_test.go
+++ b/tunnelrpc/pogs/auth_test.go
@@ -9,7 +9,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	capnp "zombiezen.com/go/capnproto2"

-	"github.com/cloudflare/cloudflared/tunnelrpc"
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
 )

 // Ensure the AuthOutcome sum is correct
@@ -119,7 +119,7 @@ func TestSerializeAuthenticationResponse(t *testing.T) {
 	for i, testCase := range tests {
 		_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
 		assert.NoError(t, err)
-		capnpEntity, err := tunnelrpc.NewAuthenticateResponse(seg)
+		capnpEntity, err := proto.NewAuthenticateResponse(seg)
 		if !assert.NoError(t, err) {
 			t.Fatal("Couldn't initialize a new message")
 		}
--- a/tunnelrpc/pogs/cloudflaredrpc.go
+++ b/tunnelrpc/pogs/cloudflaredrpc.go
@@ -1,9 +1,10 @@
 package pogs

 import (
-	"github.com/cloudflare/cloudflared/tunnelrpc"
 	capnp "zombiezen.com/go/capnproto2"
 	"zombiezen.com/go/capnproto2/rpc"
+
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
 )

 type CloudflaredServer interface {
@@ -16,8 +17,8 @@ type CloudflaredServer_PogsImpl struct {
 	ConfigurationManager_PogsImpl
 }

-func CloudflaredServer_ServerToClient(s SessionManager, c ConfigurationManager) tunnelrpc.CloudflaredServer {
-	return tunnelrpc.CloudflaredServer_ServerToClient(CloudflaredServer_PogsImpl{
+func CloudflaredServer_ServerToClient(s SessionManager, c ConfigurationManager) proto.CloudflaredServer {
+	return proto.CloudflaredServer_ServerToClient(CloudflaredServer_PogsImpl{
 		SessionManager_PogsImpl:       SessionManager_PogsImpl{s},
 		ConfigurationManager_PogsImpl: ConfigurationManager_PogsImpl{c},
 	})
--- a/tunnelrpc/pogs/configurationrpc.go
+++ b/tunnelrpc/pogs/configurationrpc.go
@@ -4,10 +4,11 @@ import (
 	"context"
 	"fmt"

-	"github.com/cloudflare/cloudflared/tunnelrpc"
 	capnp "zombiezen.com/go/capnproto2"
 	"zombiezen.com/go/capnproto2/rpc"
 	"zombiezen.com/go/capnproto2/server"
+
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
 )

 type ConfigurationManager interface {
@@ -18,11 +19,11 @@ type ConfigurationManager_PogsImpl struct {
 	impl ConfigurationManager
 }

-func ConfigurationManager_ServerToClient(c ConfigurationManager) tunnelrpc.ConfigurationManager {
-	return tunnelrpc.ConfigurationManager_ServerToClient(ConfigurationManager_PogsImpl{c})
+func ConfigurationManager_ServerToClient(c ConfigurationManager) proto.ConfigurationManager {
+	return proto.ConfigurationManager_ServerToClient(ConfigurationManager_PogsImpl{c})
 }

-func (i ConfigurationManager_PogsImpl) UpdateConfiguration(p tunnelrpc.ConfigurationManager_updateConfiguration) error {
+func (i ConfigurationManager_PogsImpl) UpdateConfiguration(p proto.ConfigurationManager_updateConfiguration) error {
 	server.Ack(p.Options)

 	version := p.Params.Version()
@@ -51,8 +52,8 @@ func (c ConfigurationManager_PogsClient) Close() error {
 }

 func (c ConfigurationManager_PogsClient) UpdateConfiguration(ctx context.Context, version int32, config []byte) (*UpdateConfigurationResponse, error) {
-	client := tunnelrpc.ConfigurationManager{Client: c.Client}
-	promise := client.UpdateConfiguration(ctx, func(p tunnelrpc.ConfigurationManager_updateConfiguration_Params) error {
+	client := proto.ConfigurationManager{Client: c.Client}
+	promise := client.UpdateConfiguration(ctx, func(p proto.ConfigurationManager_updateConfiguration_Params) error {
 		p.SetVersion(version)
 		return p.SetConfig(config)
 	})
@@ -74,7 +75,7 @@ type UpdateConfigurationResponse struct {
 	Err                error `json:"err"`
 }

-func (p *UpdateConfigurationResponse) Marshal(s tunnelrpc.UpdateConfigurationResponse) error {
+func (p *UpdateConfigurationResponse) Marshal(s proto.UpdateConfigurationResponse) error {
 	s.SetLatestAppliedVersion(p.LastAppliedVersion)
 	if p.Err != nil {
 		return s.SetErr(p.Err.Error())
@@ -82,7 +83,7 @@ func (p *UpdateConfigurationResponse) Marshal(s tunnelrpc.UpdateConfigurationRes
 	return nil
 }

-func (p *UpdateConfigurationResponse) Unmarshal(s tunnelrpc.UpdateConfigurationResponse) error {
+func (p *UpdateConfigurationResponse) Unmarshal(s proto.UpdateConfigurationResponse) error {
 	p.LastAppliedVersion = s.LatestAppliedVersion()
 	respErr, err := s.Err()
 	if err != nil {
--- a/tunnelrpc/pogs/connectionrpc.go
+++ b/tunnelrpc/pogs/connectionrpc.go
@@ -12,7 +12,7 @@ import (
 	"zombiezen.com/go/capnproto2/rpc"
 	"zombiezen.com/go/capnproto2/server"

-	"github.com/cloudflare/cloudflared/tunnelrpc"
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
 )

 type RegistrationServer interface {
@@ -25,11 +25,11 @@ type RegistrationServer_PogsImpl struct {
 	impl RegistrationServer
 }

-func RegistrationServer_ServerToClient(s RegistrationServer) tunnelrpc.RegistrationServer {
-	return tunnelrpc.RegistrationServer_ServerToClient(RegistrationServer_PogsImpl{s})
+func RegistrationServer_ServerToClient(s RegistrationServer) proto.RegistrationServer {
+	return proto.RegistrationServer_ServerToClient(RegistrationServer_PogsImpl{s})
 }

-func (i RegistrationServer_PogsImpl) RegisterConnection(p tunnelrpc.RegistrationServer_registerConnection) error {
+func (i RegistrationServer_PogsImpl) RegisterConnection(p proto.RegistrationServer_registerConnection) error {
 	server.Ack(p.Options)

 	auth, err := p.Params.Auth()
@@ -82,14 +82,14 @@ func (i RegistrationServer_PogsImpl) RegisterConnection(p tunnelrpc.Registration
 	}
 }

-func (i RegistrationServer_PogsImpl) UnregisterConnection(p tunnelrpc.RegistrationServer_unregisterConnection) error {
+func (i RegistrationServer_PogsImpl) UnregisterConnection(p proto.RegistrationServer_unregisterConnection) error {
 	server.Ack(p.Options)

 	i.impl.UnregisterConnection(p.Ctx)
 	return nil
 }

-func (i RegistrationServer_PogsImpl) UpdateLocalConfiguration(c tunnelrpc.RegistrationServer_updateLocalConfiguration) error {
+func (i RegistrationServer_PogsImpl) UpdateLocalConfiguration(c proto.RegistrationServer_updateLocalConfiguration) error {
 	server.Ack(c.Options)

 	configBytes, err := c.Params.Config()
@@ -111,8 +111,8 @@ func (c RegistrationServer_PogsClient) Close() error {
 }

 func (c RegistrationServer_PogsClient) RegisterConnection(ctx context.Context, auth TunnelAuth, tunnelID uuid.UUID, connIndex byte, options *ConnectionOptions) (*ConnectionDetails, error) {
-	client := tunnelrpc.TunnelServer{Client: c.Client}
-	promise := client.RegisterConnection(ctx, func(p tunnelrpc.RegistrationServer_registerConnection_Params) error {
+	client := proto.TunnelServer{Client: c.Client}
+	promise := client.RegisterConnection(ctx, func(p proto.RegistrationServer_registerConnection_Params) error {
 		tunnelAuth, err := p.NewAuth()
 		if err != nil {
 			return err
@@ -145,7 +145,7 @@ func (c RegistrationServer_PogsClient) RegisterConnection(ctx context.Context, a
 	}
 	result := response.Result()
 	switch result.Which() {
-	case tunnelrpc.ConnectionResponse_result_Which_error:
+	case proto.ConnectionResponse_result_Which_error:
 		resultError, err := result.Error()
 		if err != nil {
 			return nil, wrapRPCError(err)
@@ -160,7 +160,7 @@ func (c RegistrationServer_PogsClient) RegisterConnection(ctx context.Context, a
 		}
 		return nil, err

-	case tunnelrpc.ConnectionResponse_result_Which_connectionDetails:
+	case proto.ConnectionResponse_result_Which_connectionDetails:
 		connDetails, err := result.ConnectionDetails()
 		if err != nil {
 			return nil, wrapRPCError(err)
@@ -176,8 +176,8 @@ func (c RegistrationServer_PogsClient) RegisterConnection(ctx context.Context, a
 }

 func (c RegistrationServer_PogsClient) SendLocalConfiguration(ctx context.Context, config []byte) error {
-	client := tunnelrpc.TunnelServer{Client: c.Client}
-	promise := client.UpdateLocalConfiguration(ctx, func(p tunnelrpc.RegistrationServer_updateLocalConfiguration_Params) error {
+	client := proto.TunnelServer{Client: c.Client}
+	promise := client.UpdateLocalConfiguration(ctx, func(p proto.RegistrationServer_updateLocalConfiguration_Params) error {
 		if err := p.SetConfig(config); err != nil {
 			return err
 		}
@@ -194,8 +194,8 @@ func (c RegistrationServer_PogsClient) SendLocalConfiguration(ctx context.Contex
 }

 func (c RegistrationServer_PogsClient) UnregisterConnection(ctx context.Context) error {
-	client := tunnelrpc.TunnelServer{Client: c.Client}
-	promise := client.UnregisterConnection(ctx, func(p tunnelrpc.RegistrationServer_unregisterConnection_Params) error {
+	client := proto.TunnelServer{Client: c.Client}
+	promise := client.UnregisterConnection(ctx, func(p proto.RegistrationServer_unregisterConnection_Params) error {
 		return nil
 	})
 	_, err := promise.Struct()
@@ -225,20 +225,20 @@ type TunnelAuth struct {
 	TunnelSecret []byte
 }

-func (p *ConnectionOptions) MarshalCapnproto(s tunnelrpc.ConnectionOptions) error {
-	return pogs.Insert(tunnelrpc.ConnectionOptions_TypeID, s.Struct, p)
+func (p *ConnectionOptions) MarshalCapnproto(s proto.ConnectionOptions) error {
+	return pogs.Insert(proto.ConnectionOptions_TypeID, s.Struct, p)
 }

-func (p *ConnectionOptions) UnmarshalCapnproto(s tunnelrpc.ConnectionOptions) error {
-	return pogs.Extract(p, tunnelrpc.ConnectionOptions_TypeID, s.Struct)
+func (p *ConnectionOptions) UnmarshalCapnproto(s proto.ConnectionOptions) error {
+	return pogs.Extract(p, proto.ConnectionOptions_TypeID, s.Struct)
 }

-func (a *TunnelAuth) MarshalCapnproto(s tunnelrpc.TunnelAuth) error {
-	return pogs.Insert(tunnelrpc.TunnelAuth_TypeID, s.Struct, a)
+func (a *TunnelAuth) MarshalCapnproto(s proto.TunnelAuth) error {
+	return pogs.Insert(proto.TunnelAuth_TypeID, s.Struct, a)
 }

-func (a *TunnelAuth) UnmarshalCapnproto(s tunnelrpc.TunnelAuth) error {
-	return pogs.Extract(a, tunnelrpc.TunnelAuth_TypeID, s.Struct)
+func (a *TunnelAuth) UnmarshalCapnproto(s proto.TunnelAuth) error {
+	return pogs.Extract(a, proto.TunnelAuth_TypeID, s.Struct)
 }

 type ConnectionDetails struct {
@@ -247,7 +247,7 @@ type ConnectionDetails struct {
 	TunnelIsRemotelyManaged bool
 }

-func (details *ConnectionDetails) MarshalCapnproto(s tunnelrpc.ConnectionDetails) error {
+func (details *ConnectionDetails) MarshalCapnproto(s proto.ConnectionDetails) error {
 	if err := s.SetUuid(details.UUID[:]); err != nil {
 		return err
 	}
@@ -259,7 +259,7 @@ func (details *ConnectionDetails) MarshalCapnproto(s tunnelrpc.ConnectionDetails
 	return nil
 }

-func (details *ConnectionDetails) UnmarshalCapnproto(s tunnelrpc.ConnectionDetails) error {
+func (details *ConnectionDetails) UnmarshalCapnproto(s proto.ConnectionDetails) error {
 	uuidBytes, err := s.Uuid()
 	if err != nil {
 		return err
@@ -277,7 +277,7 @@ func (details *ConnectionDetails) UnmarshalCapnproto(s tunnelrpc.ConnectionDetai
 	return err
 }

-func MarshalError(s tunnelrpc.ConnectionError, err error) error {
+func MarshalError(s proto.ConnectionError, err error) error {
 	if err := s.SetCause(err.Error()); err != nil {
 		return err
 	}
--- a/tunnelrpc/pogs/connectionrpc_test.go
+++ b/tunnelrpc/pogs/connectionrpc_test.go
@@ -13,7 +13,7 @@ import (
 	capnp "zombiezen.com/go/capnproto2"
 	"zombiezen.com/go/capnproto2/rpc"

-	"github.com/cloudflare/cloudflared/tunnelrpc"
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
 )

 const testAccountTag = "abc123"
@@ -34,7 +34,7 @@ func TestMarshalConnectionOptions(t *testing.T) {

 	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
 	require.NoError(t, err)
-	capnpOpts, err := tunnelrpc.NewConnectionOptions(seg)
+	capnpOpts, err := proto.NewConnectionOptions(seg)
 	require.NoError(t, err)

 	err = orig.MarshalCapnproto(capnpOpts)
--- a/tunnelrpc/pogs/quic_metadata_protocol.go
+++ b/tunnelrpc/pogs/quic_metadata_protocol.go
@@ -0,0 +1,111 @@
+package pogs
+
+import (
+	"fmt"
+
+	capnp "zombiezen.com/go/capnproto2"
+	"zombiezen.com/go/capnproto2/pogs"
+
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
+)
+
+// ConnectionType indicates the type of underlying connection proxied within the QUIC stream.
+type ConnectionType uint16
+
+const (
+	ConnectionTypeHTTP ConnectionType = iota
+	ConnectionTypeWebsocket
+	ConnectionTypeTCP
+)
+
+func (c ConnectionType) String() string {
+	switch c {
+	case ConnectionTypeHTTP:
+		return "http"
+	case ConnectionTypeWebsocket:
+		return "ws"
+	case ConnectionTypeTCP:
+		return "tcp"
+	}
+	panic(fmt.Sprintf("invalid ConnectionType: %d", c))
+}
+
+// ConnectRequest is the representation of metadata sent at the start of a QUIC application handshake.
+type ConnectRequest struct {
+	Dest     string         `capnp:"dest"`
+	Type     ConnectionType `capnp:"type"`
+	Metadata []Metadata     `capnp:"metadata"`
+}
+
+// Metadata is a representation of key value based data sent via RequestMeta.
+type Metadata struct {
+	Key string `capnp:"key"`
+	Val string `capnp:"val"`
+}
+
+// MetadataMap returns a map format of []Metadata.
+func (r *ConnectRequest) MetadataMap() map[string]string {
+	metadataMap := make(map[string]string)
+	for _, metadata := range r.Metadata {
+		metadataMap[metadata.Key] = metadata.Val
+	}
+	return metadataMap
+}
+
+func (r *ConnectRequest) FromPogs(msg *capnp.Message) error {
+	metadata, err := proto.ReadRootConnectRequest(msg)
+	if err != nil {
+		return err
+	}
+	return pogs.Extract(r, proto.ConnectRequest_TypeID, metadata.Struct)
+}
+
+func (r *ConnectRequest) ToPogs() (*capnp.Message, error) {
+	msg, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
+	if err != nil {
+		return nil, err
+	}
+
+	root, err := proto.NewRootConnectRequest(seg)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := pogs.Insert(proto.ConnectRequest_TypeID, root.Struct, r); err != nil {
+		return nil, err
+	}
+
+	return msg, nil
+}
+
+// ConnectResponse is a representation of metadata sent as a response to a QUIC application handshake.
+type ConnectResponse struct {
+	Error    string     `capnp:"error"`
+	Metadata []Metadata `capnp:"metadata"`
+}
+
+func (r *ConnectResponse) FromPogs(msg *capnp.Message) error {
+	metadata, err := proto.ReadRootConnectResponse(msg)
+	if err != nil {
+		return err
+	}
+	return pogs.Extract(r, proto.ConnectResponse_TypeID, metadata.Struct)
+}
+
+func (r *ConnectResponse) ToPogs() (*capnp.Message, error) {
+	msg, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
+	if err != nil {
+		return nil, err
+	}
+
+	root, err := proto.NewRootConnectResponse(seg)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := pogs.Insert(proto.ConnectResponse_TypeID, root.Struct, r); err != nil {
+		return nil, err
+	}
+
+	return msg, nil
+}
--- a/tunnelrpc/pogs/reconnect_tunnel.go
+++ b/tunnelrpc/pogs/reconnect_tunnel.go
@@ -5,10 +5,10 @@ import (

 	"zombiezen.com/go/capnproto2/server"

-	"github.com/cloudflare/cloudflared/tunnelrpc"
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
 )

-func (i TunnelServer_PogsImpl) ReconnectTunnel(p tunnelrpc.TunnelServer_reconnectTunnel) error {
+func (i TunnelServer_PogsImpl) ReconnectTunnel(p proto.TunnelServer_reconnectTunnel) error {
 	jwt, err := p.Params.Jwt()
 	if err != nil {
 		return err
@@ -53,8 +53,8 @@ func (c TunnelServer_PogsClient) ReconnectTunnel(
 	hostname string,
 	options *RegistrationOptions,
 ) *TunnelRegistration {
-	client := tunnelrpc.TunnelServer{Client: c.Client}
-	promise := client.ReconnectTunnel(ctx, func(p tunnelrpc.TunnelServer_reconnectTunnel_Params) error {
+	client := proto.TunnelServer{Client: c.Client}
+	promise := client.ReconnectTunnel(ctx, func(p proto.TunnelServer_reconnectTunnel_Params) error {
 		err := p.SetJwt(jwt)
 		if err != nil {
 			return err
--- a/tunnelrpc/pogs/sessionrpc.go
+++ b/tunnelrpc/pogs/sessionrpc.go
@@ -6,11 +6,12 @@ import (
 	"net"
 	"time"

-	"github.com/cloudflare/cloudflared/tunnelrpc"
 	"github.com/google/uuid"
 	capnp "zombiezen.com/go/capnproto2"
 	"zombiezen.com/go/capnproto2/rpc"
 	"zombiezen.com/go/capnproto2/server"
+
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
 )

 type SessionManager interface {
@@ -26,11 +27,11 @@ type SessionManager_PogsImpl struct {
 	impl SessionManager
 }

-func SessionManager_ServerToClient(s SessionManager) tunnelrpc.SessionManager {
-	return tunnelrpc.SessionManager_ServerToClient(SessionManager_PogsImpl{s})
+func SessionManager_ServerToClient(s SessionManager) proto.SessionManager {
+	return proto.SessionManager_ServerToClient(SessionManager_PogsImpl{s})
 }

-func (i SessionManager_PogsImpl) RegisterUdpSession(p tunnelrpc.SessionManager_registerUdpSession) error {
+func (i SessionManager_PogsImpl) RegisterUdpSession(p proto.SessionManager_registerUdpSession) error {
 	server.Ack(p.Options)

 	sessionIDRaw, err := p.Params.SessionId()
@@ -76,7 +77,7 @@ func (i SessionManager_PogsImpl) RegisterUdpSession(p tunnelrpc.SessionManager_r
 	return resp.Marshal(result)
 }

-func (i SessionManager_PogsImpl) UnregisterUdpSession(p tunnelrpc.SessionManager_unregisterUdpSession) error {
+func (i SessionManager_PogsImpl) UnregisterUdpSession(p proto.SessionManager_unregisterUdpSession) error {
 	server.Ack(p.Options)

 	sessionIDRaw, err := p.Params.SessionId()
@@ -101,7 +102,7 @@ type RegisterUdpSessionResponse struct {
 	Spans []byte // Spans in protobuf format
 }

-func (p *RegisterUdpSessionResponse) Marshal(s tunnelrpc.RegisterUdpSessionResponse) error {
+func (p *RegisterUdpSessionResponse) Marshal(s proto.RegisterUdpSessionResponse) error {
 	if p.Err != nil {
 		return s.SetErr(p.Err.Error())
 	}
@@ -111,7 +112,7 @@ func (p *RegisterUdpSessionResponse) Marshal(s tunnelrpc.RegisterUdpSessionRespo
 	return nil
 }

-func (p *RegisterUdpSessionResponse) Unmarshal(s tunnelrpc.RegisterUdpSessionResponse) error {
+func (p *RegisterUdpSessionResponse) Unmarshal(s proto.RegisterUdpSessionResponse) error {
 	respErr, err := s.Err()
 	if err != nil {
 		return err
@@ -131,14 +132,21 @@ type SessionManager_PogsClient struct {
 	Conn   *rpc.Conn
 }

+func NewSessionManager_PogsClient(client capnp.Client, conn *rpc.Conn) SessionManager_PogsClient {
+	return SessionManager_PogsClient{
+		Client: client,
+		Conn:   conn,
+	}
+}
+
 func (c SessionManager_PogsClient) Close() error {
 	c.Client.Close()
 	return c.Conn.Close()
 }

 func (c SessionManager_PogsClient) RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16, closeAfterIdleHint time.Duration, traceContext string) (*RegisterUdpSessionResponse, error) {
-	client := tunnelrpc.SessionManager{Client: c.Client}
-	promise := client.RegisterUdpSession(ctx, func(p tunnelrpc.SessionManager_registerUdpSession_Params) error {
+	client := proto.SessionManager{Client: c.Client}
+	promise := client.RegisterUdpSession(ctx, func(p proto.SessionManager_registerUdpSession_Params) error {
 		if err := p.SetSessionId(sessionID[:]); err != nil {
 			return err
 		}
@@ -164,8 +172,8 @@ func (c SessionManager_PogsClient) RegisterUdpSession(ctx context.Context, sessi
 }

 func (c SessionManager_PogsClient) UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID, message string) error {
-	client := tunnelrpc.SessionManager{Client: c.Client}
-	promise := client.UnregisterUdpSession(ctx, func(p tunnelrpc.SessionManager_unregisterUdpSession_Params) error {
+	client := proto.SessionManager{Client: c.Client}
+	promise := client.UnregisterUdpSession(ctx, func(p proto.SessionManager_unregisterUdpSession_Params) error {
 		if err := p.SetSessionId(sessionID[:]); err != nil {
 			return err
 		}
--- a/tunnelrpc/pogs/tunnelrpc.go
+++ b/tunnelrpc/pogs/tunnelrpc.go
@@ -9,7 +9,7 @@ import (
 	"zombiezen.com/go/capnproto2/rpc"
 	"zombiezen.com/go/capnproto2/server"

-	"github.com/cloudflare/cloudflared/tunnelrpc"
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
 )

 const (
@@ -22,13 +22,13 @@ type Authentication struct {
 	OriginCAKey string
 }

-func MarshalAuthentication(s tunnelrpc.Authentication, p *Authentication) error {
-	return pogs.Insert(tunnelrpc.Authentication_TypeID, s.Struct, p)
+func MarshalAuthentication(s proto.Authentication, p *Authentication) error {
+	return pogs.Insert(proto.Authentication_TypeID, s.Struct, p)
 }

-func UnmarshalAuthentication(s tunnelrpc.Authentication) (*Authentication, error) {
+func UnmarshalAuthentication(s proto.Authentication) (*Authentication, error) {
 	p := new(Authentication)
-	err := pogs.Extract(p, tunnelrpc.Authentication_TypeID, s.Struct)
+	err := pogs.Extract(p, proto.Authentication_TypeID, s.Struct)
 	return p, err
 }

@@ -144,13 +144,13 @@ func (*RetryableRegistrationError) IsPermanent() bool {
 	return false
 }

-func MarshalTunnelRegistration(s tunnelrpc.TunnelRegistration, p *TunnelRegistration) error {
-	return pogs.Insert(tunnelrpc.TunnelRegistration_TypeID, s.Struct, p)
+func MarshalTunnelRegistration(s proto.TunnelRegistration, p *TunnelRegistration) error {
+	return pogs.Insert(proto.TunnelRegistration_TypeID, s.Struct, p)
 }

-func UnmarshalTunnelRegistration(s tunnelrpc.TunnelRegistration) (*TunnelRegistration, error) {
+func UnmarshalTunnelRegistration(s proto.TunnelRegistration) (*TunnelRegistration, error) {
 	p := new(TunnelRegistration)
-	err := pogs.Extract(p, tunnelrpc.TunnelRegistration_TypeID, s.Struct)
+	err := pogs.Extract(p, proto.TunnelRegistration_TypeID, s.Struct)
 	return p, err
 }

@@ -158,7 +158,7 @@ type RegistrationOptions struct {
 	ClientID             string `capnp:"clientId"`
 	Version              string
 	OS                   string `capnp:"os"`
-	ExistingTunnelPolicy tunnelrpc.ExistingTunnelPolicy
+	ExistingTunnelPolicy proto.ExistingTunnelPolicy
 	PoolName             string `capnp:"poolName"`
 	Tags                 []Tag
 	ConnectionID         uint8  `capnp:"connectionId"`
@@ -171,13 +171,13 @@ type RegistrationOptions struct {
 	Features             []string
 }

-func MarshalRegistrationOptions(s tunnelrpc.RegistrationOptions, p *RegistrationOptions) error {
-	return pogs.Insert(tunnelrpc.RegistrationOptions_TypeID, s.Struct, p)
+func MarshalRegistrationOptions(s proto.RegistrationOptions, p *RegistrationOptions) error {
+	return pogs.Insert(proto.RegistrationOptions_TypeID, s.Struct, p)
 }

-func UnmarshalRegistrationOptions(s tunnelrpc.RegistrationOptions) (*RegistrationOptions, error) {
+func UnmarshalRegistrationOptions(s proto.RegistrationOptions) (*RegistrationOptions, error) {
 	p := new(RegistrationOptions)
-	err := pogs.Extract(p, tunnelrpc.RegistrationOptions_TypeID, s.Struct)
+	err := pogs.Extract(p, proto.RegistrationOptions_TypeID, s.Struct)
 	return p, err
 }

@@ -190,13 +190,13 @@ type ServerInfo struct {
 	LocationName string
 }

-func MarshalServerInfo(s tunnelrpc.ServerInfo, p *ServerInfo) error {
-	return pogs.Insert(tunnelrpc.ServerInfo_TypeID, s.Struct, p)
+func MarshalServerInfo(s proto.ServerInfo, p *ServerInfo) error {
+	return pogs.Insert(proto.ServerInfo_TypeID, s.Struct, p)
 }

-func UnmarshalServerInfo(s tunnelrpc.ServerInfo) (*ServerInfo, error) {
+func UnmarshalServerInfo(s proto.ServerInfo) (*ServerInfo, error) {
 	p := new(ServerInfo)
-	err := pogs.Extract(p, tunnelrpc.ServerInfo_TypeID, s.Struct)
+	err := pogs.Extract(p, proto.ServerInfo_TypeID, s.Struct)
 	return p, err
 }

@@ -209,8 +209,8 @@ type TunnelServer interface {
 	ReconnectTunnel(ctx context.Context, jwt, eventDigest, connDigest []byte, hostname string, options *RegistrationOptions) (*TunnelRegistration, error)
 }

-func TunnelServer_ServerToClient(s TunnelServer) tunnelrpc.TunnelServer {
-	return tunnelrpc.TunnelServer_ServerToClient(TunnelServer_PogsImpl{RegistrationServer_PogsImpl{s}, s})
+func TunnelServer_ServerToClient(s TunnelServer) proto.TunnelServer {
+	return proto.TunnelServer_ServerToClient(TunnelServer_PogsImpl{RegistrationServer_PogsImpl{s}, s})
 }

 type TunnelServer_PogsImpl struct {
@@ -218,7 +218,7 @@ type TunnelServer_PogsImpl struct {
 	impl TunnelServer
 }

-func (i TunnelServer_PogsImpl) RegisterTunnel(p tunnelrpc.TunnelServer_registerTunnel) error {
+func (i TunnelServer_PogsImpl) RegisterTunnel(p proto.TunnelServer_registerTunnel) error {
 	originCert, err := p.Params.OriginCert()
 	if err != nil {
 		return err
@@ -245,7 +245,7 @@ func (i TunnelServer_PogsImpl) RegisterTunnel(p tunnelrpc.TunnelServer_registerT
 	return MarshalTunnelRegistration(result, registration)
 }

-func (i TunnelServer_PogsImpl) GetServerInfo(p tunnelrpc.TunnelServer_getServerInfo) error {
+func (i TunnelServer_PogsImpl) GetServerInfo(p proto.TunnelServer_getServerInfo) error {
 	server.Ack(p.Options)
 	serverInfo, err := i.impl.GetServerInfo(p.Ctx)
 	if err != nil {
@@ -258,13 +258,13 @@ func (i TunnelServer_PogsImpl) GetServerInfo(p tunnelrpc.TunnelServer_getServerI
 	return MarshalServerInfo(result, serverInfo)
 }

-func (i TunnelServer_PogsImpl) UnregisterTunnel(p tunnelrpc.TunnelServer_unregisterTunnel) error {
+func (i TunnelServer_PogsImpl) UnregisterTunnel(p proto.TunnelServer_unregisterTunnel) error {
 	gracePeriodNanoSec := p.Params.GracePeriodNanoSec()
 	server.Ack(p.Options)
 	return i.impl.UnregisterTunnel(p.Ctx, gracePeriodNanoSec)
 }

-func (i TunnelServer_PogsImpl) ObsoleteDeclarativeTunnelConnect(p tunnelrpc.TunnelServer_obsoleteDeclarativeTunnelConnect) error {
+func (i TunnelServer_PogsImpl) ObsoleteDeclarativeTunnelConnect(p proto.TunnelServer_obsoleteDeclarativeTunnelConnect) error {
 	return fmt.Errorf("RPC to create declarative tunnel connection has been deprecated")
 }

@@ -280,8 +280,8 @@ func (c TunnelServer_PogsClient) Close() error {
 }

 func (c TunnelServer_PogsClient) RegisterTunnel(ctx context.Context, originCert []byte, hostname string, options *RegistrationOptions) *TunnelRegistration {
-	client := tunnelrpc.TunnelServer{Client: c.Client}
-	promise := client.RegisterTunnel(ctx, func(p tunnelrpc.TunnelServer_registerTunnel_Params) error {
+	client := proto.TunnelServer{Client: c.Client}
+	promise := client.RegisterTunnel(ctx, func(p proto.TunnelServer_registerTunnel_Params) error {
 		err := p.SetOriginCert(originCert)
 		if err != nil {
 			return err
@@ -312,8 +312,8 @@ func (c TunnelServer_PogsClient) RegisterTunnel(ctx context.Context, originCert
 }

 func (c TunnelServer_PogsClient) GetServerInfo(ctx context.Context) (*ServerInfo, error) {
-	client := tunnelrpc.TunnelServer{Client: c.Client}
-	promise := client.GetServerInfo(ctx, func(p tunnelrpc.TunnelServer_getServerInfo_Params) error {
+	client := proto.TunnelServer{Client: c.Client}
+	promise := client.GetServerInfo(ctx, func(p proto.TunnelServer_getServerInfo_Params) error {
 		return nil
 	})
 	retval, err := promise.Result().Struct()
@@ -324,8 +324,8 @@ func (c TunnelServer_PogsClient) GetServerInfo(ctx context.Context) (*ServerInfo
 }

 func (c TunnelServer_PogsClient) UnregisterTunnel(ctx context.Context, gracePeriodNanoSec int64) error {
-	client := tunnelrpc.TunnelServer{Client: c.Client}
-	promise := client.UnregisterTunnel(ctx, func(p tunnelrpc.TunnelServer_unregisterTunnel_Params) error {
+	client := proto.TunnelServer{Client: c.Client}
+	promise := client.UnregisterTunnel(ctx, func(p proto.TunnelServer_unregisterTunnel_Params) error {
 		p.SetGracePeriodNanoSec(gracePeriodNanoSec)
 		return nil
 	})
--- a/tunnelrpc/pogs/tunnelrpc_test.go
+++ b/tunnelrpc/pogs/tunnelrpc_test.go
@@ -7,7 +7,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	capnp "zombiezen.com/go/capnproto2"

-	"github.com/cloudflare/cloudflared/tunnelrpc"
+	"github.com/cloudflare/cloudflared/tunnelrpc/proto"
 )

 const (
@@ -39,7 +39,7 @@ func TestTunnelRegistration(t *testing.T) {
 	for i, testCase := range testCases {
 		_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
 		assert.NoError(t, err)
-		capnpEntity, err := tunnelrpc.NewTunnelRegistration(seg)
+		capnpEntity, err := proto.NewTunnelRegistration(seg)
 		if !assert.NoError(t, err) {
 			t.Fatal("Couldn't initialize a new message")
 		}