tunserver

type DataCallback interface {
	Header([]*rpc.MetadataKV) error
	Message([]byte) error
	Trailer([]*rpc.MetadataKV) error
	Error(*statuspb.Status) error
}

type FindHandle interface {
	// Get finds a tunnel to an agentk.
	// It waits for a matching tunnel to proxy a connection through. When a matching tunnel is found, it is returned.
	// It returns gRPC status errors only, ready to return from RPC handler.
	Get(ctx context.Context) (Tunnel, error)
	// Done must be called to free resources of this FindHandle instance.
	// ctx is used for tracing only.
	Done(ctx context.Context)
}

type GatewayFinder[T Target] interface {
	Find(ctx context.Context) (ReadyGateway[T], error)
}

type PollGatewayURLsCallback[T Target] func(gatewayURLs []T)

type PollingGatewayURLQuerier[T Target] interface {
	PollGatewayURLs(ctx context.Context, agentKey api.AgentKey, cb PollGatewayURLsCallback[T])
	CachedGatewayURLs(agentKey api.AgentKey) []T
}

type ReadyGateway[T Target] struct {
	URL          T
	Stream       grpc.ClientStream
	StreamCancel context.CancelFunc
	Codec        encoding.CodecV2
}

type Router[T Target] struct {
	Plugin RouterPlugin[T]
	// PrivateAPIServer is the gRPC server that other kas instances can talk to.
	// Request handlers can obtain the per-request logger using modshared.RPCAPIFromContext(requestContext).
	PrivateAPIServer grpc.ServiceRegistrar
}

type RouterPlugin[T Target] interface {
	// FindReadyGateway finds a ready gateway tunserver to route the request through.
	// It returns an error, compatible with gRPC status package.
	FindReadyGateway(ctx context.Context, log *slog.Logger, method string) (ReadyGateway[T], *slog.Logger, api.AgentKey, error)
	// FindTunnel starts searching for a tunnel to a matching tunclient.
	// Tunnel found boolean indicates whether a suitable tunnel is immediately available from the
	// returned FindHandle object.
	FindTunnel(grpc.ServerStream, modshared.RPCAPI) (bool, *slog.Logger, FindHandle, error)
	// PrepareStreamForForwarding is used to perform any actions, like wrapping, on the stream just before
	// forwarding it to tunclient.
	// Returned error is compatible with gRPC-status error.
	PrepareStreamForForwarding(stream grpc.ServerStream) (grpc.ServerStream, error)
}

type RoutingClientConn[T Target] struct {
	Log    *slog.Logger
	API    modshared.API
	Plugin RouterPlugin[T]
}

type StateType int

type Target interface {
	comparable
	String() string
}

type Tunnel interface {
	// ForwardStream performs bi-directional message forwarding between incomingStream and the tunnel.
	// cb is called with header, messages and trailer coming from the tunnel. It's the callers
	// responsibility to forward them into the incomingStream.
	ForwardStream(log *slog.Logger, rpcAPI modshared.RPCAPI, incomingStream grpc.ServerStream, cb DataCallback) error
	// Done must be called when the caller is done with the Tunnel.
	// ctx is used for tracing only.
	Done(ctx context.Context)
}

type TunnelImpl struct {
	Tunnel       grpc.BidiStreamingServer[rpc.ConnectRequest, rpc.ConnectResponse]
	TunnelRetErr chan<- error
	AgentKey     api.AgentKey
	Descriptor   *info.APIDescriptor
	State        StateType

	OnForward func(*TunnelImpl) error
	OnDone    func(context.Context, *TunnelImpl)
}