banger/internal/daemon/ports.go

package daemon

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"banger/internal/api"
	"banger/internal/model"
	"banger/internal/system"
	"banger/internal/vmdns"
	"banger/internal/vsockagent"
)

const httpProbeTimeout = 750 * time.Millisecond

func (d *Daemon) PortsVM(ctx context.Context, idOrName string) (result api.VMPortsResult, err error) {
	_, err = d.withVMLockByRef(ctx, idOrName, func(vm model.VMRecord) (model.VMRecord, error) {
		result.Name = vm.Name
		result.DNSName = strings.TrimSpace(vm.Runtime.DNSName)
		if result.DNSName == "" && strings.TrimSpace(vm.Name) != "" {
			result.DNSName = vmdns.RecordName(vm.Name)
		}
		if vm.State != model.VMStateRunning || !system.ProcessRunning(vm.Runtime.PID, vm.Runtime.APISockPath) {
			return model.VMRecord{}, fmt.Errorf("vm %s is not running", vm.Name)
		}
		if strings.TrimSpace(vm.Runtime.GuestIP) == "" {
			return model.VMRecord{}, errors.New("vm has no guest IP")
		}
		if strings.TrimSpace(vm.Runtime.VSockPath) == "" {
			return model.VMRecord{}, errors.New("vm has no vsock path")
		}
		if vm.Runtime.VSockCID == 0 {
			return model.VMRecord{}, errors.New("vm has no vsock cid")
		}
		if err := d.ensureSocketAccess(ctx, vm.Runtime.VSockPath, "firecracker vsock socket"); err != nil {
			return model.VMRecord{}, err
		}
		portsCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
		defer cancel()
		listeners, err := vsockagent.Ports(portsCtx, d.logger, vm.Runtime.VSockPath)
		if err != nil {
			return model.VMRecord{}, err
		}
		result.Ports = buildVMPorts(vm, listeners)
		return vm, nil
	})
	return result, err
}

func buildVMPorts(vm model.VMRecord, listeners []vsockagent.PortListener) []api.VMPort {
	endpointHost := strings.TrimSpace(vm.Runtime.DNSName)
	if endpointHost == "" {
		endpointHost = strings.TrimSpace(vm.Runtime.GuestIP)
	}
	probeHost := strings.TrimSpace(vm.Runtime.GuestIP)
	ports := make([]api.VMPort, 0, len(listeners))
	for _, listener := range listeners {
		if listener.Port <= 0 {
			continue
		}
		port := api.VMPort{
			Proto:       strings.ToLower(strings.TrimSpace(listener.Proto)),
			BindAddress: strings.TrimSpace(listener.BindAddress),
			Port:        listener.Port,
			PID:         listener.PID,
			Process:     strings.TrimSpace(listener.Process),
			Command:     strings.TrimSpace(listener.Command),
			Endpoint:    net.JoinHostPort(endpointHost, strconv.Itoa(listener.Port)),
		}
		if port.Command == "" {
			port.Command = port.Process
		}
		if port.Proto == "tcp" && probeHost != "" && endpointHost != "" && probeHTTPListener(probeHost, listener.Port) {
			port.WebURL = "http://" + net.JoinHostPort(endpointHost, strconv.Itoa(listener.Port)) + "/"
		}
		ports = append(ports, port)
	}
	sort.Slice(ports, func(i, j int) bool {
		if ports[i].Proto != ports[j].Proto {
			return ports[i].Proto < ports[j].Proto
		}
		if ports[i].Port != ports[j].Port {
			return ports[i].Port < ports[j].Port
		}
		if ports[i].Process != ports[j].Process {
			return ports[i].Process < ports[j].Process
		}
		return ports[i].Command < ports[j].Command
	})
	return ports
}

func probeHTTPListener(guestIP string, port int) bool {
	if strings.TrimSpace(guestIP) == "" || port <= 0 {
		return false
	}
	url := "http://" + net.JoinHostPort(strings.TrimSpace(guestIP), strconv.Itoa(port)) + "/"
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return false
	}
	client := &http.Client{
		Timeout: httpProbeTimeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
		Transport: &http.Transport{
			Proxy: nil,
		},
	}
	resp, err := client.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1))
	return resp.ProtoMajor >= 1
}