50 lines
1.3 KiB
Go
50 lines
1.3 KiB
Go
package connection
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/sirupsen/logrus"
|
|
rpc "zombiezen.com/go/capnproto2/rpc"
|
|
|
|
"github.com/cloudflare/cloudflared/h2mux"
|
|
"github.com/cloudflare/cloudflared/tunnelrpc"
|
|
tunnelpogs "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
|
|
)
|
|
|
|
// NewRPCClient creates and returns a new RPC client, which will communicate
|
|
// using a stream on the given muxer
|
|
func NewRPCClient(
|
|
ctx context.Context,
|
|
muxer *h2mux.Muxer,
|
|
logger *logrus.Entry,
|
|
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}
|
|
return client, nil
|
|
}
|
|
|
|
func isRPCStreamResponse(headers []h2mux.Header) bool {
|
|
return len(headers) == 1 &&
|
|
headers[0].Name == ":status" &&
|
|
headers[0].Value == "200"
|
|
}
|