zt/tunnel/intercept/proxy/proxy.go

/*
	Copyright NetFoundry Inc.

	Licensed under the Apache License, Version 2.0 (the "License");
	you may not use this file except in compliance with the License.
	You may obtain a copy of the License at

	https://www.apache.org/licenses/LICENSE-2.0

	Unless required by applicable law or agreed to in writing, software
	distributed under the License is distributed on an "AS IS" BASIS,
	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	See the License for the specific language governing permissions and
	limitations under the License.
*/

package proxy

import (
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"sync"

	"github.com/michaelquigley/pfxlog"
	"github.com/hanzozt/foundation/v2/info"
	"github.com/hanzozt/foundation/v2/mempool"
	"github.com/hanzozt/transport/v2"
	"github.com/hanzozt/zt/v2/common/alert"
	"github.com/hanzozt/zt/v2/tunnel"
	"github.com/hanzozt/zt/v2/tunnel/dns"
	"github.com/hanzozt/zt/v2/tunnel/entities"
	"github.com/hanzozt/zt/v2/tunnel/intercept"
	"github.com/hanzozt/zt/v2/tunnel/udp_vconn"
)

type Service struct {
	Name      string
	Port      int
	Protocols []intercept.Protocol
	Binding   string
	Closers   map[string]io.Closer
	sync.Mutex
	TunnelService *entities.Service
}

func (self *Service) addCloser(closerType string, c io.Closer) {
	self.Lock()
	defer self.Unlock()
	self.Closers[closerType] = c
}

func (self *Service) Stop() error {
	self.Lock()
	defer self.Unlock()

	logger := pfxlog.Logger().WithField("service", self.Name)

	var errList []error
	for listenerType, closer := range self.Closers {
		logger.Infof("closing %s listener for service of type", listenerType)
		if err := closer.Close(); err != nil {
			errList = append(errList, err)
		}
	}
	clear(self.Closers)
	return errors.Join(errList...)
}

type interceptor struct {
	interceptIP net.IP
	services    map[string]*Service
	lock        sync.Mutex
	alerter     Alerter
}

type Alerter interface {
	ReportError(message string, details []string, relatedEntities map[string]string)
}

type DefaultAlerter struct{}

func (d DefaultAlerter) ReportError(message string, details []string, relatedEntities map[string]string) {
	pfxlog.Logger().WithField("details", details).WithField("relatedEntities", relatedEntities).Errorf("alert: %s", message)
}

func NewDelegate(alerter Alerter) intercept.Interceptor {
	return &interceptor{
		interceptIP: net.IPv4(127, 0, 0, 1),
		services:    map[string]*Service{},
		alerter:     alerter,
	}
}

func New(alerter Alerter, ip net.IP, serviceList []string) (intercept.Interceptor, error) {
	services := make(map[string]*Service, len(serviceList))

	for _, arg := range serviceList {
		parts := strings.Split(arg, ":")
		if len(parts) < 2 || len(parts) > 3 {
			return nil, fmt.Errorf("invalid argument '%s'", arg)
		}

		port, err := strconv.Atoi(parts[1])
		if err != nil {
			return nil, fmt.Errorf("invalid port specified in '%s'", arg)
		}

		service := &Service{
			Name:    parts[0],
			Port:    port,
			Closers: map[string]io.Closer{},
		}

		if len(parts) == 3 {
			protocol := parts[2]
			switch protocol {
			case "tcp":
				service.Protocols = append(service.Protocols, intercept.TCP)
			case "udp":
				service.Protocols = append(service.Protocols, intercept.UDP)
			default:
				return nil, fmt.Errorf("invalid protocol specified in '%s', must be tcp or udp", arg)
			}
		} else {
			// Default to TCP if no protocol specified
			service.Protocols = append(service.Protocols, intercept.TCP)
		}
		services[parts[0]] = service
	}

	p := interceptor{
		interceptIP: ip,
		services:    services,
		alerter:     alerter,
	}
	return &p, nil
}

