uopi/pkg/ca/search.go

package ca

import (
	"context"
	"encoding/binary"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/uopi/goca/proto"
)

// searchResult is returned to a waiter when a search succeeds.
type searchResult struct {
	addr string // "ip:port" TCP address of the IOC
}

// searchWaiter tracks an outstanding PV search request.
type searchWaiter struct {
	pvName string
	ch     chan searchResult // closed or sent to exactly once
}

// searchEngine performs UDP-based PV name → IOC address resolution.
// It sends CA_PROTO_SEARCH datagrams to all configured server addresses and
// listens for replies, matching them to pending waiters via a searchID.
type searchEngine struct {
	addrs []string // "ip:port" strings to search

	mu      sync.Mutex
	waiters map[uint32]*searchWaiter // searchID → waiter

	idSeq atomic.Uint32
	conn  *net.UDPConn // shared UDP socket (nil if not started)
}

func newSearchEngine(addrs []string) *searchEngine {
	return &searchEngine{
		addrs:   addrs,
		waiters: make(map[uint32]*searchWaiter),
	}
}

// start opens the shared UDP socket and launches the receive goroutine.
// ctx cancellation stops the receive loop and closes the socket.
func (se *searchEngine) start(ctx context.Context) error {
	conn, err := net.ListenUDP("udp4", &net.UDPAddr{})
	if err != nil {
		return fmt.Errorf("ca search: open UDP socket: %w", err)
	}
	se.conn = conn

	go se.recvLoop(ctx)
	return nil
}

// lookup resolves pvName to an IOC TCP address.
// It retries with exponential back-off until ctx is cancelled or done.
func (se *searchEngine) lookup(ctx context.Context, pvName string) (string, error) {
	id := se.idSeq.Add(1)
	w := &searchWaiter{
		pvName: pvName,
		ch:     make(chan searchResult, 1),
	}

	se.mu.Lock()
	se.waiters[id] = w
	se.mu.Unlock()

	defer func() {
		se.mu.Lock()
		delete(se.waiters, id)
		se.mu.Unlock()
	}()

	delay := 100 * time.Millisecond
	const maxDelay = 30 * time.Second

	for {
		if err := se.sendSearch(id, pvName); err != nil {
			return "", err
		}

		select {
		case res := <-w.ch:
			return res.addr, nil
		case <-time.After(delay):
			delay *= 2
			if delay > maxDelay {
				delay = maxDelay
			}
		case <-ctx.Done():
			return "", fmt.Errorf("ca search %q: %w", pvName, ctx.Err())
		}
	}
}

// sendSearch sends a CA_PROTO_VERSION + CA_PROTO_SEARCH burst to all addresses.
// The CA protocol requires both messages to be in the same UDP datagram.
func (se *searchEngine) sendSearch(id uint32, pvName string) error {
	payload := proto.BuildStringPayload(pvName)

	// VERSION message (no payload).
	verHdr := proto.Header{
		Command:   proto.CmdVersion,
		DataCount: proto.MinorVersion,
	}

	// SEARCH message.
	searchHdr := proto.Header{
		Command:    proto.CmdSearch,
		DataType:   proto.SearchReply,
		DataCount:  proto.MinorVersion,
		Parameter1: id,
		Parameter2: id,
	}
	searchMsg := proto.BuildMessage(searchHdr, payload)

	// Bundle into a single datagram — softIoc (and the CA Repeater) require
	// VERSION and SEARCH to arrive together in one UDP packet.
	pkt := make([]byte, 0, proto.HeaderSize+len(searchMsg))
	pkt = append(pkt, verHdr.Bytes()...)
	pkt = append(pkt, searchMsg...)

	for _, addr := range se.addrs {
		udpAddr, err := net.ResolveUDPAddr("udp4", addr)
		if err != nil {
			continue
		}
		_, _ = se.conn.WriteTo(pkt, udpAddr)
		dbg("CA search sent", "pv", pvName, "id", id, "dst", udpAddr)
	}
	return nil
}

