ctrld/cmd/ctrld_library/packet_capture.go

package ctrld_library

import (
	"fmt"
	"net/netip"
	"runtime"
	"time"

	"github.com/Control-D-Inc/ctrld"
	"github.com/Control-D-Inc/ctrld/cmd/cli"
	"github.com/Control-D-Inc/ctrld/cmd/ctrld_library/netstack"
	"github.com/miekg/dns"
)

// PacketAppCallback extends AppCallback with packet read/write capabilities.
// Mobile platforms implementing full packet capture should use this interface.
type PacketAppCallback interface {
	AppCallback

	// ReadPacket reads a raw IP packet from the TUN interface.
	// This should be a blocking call that returns when a packet is available.
	ReadPacket() ([]byte, error)

	// WritePacket writes a raw IP packet back to the TUN interface.
	WritePacket(packet []byte) error

	// ClosePacketIO closes packet I/O resources.
	ClosePacketIO() error

	// ProtectSocket protects a socket file descriptor from being routed through the VPN.
	// On Android, this calls VpnService.protect() to prevent routing loops.
	// On iOS, this marks the socket to bypass the VPN.
	// Returns nil on success, error on failure.
	ProtectSocket(fd int) error
}

// PacketCaptureController holds state for packet capture mode
type PacketCaptureController struct {
	baseController *Controller

	// Packet capture mode fields
	netstackCtrl *netstack.NetstackController
	dnsBridge    *netstack.DNSBridge
	packetStopCh chan struct{}
}

// NewPacketCaptureController creates a new packet capture controller
func NewPacketCaptureController(appCallback PacketAppCallback) *PacketCaptureController {
	return &PacketCaptureController{
		baseController: &Controller{AppCallback: appCallback},
		packetStopCh:   make(chan struct{}),
	}
}

// StartWithPacketCapture starts ctrld in full packet capture mode for mobile.
// This method enables full IP packet processing with DNS filtering and upstream routing.
// It requires a PacketAppCallback that provides packet read/write capabilities.
func (pc *PacketCaptureController) StartWithPacketCapture(
	packetCallback PacketAppCallback,
	CdUID string,
	ProvisionID string,
	CustomHostname string,
	HomeDir string,
	UpstreamProto string,
	logLevel int,
	logPath string,
) error {
	if pc.baseController.stopCh != nil {
		return fmt.Errorf("controller already running")
	}

	// Set up configuration
	pc.baseController.Config = cli.AppConfig{
		CdUID:          CdUID,
		ProvisionID:    ProvisionID,
		CustomHostname: CustomHostname,
		HomeDir:        HomeDir,
		UpstreamProto:  UpstreamProto,
		Verbose:        logLevel,
		LogPath:        logPath,
	}
	pc.baseController.AppCallback = packetCallback

	// Set global socket protector for HTTP client sockets (API calls, etc)
	// This prevents routing loops when ctrld makes HTTP requests to api.controld.com
	ctrld.SetSocketProtector(packetCallback.ProtectSocket)

	// Create DNS bridge for communication between netstack and DNS proxy
	pc.dnsBridge = netstack.NewDNSBridge()
	pc.dnsBridge.Start()

	// Create packet handler that wraps the mobile callbacks
	packetHandler := netstack.NewMobilePacketHandler(
		packetCallback.ReadPacket,
		packetCallback.WritePacket,
		packetCallback.ClosePacketIO,
		packetCallback.ProtectSocket,
	)

	// Create DNS handler that uses the bridge
	dnsHandler := func(query []byte) ([]byte, error) {
		// Extract source IP from query context if available
		// For now, use a placeholder
		return pc.dnsBridge.ProcessQuery(query, "10.0.0.2", 0)
	}

	// Auto-detect platform and use appropriate TUN IP
	// Android: TUN=10.0.0.1, Device=10.0.0.2
	// iOS: TUN=10.0.0.2, Device=10.0.0.1
	tunIP := "10.0.0.1" // Default for Android
	// Check if running on iOS (no reliable way, so we'll make it configurable)
	// For now, use Android config. iOS should update their VPN settings to match.

	tunIPv4, err := netip.ParseAddr(tunIP)
	if err != nil {
		return fmt.Errorf("failed to parse TUN IPv4: %v", err)
	}

	netstackCfg := &netstack.Config{
		MTU:               1500,
		TUNIPv4:           tunIPv4,
		DNSHandler:        dnsHandler,
		UpstreamInterface: nil,  // Will use default interface
		EnableIPBlocking:  true, // Enable IP whitelisting in firewall mode
	}

	ctrld.ProxyLogger.Load().Info().Msgf("[PacketCapture] Netstack TUN IP: %s", tunIP)

	// Create netstack controller
	netstackCtrl, err := netstack.NewNetstackController(packetHandler, netstackCfg)
	if err != nil {
		pc.dnsBridge.Stop()
		return fmt.Errorf("failed to create netstack controller: %v", err)
	}

	pc.netstackCtrl = netstackCtrl

	// Start netstack processing
	if err := pc.netstackCtrl.Start(); err != nil {
		pc.dnsBridge.Stop()
		return fmt.Errorf("failed to start netstack: %v", err)
	}

	// Start regular ctrld DNS processing in background
	// This allows us to use existing DNS filtering logic
	pc.baseController.stopCh = make(chan struct{})

	// Start DNS query processor that receives queries from the bridge
	// and sends them to the ctrld DNS proxy
	go pc.processDNSQueries()

	// Start the main ctrld mobile runner
	go func() {
		appCallback := mapCallback(pc.baseController.AppCallback)
		cli.RunMobile(&pc.baseController.Config, &appCallback, pc.baseController.stopCh)
	}()

	// Log platform detection for DNS proxy port
	dnsPort := "5354"
	if runtime.GOOS == "ios" || runtime.GOOS == "darwin" {
		dnsPort = "53"
	}
	ctrld.ProxyLogger.Load().Info().Msgf("[PacketCapture] Platform: %s, DNS proxy port: %s", runtime.GOOS, dnsPort)

	return nil
}