func (p *interceptor) Intercept(service *entities.Service, _ dns.Resolver, _ intercept.AddressTracker) error {
	p.lock.Lock()
	defer p.lock.Unlock()

	log := pfxlog.Logger().WithField("service", service.Name)

	proxiedService, hasStaticProxyConfig := p.services[*service.Name]

	if proxiedService != nil {
		if err := proxiedService.Stop(); err != nil {
			log.WithError(err).Error("error stopping existing proxied service listeners")
		}
	}

	proxyConfig := &entities.ProxyV1Config{}
	hasProxyConfig, err := service.GetConfigOfType(entities.ProxyV1, proxyConfig)
	if err != nil {
		if !hasStaticProxyConfig {
			return err
		} else {
			log.WithError(err).Error("error loading dynamic service proxy configuration")
		}
	} else if hasProxyConfig {
		proxiedService = &Service{
			Name:    *service.Name,
			Port:    int(proxyConfig.Port),
			Binding: proxyConfig.Binding,
			Closers: map[string]io.Closer{},
		}

		for _, protocol := range proxyConfig.Protocols {
			if protocol == "tcp" {
				proxiedService.Protocols = append(proxiedService.Protocols, intercept.TCP)
			} else if protocol == "udp" {
				proxiedService.Protocols = append(proxiedService.Protocols, intercept.UDP)
			} else {
				return fmt.Errorf("unexpected proxy protocol '%s' for service '%s'", protocol, proxiedService.Name)
			}
		}

		p.services[proxiedService.Name] = proxiedService
	}

	if proxiedService == nil {
		log.Debugf("service %v was not specified at initialization and has no proxy config, no proxy listener being created", service.Name)
		return nil
	}

	proxiedService.TunnelService = service

	// pre-fetch network session todo move this to service poller?
	service.FabricProvider.PrepForUse(*service.ID)

	return p.runServiceListener(proxiedService)
}

func (p *interceptor) runServiceListener(service *Service) error {
	var errList []error
	for _, protocol := range service.Protocols {
		if protocol == intercept.TCP {
			if err := p.handleTCP(service); err != nil {
				errList = append(errList, err)
			}
		} else if protocol == intercept.UDP {
			if err := p.handleUDP(service); err != nil {
				errList = append(errList, err)
			}
		}
	}

	if len(errList) > 0 {
		p.alerter.ReportError("unable to proxy service", alert.ErrListToDetailList(errList), map[string]string{
			alert.EntityTypeService: service.TunnelService.GetId(),
			alert.EntityTypeSelf:    alert.EntityTypeSelfErt,
		})
	}

	return errors.Join(errList...)
}

func (p *interceptor) handleTCP(service *Service) error {
	log := pfxlog.Logger().WithField("service", service.Name)

	serviceBinding := p.interceptIP
	if service.Binding != "" {
		// Special-case wildcard bindings. ResolveLocalBinding expects an interface-mappable IP,
		// but 0.0.0.0 is a valid local bind address meaning "all interfaces".
		if ip := net.ParseIP(service.Binding); ip != nil && ip.IsUnspecified() {
			if ip.To4() == nil {
				return fmt.Errorf("unable to resolve binding '%s' on service '%s' for tcp proxy (tcp4 listener does not support IPv6 wildcard)",
					service.Binding, service.Name)
			}
			serviceBinding = ip
		} else {
			binding, err := transport.ResolveLocalBinding(service.Binding)
			if err != nil {
				return fmt.Errorf("unable to resolve binding '%s' on service '%s' for tcp proxy (%w)",
					service.Binding, service.Name, err)
			}
			serviceBinding = binding
		}
	}

	listenAddr := net.TCPAddr{IP: serviceBinding, Port: service.Port}
	server, err := net.Listen("tcp4", listenAddr.String())
	if err != nil {
		return err
	}
	service.addCloser("tcp", server)

	log = log.WithField("addr", server.Addr().String())

	log.Info("service is listening")

	go func() {
		defer log.Info("service stopped")

		for {
			conn, err := server.Accept()
			if err != nil {
				log.WithError(err).Error("accept failed")
				return
			}
			sourceAddr := service.TunnelService.GetSourceAddr(conn.RemoteAddr(), conn.LocalAddr())
			appInfo := tunnel.GetAppInfo("tcp", "", p.interceptIP.String(), strconv.Itoa(service.Port), sourceAddr)
			identity := service.TunnelService.GetDialIdentity(conn.RemoteAddr(), conn.LocalAddr())
			go tunnel.DialAndRun(service.TunnelService.FabricProvider, service.TunnelService, identity, conn, appInfo, true)
		}
	}()

	return nil
}