// recvLoop reads UDP datagrams and dispatches search replies.
func (se *searchEngine) recvLoop(ctx context.Context) {
	buf := make([]byte, 65536)
	for {
		// Use a short read deadline so we can check ctx.Done() periodically.
		_ = se.conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
		n, srcAddr, err := se.conn.ReadFromUDP(buf)
		if err != nil {
			if ctx.Err() != nil {
				return
			}
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue
			}
			continue
		}

		dbg("CA UDP datagram received", "src", srcAddr, "bytes", n)
		se.parseReply(buf[:n], srcAddr)
	}
}

// parseReply scans a UDP datagram for CA_PROTO_SEARCH response messages and
// wakes the matching waiter.
func (se *searchEngine) parseReply(data []byte, src *net.UDPAddr) {
	for len(data) >= proto.HeaderSize {
		h, n, err := proto.DecodeHeader(newBytesReader(data))
		if err != nil {
			return
		}
		payloadEnd := n + int(h.PayloadSize)
		if payloadEnd > len(data) {
			return
		}
		payload := data[n:payloadEnd]
		data = data[payloadEnd:]

		if h.Command != proto.CmdSearch {
			continue
		}

		// data_type field = TCP port of the CA server.
		tcpPort := int(h.DataType)
		searchID := h.Parameter2

		// Resolve server IP.
		//
		// Two payload formats exist in the wild:
		//
		//   New (≥ CA 4.9):  [IP(4)][pad(4)]
		//   Old (< CA 4.9):  [minor_version(2)][pad(2)][IP(4)]
		//
		// The old format is identified by payload[0] == 0x00 (the high byte of a
		// uint16 minor version is always zero for CA versions ≤ 255, whereas a
		// real IPv4 address never starts with 0x00).
		//
		// In both formats, 0xFFFFFFFF in the IP position means "use sender IP".
		var serverIP net.IP
		if len(payload) >= 8 && payload[0] == 0x00 {
			// Old format: minor version at [0:2], IP at [4:8].
			ipSlice := payload[4:8]
			if useSenderIP(ipSlice) {
				serverIP = src.IP
			} else {
				ip := make(net.IP, 4)
				copy(ip, ipSlice)
				serverIP = ip
			}
		} else if len(payload) >= 4 {
			// New format: IP at [0:4].
			if useSenderIP(payload[:4]) {
				serverIP = src.IP
			} else {
				ip := make(net.IP, 4)
				copy(ip, payload[:4])
				serverIP = ip
			}
		} else {
			serverIP = src.IP
		}

		tcpAddr := fmt.Sprintf("%s:%d", serverIP.String(), tcpPort)
		dbg("CA search reply", "searchID", searchID, "tcpAddr", tcpAddr, "src", src)

		se.mu.Lock()
		w, ok := se.waiters[searchID]
		if ok {
			delete(se.waiters, searchID)
		}
		se.mu.Unlock()

		if ok {
			dbg("CA search matched", "pv", w.pvName, "ioc", tcpAddr)
			select {
			case w.ch <- searchResult{addr: tcpAddr}:
			default:
			}
		} else {
			dbg("CA search reply unmatched", "searchID", searchID, "known_waiters", func() int {
				se.mu.Lock()
				defer se.mu.Unlock()
				return len(se.waiters)
			}())
		}
	}
}

// useSenderIP reports whether an IP field in a CA search reply means
// "use the sender's address".  Both 0.0.0.0 (all-zeros) and 255.255.255.255
// (all-ones / 0xFFFFFFFF) are used by different EPICS Base versions.
func useSenderIP(b []byte) bool {
	allFF := true
	allZero := true
	for _, v := range b {
		if v != 0xFF {
			allFF = false
		}
		if v != 0x00 {
			allZero = false
		}
	}
	return allFF || allZero
}

func isAllFF(b []byte) bool {
	for _, v := range b {
		if v != 0xFF {
			return false
		}
	}
	return true
}

