tunnel/provider.go

Rama main · réplica de lectura del repositorio original.

package tunnel

import (
	"encoding/json"
	"io"
	"net"
	"net/netip"
	"time"

	"github.com/openziti/edge-api/rest_model"
	"github.com/openziti/sdk-golang/v2/ziti"
	"github.com/openziti/sdk-golang/v2/ziti/edge"
	"github.com/openziti/ziti/v2/tunnel/health"
	"github.com/sirupsen/logrus"
)

type PortRange struct {
	Low  uint16
	High uint16
}

type AllowConfig interface {
	GetAllowedProtocols() []string
	GetAllowedPortRanges() []PortRange
	GetAllowedAddresses() []string
}

type Dialer interface {
	Dial(network string, address string) (net.Conn, error)
}

type DialWrapper interface {
	Dial(hostContext HostingContext, localAddr *netip.AddrPort, dialer Dialer, protocol string, address string) (net.Conn, error)
}

type HostedService interface {
	GetName() string
	GetId() string
	IsEncryptionRequired() bool
	GetDialTimeout() time.Duration
	GetConfigOfType(configType string, target interface{}) (bool, error)
}

type HostingContext interface {
	Service() HostedService
	ServiceId() string
	GetTerminatorIdCacheKey() string
	ServiceName() string
	ListenOptions() *ziti.ListenOptions
	// Dial connects to the upstream selected by options, an accepted conn's app data, and reports
	// whether the returned conn supports half close. It must be safe for concurrent use.
	Dial(options map[string]interface{}) (net.Conn, bool, error)
	GetHealthChecks() []health.CheckDefinition
	GetInitialHealthState() (ziti.Precedence, uint16)
	OnClose()
	SetCloseCallback(func())
	GetAllowConfig() AllowConfig
	SetDialWrapper(listener DialWrapper)
}

type HostControl interface {
	io.Closer
	UpdateCost(cost uint16) error
	UpdatePrecedence(precedence edge.Precedence) error
	UpdateCostAndPrecedence(cost uint16, precedence edge.Precedence) error
	SendHealthEvent(pass bool) error
}

type FabricProvider interface {
	PrepForUse(serviceId string)
	GetCurrentIdentity() (*rest_model.IdentityDetail, error)
	GetCurrentIdentityWithBackoff() (*rest_model.IdentityDetail, error)
	TunnelService(service Service, identity string, conn net.Conn, halfClose bool, appInfo []byte) error
	HostService(hostCtx HostingContext) (HostControl, error)
}

func AppDataToMap(appData []byte) (map[string]interface{}, error) {
	result := map[string]interface{}{}
	if len(appData) != 0 {
		if err := json.Unmarshal(appData, &result); err != nil {
			return nil, err
		}
	}
	return result, nil
}

func NewContextProvider(context ziti.Context) FabricProvider {
	return &contextProvider{
		Context: context,
	}
}

type contextProvider struct {
	ziti.Context
}

func (self *contextProvider) PrepForUse(serviceId string) {
	if _, err := self.GetSession(serviceId); err != nil {
		logrus.WithError(err).Error("failed to acquire network session")
	} else {
		logrus.Debug("acquired network session")
	}
}

func (self *contextProvider) TunnelService(service Service, identity string, conn net.Conn, halfClose bool, appData []byte) error {
	options := &ziti.DialOptions{
		ConnectTimeout: service.GetDialTimeout(),
		AppData:        appData,
		Identity:       identity,
	}

	zitiConn, err := self.DialWithOptions(service.GetName(), options)
	if err != nil {
		return err
	}

	Run(zitiConn, conn, halfClose)
	return nil
}

func (self *contextProvider) HostService(hostCtx HostingContext) (HostControl, error) {
	logger := logrus.WithField("service", hostCtx.ServiceName())
	listener, err := self.ListenWithOptions(hostCtx.ServiceName(), hostCtx.ListenOptions())
	if err != nil {
		logger.WithError(err).Error("error listening for service")
		return nil, err
	}

	go self.accept(listener, hostCtx)

	return listener, nil
}

// accept hands each conn accepted on listener to its own goroutine, so a slow upstream dial or
// accept completion holds up only that conn. It calls hostCtx.OnClose once the listener closes,
// which may be before conns already handed off have finished.
func (self *contextProvider) accept(listener edge.Listener, hostCtx HostingContext) {
	defer hostCtx.OnClose()

	logger := logrus.WithField("service", hostCtx.ServiceName())
	for {
		logger.Info("hosting service, waiting for connections")
		conn, err := listener.AcceptEdge()
		if err != nil {
			logger.WithError(err).Error("closing listener for service")
			return
		}

		go handleAccepted(conn, hostCtx, logger)
	}
}

// handleAccepted dials the upstream for an accepted conn, completes the accept, and runs the
// connection until both directions finish. On failure it closes conn and any upstream conn it
// opened.
func handleAccepted(conn edge.Conn, hostCtx HostingContext, logger *logrus.Entry) {
	logger = logger.WithField("circuitId", conn.GetCircuitId())

	options, err := AppDataToMap(conn.GetAppData())
	if err != nil {
		logger.WithError(err).Error("dial failed")
		conn.CompleteAcceptFailed(err)
		if closeErr := conn.Close(); closeErr != nil {
			logger.WithError(closeErr).Error("close of ziti connection failed")
		}
		return
	}

	externalConn, halfClose, err := hostCtx.Dial(options)
	if err != nil {
		logger.WithError(err).Error("dial failed")
		conn.CompleteAcceptFailed(err)
		if closeErr := conn.Close(); closeErr != nil {
			logger.WithError(closeErr).Error("close of ziti connection failed")
		}
		return
	}

	logger.Infof("successful connection %v->%v", conn.LocalAddr(), conn.RemoteAddr())

	if err := conn.CompleteAcceptSuccess(); err != nil {
		logger.WithError(err).Error("complete accept success failed")

		if closeErr := conn.Close(); closeErr != nil {
			logger.WithError(closeErr).Error("close of ziti connection failed")
		}

		if closeErr := externalConn.Close(); closeErr != nil {
			logger.WithError(closeErr).Error("close of external connection failed")
		}
		return
	}

	Run(conn, externalConn, halfClose)
}