func (p *interceptor) handleUDP(service *Service) error {
	log := pfxlog.Logger().WithField("service", service.Name)

	serviceBinding := p.interceptIP
	if service.Binding != "" {
		// Special-case wildcard binding. 0.0.0.0 is a valid local bind address meaning "all interfaces".
		if ip := net.ParseIP(service.Binding); ip != nil && ip.IsUnspecified() {
			serviceBinding = ip
		} else {
			binding, err := transport.ResolveLocalBinding(service.Binding)
			if err != nil {
				return fmt.Errorf("unable to resolve binding '%s' on service '%s' for udp proxy (%w)", service.Binding, service.Name, err)
			}
			serviceBinding = binding
		}
	}

	if serviceBinding.To4() == nil {
		return fmt.Errorf("unable to resolve binding '%s' on service '%s' for udp proxy (IPv6 is not supported)",
			service.Binding, service.Name)
	}
	serviceBinding = serviceBinding.To4()

	listenAddr := &net.UDPAddr{IP: serviceBinding, Port: service.Port}
	udpPacketConn, err := net.ListenUDP("udp4", listenAddr)
	if err != nil {
		return err
	}

	service.addCloser("udp", udpPacketConn)

	log = log.WithField("addr", udpPacketConn.LocalAddr().String())

	log.Infof("service %v is listening", service.Name)
	reader := &udpReader{
		service: service.TunnelService,
		conn:    udpPacketConn,
	}
	vconnManager := udp_vconn.NewManager(service.TunnelService.FabricProvider, udp_vconn.NewUnlimitedConnectionPolicy(), udp_vconn.NewDefaultExpirationPolicy(), true)
	go reader.generateReadEvents(vconnManager)
	return nil
}

func (p *interceptor) Stop() {
	pfxlog.Logger().Info("stopping proxy interceptor")

	for _, service := range p.services {
		_ = service.Stop()
	}
}

func (p *interceptor) StopIntercepting(serviceName string, _ intercept.AddressTracker) error {
	p.lock.Lock()
	defer p.lock.Unlock()

	pfxlog.Logger().WithField("service", serviceName).Info("stopping proxy interceptor for service")
	if service, ok := p.services[serviceName]; ok {
		return service.Stop()
	}
	return nil
}

type udpReader struct {
	service *entities.Service
	conn    *net.UDPConn
}

func (reader *udpReader) generateReadEvents(manager udp_vconn.Manager) {
	log := pfxlog.Logger().WithField("service", reader.service.Name)
	bufPool := mempool.NewPool(16, info.MaxUdpPacketSize)
	for {
		buf := bufPool.AcquireBuffer()
		n, srcAddr, err := reader.conn.ReadFromUDP(buf.Buf)
		if err != nil {
			log.WithError(err).Error("failure while reading udp message. stopping UDP read loop")
			manager.QueueError(err)
			return
		}

		log.Debugf("read %v bytes from udp, queuing", len(buf.GetPayload()))
		buf.Buf = buf.Buf[:n]
		event := &udpReadEvent{
			reader:  reader,
			buf:     buf,
			srcAddr: srcAddr,
		}
		manager.QueueEvent(event)
	}
}

type udpReadEvent struct {
	reader  *udpReader
	buf     *mempool.DefaultPooledBuffer
	srcAddr net.Addr
}

func (event *udpReadEvent) Handle(manager udp_vconn.Manager) error {
	log := pfxlog.Logger()

	writeQueue := manager.GetWriteQueue(event.srcAddr)

	if writeQueue == nil {
		log.Infof("received connection for %v --> %v, which maps to intercepted service %v",
			event.srcAddr, event.reader.conn.LocalAddr(), event.reader.service)
		var err error
		writeQueue, err = manager.CreateWriteQueue(event.srcAddr.(*net.UDPAddr), event.srcAddr, event.reader.service, event.reader.conn)
		if err != nil {
			return err
		}
	}

	log.Infof("received %v bytes for conn %v -> %v", len(event.buf.Buf), writeQueue.LocalAddr(), writeQueue.Service())
	writeQueue.Accept(event.buf)

	return nil
}