71 lines
2.2 KiB
Go
71 lines
2.2 KiB
Go
package quic
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"time"
|
|
|
|
"zombiezen.com/go/capnproto2/rpc"
|
|
|
|
"github.com/cloudflare/cloudflared/tunnelrpc"
|
|
"github.com/cloudflare/cloudflared/tunnelrpc/pogs"
|
|
)
|
|
|
|
// HandleRequestFunc wraps the proxied request from the upstream and also provides methods on the stream to
|
|
// handle the response back.
|
|
type HandleRequestFunc = func(ctx context.Context, stream *RequestServerStream) error
|
|
|
|
// CloudflaredServer 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 CloudflaredServer struct {
|
|
handleRequest HandleRequestFunc
|
|
sessionManager pogs.SessionManager
|
|
configManager pogs.ConfigurationManager
|
|
responseTimeout time.Duration
|
|
}
|
|
|
|
func NewCloudflaredServer(handleRequest HandleRequestFunc, sessionManager pogs.SessionManager, configManager pogs.ConfigurationManager, responseTimeout time.Duration) *CloudflaredServer {
|
|
return &CloudflaredServer{
|
|
handleRequest: handleRequest,
|
|
sessionManager: sessionManager,
|
|
configManager: configManager,
|
|
responseTimeout: responseTimeout,
|
|
}
|
|
}
|
|
|
|
// Serve executes the defined handlers in ServerStream on the provided stream if it is a proper RPC stream with the
|
|
// correct preamble protocol signature.
|
|
func (s *CloudflaredServer) Serve(ctx context.Context, stream io.ReadWriteCloser) error {
|
|
signature, err := determineProtocol(stream)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
switch signature {
|
|
case dataStreamProtocolSignature:
|
|
return s.handleRequest(ctx, &RequestServerStream{stream})
|
|
case rpcStreamProtocolSignature:
|
|
return s.handleRPC(ctx, stream)
|
|
default:
|
|
return fmt.Errorf("unknown protocol %v", signature)
|
|
}
|
|
}
|
|
|
|
func (s *CloudflaredServer) handleRPC(ctx context.Context, stream io.ReadWriteCloser) error {
|
|
ctx, cancel := context.WithTimeout(ctx, s.responseTimeout)
|
|
defer cancel()
|
|
transport := tunnelrpc.SafeTransport(stream)
|
|
defer transport.Close()
|
|
|
|
main := pogs.CloudflaredServer_ServerToClient(s.sessionManager, s.configManager)
|
|
rpcConn := rpc.NewConn(transport, rpc.MainInterface(main.Client))
|
|
defer rpcConn.Close()
|
|
|
|
// We ignore the errors here because if cloudflared fails to handle a request, we will just move on.
|
|
select {
|
|
case <-rpcConn.Done():
|
|
case <-ctx.Done():
|
|
}
|
|
return nil
|
|
}
|