// processDNSQueries processes DNS queries from the bridge using the ctrld DNS proxy
func (pc *PacketCaptureController) processDNSQueries() {
	queryCh := pc.dnsBridge.GetQueryChannel()

	for {
		select {
		case <-pc.packetStopCh:
			return
		case <-pc.baseController.stopCh:
			return
		case query := <-queryCh:
			go pc.handleDNSQuery(query)
		}
	}
}

// handleDNSQuery handles a single DNS query
func (pc *PacketCaptureController) handleDNSQuery(query *netstack.DNSQuery) {
	// Parse DNS message
	msg := new(dns.Msg)
	if err := msg.Unpack(query.Query); err != nil {
		return
	}

	// Determine DNS proxy port based on platform
	// Android: 0.0.0.0:5354
	// iOS: 127.0.0.1:53
	dnsProxyAddr := "127.0.0.1:5354" // Default for Android
	if runtime.GOOS == "ios" || runtime.GOOS == "darwin" {
		// iOS uses port 53
		dnsProxyAddr = "127.0.0.1:53"
	}

	// Send query to actual DNS proxy
	client := &dns.Client{
		Net:     "udp",
		Timeout: 3 * time.Second,
	}

	response, _, err := client.Exchange(msg, dnsProxyAddr)
	if err != nil {
		// Create SERVFAIL response
		response = new(dns.Msg)
		response.SetReply(msg)
		response.Rcode = dns.RcodeServerFailure
	}

	// Pack response
	responseBytes, err := response.Pack()
	if err != nil {
		return
	}

	// Send response back through bridge
	pc.dnsBridge.SendResponse(query.ID, responseBytes)
}

// Stop stops the packet capture controller
func (pc *PacketCaptureController) Stop(restart bool, pin int64) int {
	ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() called - starting shutdown")
	var errorCode = 0

	// Clear global socket protector
	ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - clearing socket protector")
	ctrld.SetSocketProtector(nil)

	// Stop DNS bridge
	if pc.dnsBridge != nil {
		ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - stopping DNS bridge")
		pc.dnsBridge.Stop()
		pc.dnsBridge = nil
		ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - DNS bridge stopped")
	}

	// Stop netstack
	if pc.netstackCtrl != nil {
		ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - stopping netstack controller")
		if err := pc.netstackCtrl.Stop(); err != nil {
			// Log error but continue shutdown
			ctrld.ProxyLogger.Load().Error().Msgf("[PacketCapture] Stop() - error stopping netstack: %v", err)
		}
		pc.netstackCtrl = nil
		ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - netstack controller stopped")
	}

	// Close packet stop channel
	if pc.packetStopCh != nil {
		ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - closing packet stop channel")
		select {
		case <-pc.packetStopCh:
			// Already closed
			ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - packet stop channel already closed")
		default:
			close(pc.packetStopCh)
			ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - packet stop channel closed")
		}
		pc.packetStopCh = make(chan struct{})
	}

	// Stop base controller
	ctrld.ProxyLogger.Load().Info().Msgf("[PacketCapture] Stop() - stopping base controller (restart=%v, pin=%d)", restart, pin)
	if !restart {
		errorCode = cli.CheckDeactivationPin(pin, pc.baseController.stopCh)
		ctrld.ProxyLogger.Load().Info().Msgf("[PacketCapture] Stop() - deactivation pin check returned: %d", errorCode)
	}
	if errorCode == 0 && pc.baseController.stopCh != nil {
		ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - closing base controller stop channel")
		select {
		case <-pc.baseController.stopCh:
			// Already closed
			ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - base controller stop channel already closed")
		default:
			close(pc.baseController.stopCh)
			ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Stop() - base controller stop channel closed")
		}
		pc.baseController.stopCh = nil
	}

	ctrld.ProxyLogger.Load().Info().Msgf("[PacketCapture] Stop() - shutdown complete, errorCode=%d", errorCode)
	return errorCode
}

// IsRunning returns true if the controller is running
func (pc *PacketCaptureController) IsRunning() bool {
	return pc.baseController.stopCh != nil
}

// IsPacketMode returns true (always in packet mode for this controller)
func (pc *PacketCaptureController) IsPacketMode() bool {
	return true
}

// SetFirewallMode enables or disables firewall mode (IP whitelisting) at runtime
func (pc *PacketCaptureController) SetFirewallMode(enabled bool) {
	if pc.netstackCtrl != nil {
		pc.netstackCtrl.SetFirewallMode(enabled)
		if enabled {
			ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Firewall mode ENABLED - IP whitelisting active")
		} else {
			ctrld.ProxyLogger.Load().Info().Msg("[PacketCapture] Firewall mode DISABLED - all IPs allowed")
		}
	}
}