// -------------------------------------------------------------------------- //
// Address list helpers                                                        //
// -------------------------------------------------------------------------- //

// resolveAddrs converts a list of "host" or "host:port" strings to
// "host:port" strings using defaultPort when no port is specified.
func resolveAddrs(addrs []string, defaultPort int) []string {
	out := make([]string, 0, len(addrs))
	for _, a := range addrs {
		if a == "" {
			continue
		}
		_, _, err := net.SplitHostPort(a)
		if err != nil {
			// No port specified; append default.
			a = fmt.Sprintf("%s:%d", a, defaultPort)
		}
		out = append(out, a)
	}
	return out
}

// localBroadcastAddrs returns broadcast addresses for all local network
// interfaces, used when AutoAddrList is true.
func localBroadcastAddrs(port int) []string {
	ifaces, err := net.Interfaces()
	if err != nil {
		return nil
	}
	var out []string
	for _, iface := range ifaces {
		if iface.Flags&net.FlagBroadcast == 0 {
			continue
		}
		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}
		for _, a := range addrs {
			ipNet, ok := a.(*net.IPNet)
			if !ok {
				continue
			}
			ip4 := ipNet.IP.To4()
			if ip4 == nil {
				continue
			}
			// Compute broadcast: IP | ^mask
			mask := ipNet.Mask
			bcast := make(net.IP, 4)
			for i := range bcast {
				bcast[i] = ip4[i] | ^mask[i]
			}
			out = append(out, fmt.Sprintf("%s:%d", bcast.String(), port))
		}
	}
	return out
}

// -------------------------------------------------------------------------- //
// Minimal bytes.Reader-like for proto.DecodeHeader                           //
// -------------------------------------------------------------------------- //

type bytesReader struct {
	b []byte
	i int
}

func newBytesReader(b []byte) *bytesReader { return &bytesReader{b: b} }

func (r *bytesReader) Read(p []byte) (int, error) {
	if r.i >= len(r.b) {
		return 0, fmt.Errorf("EOF")
	}
	n := copy(p, r.b[r.i:])
	r.i += n
	return n, nil
}

// -------------------------------------------------------------------------- //
// searchAddrFromReply extracts server IP from a CA_PROTO_SEARCH response.   //
// Exported for testability.                                                   //
// -------------------------------------------------------------------------- //

// EncodedSearchPair builds the VERSION + SEARCH UDP datagram pair for pvName
// with the given searchID. Exported for use in tests.
func EncodedSearchPair(pvName string, searchID uint32) []byte {
	verHdr := proto.Header{
		Command:   proto.CmdVersion,
		DataCount: proto.MinorVersion,
	}
	payload := proto.BuildStringPayload(pvName)
	searchHdr := proto.Header{
		Command:    proto.CmdSearch,
		DataType:   proto.SearchReply,
		DataCount:  proto.MinorVersion,
		Parameter1: searchID,
		Parameter2: searchID,
	}
	searchMsg := proto.BuildMessage(searchHdr, payload)
	out := make([]byte, 0, proto.HeaderSize+len(searchMsg))
	out = append(out, verHdr.Bytes()...)
	out = append(out, searchMsg...)
	return out
}

// EncodeSearchReply builds a CA_PROTO_SEARCH UDP reply.
// serverIP may be nil to use 0xFFFFFFFF (sender IP).  port is the CA TCP port.
func EncodeSearchReply(searchID uint32, serverIP net.IP, port int) []byte {
	var payload []byte
	if serverIP != nil {
		payload = make([]byte, 8)
		copy(payload[:4], serverIP.To4())
		binary.BigEndian.PutUint32(payload[4:], 0)
	} else {
		payload = []byte{0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x00}
	}
	h := proto.Header{
		Command:    proto.CmdSearch,
		DataType:   uint16(port),
		Parameter1: proto.MinorVersion,
		Parameter2: searchID,
	}
	return proto.BuildMessage(h, payload)
}