hoelee/cloudflared
1
0
Fork 0
mirror of https://github.com/cloudflare/cloudflared.git synced 2025-07-27 14:29:58 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

TUN-3427: Define a struct that only implements RegistrationServer in tunnelpogs

Browse Source
This commit is contained in:
cthuang
2020-09-28 10:10:30 +01:00
parent 8e8513e325
commit 2c9b7361b7
9 changed files with 242 additions and 201 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
connection/rpc.go
View File

@@ -2,48 +2,40 @@ package connection
import (
"context"
"fmt"
"time"
"io"
rpc "zombiezen.com/go/capnproto2/rpc"
"github.com/cloudflare/cloudflared/h2mux"
"github.com/cloudflare/cloudflared/logger"
"github.com/cloudflare/cloudflared/tunnelrpc"
tunnelpogs "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
)
// NewRPCClient creates and returns a new RPC client, which will communicate
// NewTunnelRPCClient creates and returns a new RPC client, which will communicate
// using a stream on the given muxer
func NewRPCClient(
func NewTunnelRPCClient(
ctx context.Context,
muxer *h2mux.Muxer,
stream io.ReadWriteCloser,
logger logger.Service,
openStreamTimeout time.Duration,
) (client tunnelpogs.TunnelServer_PogsClient, err error) {
openStreamCtx, openStreamCancel := context.WithTimeout(ctx, openStreamTimeout)
defer openStreamCancel()
stream, err := muxer.OpenRPCStream(openStreamCtx)
if err != nil {
return
}
if !isRPCStreamResponse(stream.Headers) {
stream.Close()
err = fmt.Errorf("rpc: bad response headers: %v", stream.Headers)
return
}
conn := rpc.NewConn(
tunnelrpc.NewTransportLogger(logger, rpc.StreamTransport(stream)),
tunnelrpc.ConnLog(logger),
)
client = tunnelpogs.TunnelServer_PogsClient{Client: conn.Bootstrap(ctx), Conn: conn}
registrationClient := tunnelpogs.RegistrationServer_PogsClient{Client: conn.Bootstrap(ctx), Conn: conn}
client = tunnelpogs.TunnelServer_PogsClient{RegistrationServer_PogsClient: registrationClient, Client: conn.Bootstrap(ctx), Conn: conn}
return client, nil
}
func isRPCStreamResponse(headers []h2mux.Header) bool {
return len(headers) == 1 &&
headers[0].Name == ":status" &&
headers[0].Value == "200"
func NewRegistrationRPCClient(
ctx context.Context,
stream io.ReadWriteCloser,
logger logger.Service,
) (client tunnelpogs.RegistrationServer_PogsClient, err error) {
conn := rpc.NewConn(
tunnelrpc.NewTransportLogger(logger, rpc.StreamTransport(stream)),
tunnelrpc.ConnLog(logger),
)
client = tunnelpogs.RegistrationServer_PogsClient{Client: conn.Bootstrap(ctx), Conn: conn}
return client, nil
}