cloudflared-mirror/tunnelrpc/registration_server.go

41 lines
1.1 KiB
Go
Raw Normal View History

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()
}
}