router.go 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999
  1. //go:build e2e_testing
  2. // +build e2e_testing
  3. package router
  4. import (
  5. "context"
  6. "fmt"
  7. "net/netip"
  8. "os"
  9. "path/filepath"
  10. "reflect"
  11. "regexp"
  12. "sort"
  13. "sync"
  14. "sync/atomic"
  15. "testing"
  16. "time"
  17. "github.com/google/gopacket"
  18. "github.com/google/gopacket/layers"
  19. "github.com/slackhq/nebula"
  20. "github.com/slackhq/nebula/header"
  21. "github.com/slackhq/nebula/udp"
  22. "golang.org/x/exp/maps"
  23. )
  24. // outNatKey is the (from, to) pair used by outNat. Comparable struct, so it works as a map key without the
  25. // allocation cost of a string-concat key.
  26. type outNatKey struct {
  27. from, to netip.AddrPort
  28. }
  29. // fannedPacket pairs a UDP TX packet with its source control so the router can route it after popping from
  30. // the fan-in channel.
  31. type fannedPacket struct {
  32. from *nebula.Control
  33. pkt *udp.Packet
  34. }
  35. type R struct {
  36. // Simple map of the ip:port registered on a control to the control
  37. // Basically a router, right?
  38. controls map[netip.AddrPort]*nebula.Control
  39. // A map for inbound packets for a control that doesn't know about this address
  40. inNat map[netip.AddrPort]*nebula.Control
  41. // A last used map, if an inbound packet hit the inNat map then
  42. // all return packets should use the same last used inbound address for the outbound sender
  43. outNat map[outNatKey]netip.AddrPort
  44. // A map of vpn ip to the nebula control it belongs to
  45. vpnControls map[netip.Addr]*nebula.Control
  46. // Cached select infrastructure for RouteForAllUntilTxTun.
  47. // The controls map is immutable after NewR so the cases are good for the test lifetime.
  48. // We only rebuild if a different receiver is asked.
  49. selRecvCtl *nebula.Control
  50. selCases []reflect.SelectCase
  51. selCtls []*nebula.Control
  52. // Optional fan-in mode for hot-path benchmarks: one forwarder goroutine per control drains UDP TX into udpFanIn,
  53. // so RouteForAllUntilTxTun can do a fixed 2-way native select instead of paying reflect.Select per call.
  54. // Off by default (would otherwise interleave with tests that use GetFromUDP directly on the same control).
  55. // Enabled by EnableFanIn.
  56. udpFanIn chan fannedPacket
  57. stopFanIn chan struct{}
  58. fanInWG sync.WaitGroup
  59. fanInMu sync.Mutex
  60. fanInOn atomic.Bool
  61. ignoreFlows []ignoreFlow
  62. flow []flowEntry
  63. // A set of additional mermaid graphs to draw in the flow log markdown file
  64. // Currently consisting only of hostmap renders
  65. additionalGraphs []mermaidGraph
  66. // All interactions are locked to help serialize behavior
  67. sync.Mutex
  68. fn string
  69. cancelRender context.CancelFunc
  70. t testing.TB
  71. }
  72. type ignoreFlow struct {
  73. tun NullBool
  74. messageType header.MessageType
  75. subType header.MessageSubType
  76. //from
  77. //to
  78. }
  79. type mermaidGraph struct {
  80. title string
  81. content string
  82. }
  83. type NullBool struct {
  84. HasValue bool
  85. IsTrue bool
  86. }
  87. type flowEntry struct {
  88. note string
  89. packet *packet
  90. }
  91. type packet struct {
  92. from *nebula.Control
  93. to *nebula.Control
  94. packet *udp.Packet
  95. tun bool // a packet pulled off a tun device
  96. rx bool // the packet was received by a udp device
  97. // h is the nebula header, parsed once when the packet is recorded. parseErr says why there isn't one, which
  98. // the flow log reports rather than hiding. Punchy sends a single byte, so an unparseable packet is normal.
  99. h header.H
  100. parseErr error
  101. }
  102. // fromAddr and toAddr are the addresses this packet actually travelled between. Reading them off the control
  103. // instead would misreport the whole history once a test moves a node. Tun packets are synthesized without
  104. // addresses, so they fall back to the control.
  105. func (p *packet) fromAddr() netip.AddrPort {
  106. if p.tun || !p.packet.From.IsValid() {
  107. return p.from.GetUDPAddr()
  108. }
  109. return p.packet.From
  110. }
  111. func (p *packet) toAddr() netip.AddrPort {
  112. if p.tun || !p.packet.To.IsValid() {
  113. return p.to.GetUDPAddr()
  114. }
  115. return p.packet.To
  116. }
  117. func (p *packet) WasReceived() {
  118. if p != nil {
  119. p.rx = true
  120. }
  121. }
  122. type ExitType int
  123. const (
  124. // KeepRouting the function will get called again on the next packet
  125. KeepRouting ExitType = 0
  126. // ExitNow does not route this packet and exits immediately
  127. ExitNow ExitType = 1
  128. // RouteAndExit routes this packet and exits immediately afterwards
  129. RouteAndExit ExitType = 2
  130. )
  131. type ExitFunc func(packet *udp.Packet, receiver *nebula.Control) ExitType
  132. // NewR creates a new router to pass packets in a controlled fashion between the provided controllers.
  133. // The packet flow will be recorded in a file within the mermaid directory under the same name as the test.
  134. // Renders will occur automatically, roughly every 100ms, until a call to RenderFlow() is made
  135. func NewR(t testing.TB, controls ...*nebula.Control) *R {
  136. ctx, cancel := context.WithCancel(context.Background())
  137. if err := os.MkdirAll("mermaid", 0755); err != nil {
  138. panic(err)
  139. }
  140. r := &R{
  141. controls: make(map[netip.AddrPort]*nebula.Control),
  142. vpnControls: make(map[netip.Addr]*nebula.Control),
  143. inNat: make(map[netip.AddrPort]*nebula.Control),
  144. outNat: make(map[outNatKey]netip.AddrPort),
  145. flow: []flowEntry{},
  146. ignoreFlows: []ignoreFlow{},
  147. fn: filepath.Join("mermaid", fmt.Sprintf("%s.md", t.Name())),
  148. t: t,
  149. cancelRender: cancel,
  150. }
  151. // Try to remove our render file
  152. os.Remove(r.fn)
  153. for _, c := range controls {
  154. addr := c.GetUDPAddr()
  155. if _, ok := r.controls[addr]; ok {
  156. panic("Duplicate listen address: " + addr.String())
  157. }
  158. for _, vpnAddr := range c.GetVpnAddrs() {
  159. r.vpnControls[vpnAddr] = c
  160. }
  161. r.controls[addr] = c
  162. }
  163. // Spin the renderer in case we go nuts and the test never completes
  164. go func() {
  165. clockSource := time.NewTicker(time.Millisecond * 100)
  166. defer clockSource.Stop()
  167. for {
  168. select {
  169. case <-ctx.Done():
  170. return
  171. case <-clockSource.C:
  172. r.Lock()
  173. r.renderHostmaps("clock tick")
  174. r.renderFlow()
  175. r.Unlock()
  176. }
  177. }
  178. }()
  179. return r
  180. }
  181. // AddRoute will place the nebula controller at the ip and port specified.
  182. // It does not look at the addr attached to the instance.
  183. // If a route is used, this will behave like a NAT for the return path.
  184. // Rewriting the source ip:port to what was last sent to from the origin
  185. func (r *R) AddRoute(ip netip.Addr, port uint16, c *nebula.Control) {
  186. r.Lock()
  187. defer r.Unlock()
  188. inAddr := netip.AddrPortFrom(ip, port)
  189. if _, ok := r.inNat[inAddr]; ok {
  190. panic("Duplicate listen address inNat: " + inAddr.String())
  191. }
  192. r.inNat[inAddr] = c
  193. }
  194. // RenderFlow renders the packet flow seen up until now and stops further automatic renders from happening.
  195. func (r *R) RenderFlow() {
  196. r.cancelRender()
  197. r.Lock()
  198. defer r.Unlock()
  199. r.renderFlow()
  200. }
  201. // CancelFlowLogs stops flow logs from being tracked and destroys any logs already collected
  202. func (r *R) CancelFlowLogs() {
  203. r.cancelRender()
  204. r.Lock()
  205. r.flow = nil
  206. r.Unlock()
  207. }
  208. // renderFlow writes the flow log to disk. Caller must hold r.Lock. renderFlow reads r.flow / r.additionalGraphs and
  209. // the *packet pointers stashed inside, all of which are mutated under the same lock by routing paths.
  210. func (r *R) renderFlow() {
  211. if r.flow == nil {
  212. return
  213. }
  214. f, err := os.OpenFile(r.fn, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0644)
  215. if err != nil {
  216. panic(err)
  217. }
  218. var participants = map[netip.AddrPort]struct{}{}
  219. var participantsVals []string
  220. fmt.Fprintln(f, "```mermaid")
  221. fmt.Fprintln(f, "sequenceDiagram")
  222. // Assemble participants
  223. for _, e := range r.flow {
  224. if e.packet == nil {
  225. continue
  226. }
  227. addr := e.packet.fromAddr()
  228. if _, ok := participants[addr]; ok {
  229. continue
  230. }
  231. participants[addr] = struct{}{}
  232. sanAddr := normalizeName(addr.String())
  233. participantsVals = append(participantsVals, sanAddr)
  234. fmt.Fprintf(
  235. f, " participant %s as Nebula: %s<br/>UDP: %s\n",
  236. sanAddr, e.packet.from.GetVpnAddrs(), sanAddr,
  237. )
  238. }
  239. if len(participantsVals) > 2 {
  240. // Get the first and last participantVals for notes
  241. participantsVals = []string{participantsVals[0], participantsVals[len(participantsVals)-1]}
  242. }
  243. // Print packets
  244. for _, e := range r.flow {
  245. if e.packet == nil {
  246. //fmt.Fprintf(f, " note over %s: %s\n", strings.Join(participantsVals, ", "), e.note)
  247. continue
  248. }
  249. p := e.packet
  250. if p.tun {
  251. fmt.Fprintln(f, r.formatUdpPacket(p))
  252. } else {
  253. line := "--x"
  254. if p.rx {
  255. line = "->>"
  256. }
  257. detail := fmt.Sprintf("%s(%s), index %v, counter: %v",
  258. p.h.TypeName(), p.h.SubTypeName(), p.h.RemoteIndex, p.h.MessageCounter)
  259. if p.parseErr != nil {
  260. detail = fmt.Sprintf("unparsed, %v (%d bytes)", p.parseErr, len(p.packet.Data))
  261. }
  262. fmt.Fprintf(f, " %s%s%s: %s\n",
  263. normalizeName(p.fromAddr().String()),
  264. line,
  265. normalizeName(p.toAddr().String()),
  266. detail,
  267. )
  268. }
  269. }
  270. fmt.Fprintln(f, "```")
  271. for _, g := range r.additionalGraphs {
  272. fmt.Fprintf(f, "## %s\n", g.title)
  273. fmt.Fprintln(f, "```mermaid")
  274. fmt.Fprintln(f, g.content)
  275. fmt.Fprintln(f, "```")
  276. }
  277. }
  278. func normalizeName(s string) string {
  279. rx := regexp.MustCompile("[\\[\\]\\:]")
  280. return rx.ReplaceAllLiteralString(s, "_")
  281. }
  282. // IgnoreFlow tells the router to stop recording future flows that matches the provided criteria.
  283. // messageType and subType will target nebula underlay packets while tun will target nebula overlay packets
  284. // NOTE: This is a very broad system, if you set tun to true then no more tun traffic will be rendered
  285. func (r *R) IgnoreFlow(messageType header.MessageType, subType header.MessageSubType, tun NullBool) {
  286. r.Lock()
  287. defer r.Unlock()
  288. r.ignoreFlows = append(r.ignoreFlows, ignoreFlow{
  289. tun,
  290. messageType,
  291. subType,
  292. })
  293. }
  294. func (r *R) RenderHostmaps(title string, controls ...*nebula.Control) {
  295. r.Lock()
  296. defer r.Unlock()
  297. s := renderHostmaps(controls...)
  298. if len(r.additionalGraphs) > 0 {
  299. lastGraph := r.additionalGraphs[len(r.additionalGraphs)-1]
  300. if lastGraph.content == s && lastGraph.title == title {
  301. // Ignore this rendering if it matches the last rendering added
  302. // This is useful if you want to track rendering changes
  303. return
  304. }
  305. }
  306. r.additionalGraphs = append(r.additionalGraphs, mermaidGraph{
  307. title: title,
  308. content: s,
  309. })
  310. }
  311. func (r *R) renderHostmaps(title string) {
  312. c := maps.Values(r.controls)
  313. sort.SliceStable(c, func(i, j int) bool {
  314. return c[i].GetVpnAddrs()[0].Compare(c[j].GetVpnAddrs()[0]) > 0
  315. })
  316. s := renderHostmaps(c...)
  317. if len(r.additionalGraphs) > 0 {
  318. lastGraph := r.additionalGraphs[len(r.additionalGraphs)-1]
  319. if lastGraph.content == s {
  320. // Ignore this rendering if it matches the last rendering added
  321. // This is useful if you want to track rendering changes
  322. return
  323. }
  324. }
  325. r.additionalGraphs = append(r.additionalGraphs, mermaidGraph{
  326. title: title,
  327. content: s,
  328. })
  329. }
  330. // InjectFlow can be used to record packet flow if the test is handling the routing on its own.
  331. // The packet is assumed to have been received
  332. func (r *R) InjectFlow(from, to *nebula.Control, p *udp.Packet) {
  333. r.Lock()
  334. defer r.Unlock()
  335. r.unlockedInjectFlow(from, to, p, false)
  336. }
  337. func (r *R) Log(arg ...any) {
  338. if r.flow == nil {
  339. return
  340. }
  341. r.Lock()
  342. r.flow = append(r.flow, flowEntry{note: fmt.Sprint(arg...)})
  343. r.t.Log(arg...)
  344. r.Unlock()
  345. }
  346. func (r *R) Logf(format string, arg ...any) {
  347. if r.flow == nil {
  348. return
  349. }
  350. r.Lock()
  351. r.flow = append(r.flow, flowEntry{note: fmt.Sprintf(format, arg...)})
  352. r.t.Logf(format, arg...)
  353. r.Unlock()
  354. }
  355. // unlockedInjectFlow is used by the router to record a packet has been transmitted, the packet is returned and
  356. // should be marked as received AFTER it has been placed on the receivers channel.
  357. // If flow logs have been disabled this function will return nil
  358. func (r *R) unlockedInjectFlow(from, to *nebula.Control, p *udp.Packet, tun bool) *packet {
  359. if r.flow == nil {
  360. return nil
  361. }
  362. r.renderHostmaps(fmt.Sprintf("Packet %v", len(r.flow)))
  363. var h header.H
  364. var parseErr error
  365. if !tun {
  366. parseErr = h.Parse(p.Data)
  367. }
  368. // Decide before copying, the copy comes from a freelist and an ignored packet would never be released
  369. for _, i := range r.ignoreFlows {
  370. if tun {
  371. if i.tun.HasValue && i.tun.IsTrue {
  372. return nil
  373. }
  374. continue
  375. }
  376. // A packet we could not parse has no type to match against, so no rule can ignore it
  377. if parseErr == nil && i.messageType == h.Type && i.subType == h.Subtype {
  378. return nil
  379. }
  380. }
  381. fp := &packet{
  382. from: from,
  383. to: to,
  384. packet: p.Copy(),
  385. tun: tun,
  386. h: h,
  387. parseErr: parseErr,
  388. }
  389. r.flow = append(r.flow, flowEntry{packet: fp})
  390. return fp
  391. }
  392. // OnceFrom will route a single packet from sender then return
  393. // If the router doesn't have the nebula controller for that address, we panic
  394. func (r *R) OnceFrom(sender *nebula.Control) {
  395. r.RouteExitFunc(sender, func(*udp.Packet, *nebula.Control) ExitType {
  396. return RouteAndExit
  397. })
  398. }
  399. // RouteUntilTxTun will route for sender and return when a packet is seen on receivers tun
  400. // If the router doesn't have the nebula controller for that address, we panic
  401. func (r *R) RouteUntilTxTun(sender *nebula.Control, receiver *nebula.Control) []byte {
  402. tunTx := receiver.GetTunTxChan()
  403. udpTx := sender.GetUDPTxChan()
  404. for {
  405. select {
  406. // Maybe we already have something on the tun for us
  407. case b := <-tunTx:
  408. r.Lock()
  409. np := udp.Packet{Data: make([]byte, len(b))}
  410. copy(np.Data, b)
  411. r.unlockedInjectFlow(receiver, receiver, &np, true)
  412. r.Unlock()
  413. return b
  414. // Nope, lets push the sender along
  415. case p := <-udpTx:
  416. r.Lock()
  417. a := sender.GetUDPAddr()
  418. c := r.getControl(a, p.To, p)
  419. if c == nil {
  420. r.Unlock()
  421. panic("No control for udp tx " + a.String())
  422. }
  423. fp := r.unlockedInjectFlow(sender, c, p, false)
  424. c.InjectUDPPacket(p) // copies internally; original is ours to release
  425. fp.WasReceived()
  426. r.Unlock()
  427. p.Release()
  428. }
  429. }
  430. }
  431. // RouteForAllUntilTxTun will route for everyone and return when a packet is seen on the receiver's tun.
  432. // If a control's UDP TX address can't be matched to a registered control, we panic.
  433. //
  434. // For allocation-sensitive callers (hot-path benchmarks, in particular relay
  435. // benches with 3+ controls), call EnableFanIn() first.
  436. func (r *R) RouteForAllUntilTxTun(receiver *nebula.Control) []byte {
  437. if r.fanInOn.Load() {
  438. return r.routeFanIn(receiver)
  439. }
  440. return r.routeReflect(receiver)
  441. }
  442. // routeFanIn is the alloc-free path used when EnableFanIn is in effect.
  443. func (r *R) routeFanIn(receiver *nebula.Control) []byte {
  444. tunTx := receiver.GetTunTxChan()
  445. for {
  446. select {
  447. case p := <-tunTx:
  448. r.Lock()
  449. if r.flow != nil {
  450. np := udp.Packet{Data: make([]byte, len(p))}
  451. copy(np.Data, p)
  452. r.unlockedInjectFlow(receiver, receiver, &np, true)
  453. }
  454. r.Unlock()
  455. return p
  456. case fp := <-r.udpFanIn:
  457. r.routeUDP(fp.from, fp.pkt)
  458. }
  459. }
  460. }
  461. // routeReflect is the default reflect.Select-based path. Pays the boxing allocation per call but doesn't interfere
  462. // with tests that pull packets directly from controls' UDP TX channels via GetFromUDP.
  463. func (r *R) routeReflect(receiver *nebula.Control) []byte {
  464. sc, cm := r.selectCasesFor(receiver)
  465. for {
  466. x, rx, _ := reflect.Select(sc)
  467. if x == 0 {
  468. p := rx.Interface().([]byte)
  469. r.Lock()
  470. if r.flow != nil {
  471. np := udp.Packet{Data: make([]byte, len(p))}
  472. copy(np.Data, p)
  473. r.unlockedInjectFlow(cm[x], cm[x], &np, true)
  474. }
  475. r.Unlock()
  476. return p
  477. }
  478. r.routeUDP(cm[x], rx.Interface().(*udp.Packet))
  479. }
  480. }
  481. // EnableFanIn switches RouteForAllUntilTxTun to the alloc-free fan-in path.
  482. // One forwarder goroutine per registered control drains UDP TX into a shared channel that RouteForAllUntilTxTun selects
  483. // on alongside the receiver's TUN TX channel.
  484. func (r *R) EnableFanIn() {
  485. r.fanInMu.Lock()
  486. defer r.fanInMu.Unlock()
  487. if r.fanInOn.Load() {
  488. return
  489. }
  490. r.udpFanIn = make(chan fannedPacket, 32)
  491. r.stopFanIn = make(chan struct{})
  492. for _, c := range r.controls {
  493. r.startFanInWorker(c)
  494. }
  495. r.fanInOn.Store(true)
  496. r.t.Cleanup(r.stopFanInWorkers)
  497. }
  498. // startFanInWorker spawns a goroutine that drains c's UDP TX into r.udpFanIn.
  499. func (r *R) startFanInWorker(c *nebula.Control) {
  500. r.fanInWG.Add(1)
  501. udpTx := c.GetUDPTxChan()
  502. go func() {
  503. defer r.fanInWG.Done()
  504. for {
  505. select {
  506. case <-r.stopFanIn:
  507. return
  508. case p := <-udpTx:
  509. select {
  510. case <-r.stopFanIn:
  511. p.Release()
  512. return
  513. case r.udpFanIn <- fannedPacket{from: c, pkt: p}:
  514. }
  515. }
  516. }
  517. }()
  518. }
  519. // stopFanInWorkers signals the fan-in goroutines to exit and waits for them.
  520. func (r *R) stopFanInWorkers() {
  521. r.fanInMu.Lock()
  522. wasOn := r.fanInOn.Swap(false)
  523. r.fanInMu.Unlock()
  524. if !wasOn {
  525. return
  526. }
  527. close(r.stopFanIn)
  528. r.fanInWG.Wait()
  529. }
  530. // routeUDP forwards a UDP TX packet from the named source control to the destination control derived from p.To,
  531. // releasing the source packet after InjectUDPPacket has copied its bytes into a fresh pool slot.
  532. func (r *R) routeUDP(from *nebula.Control, p *udp.Packet) {
  533. r.Lock()
  534. defer r.Unlock()
  535. a := from.GetUDPAddr()
  536. c := r.getControl(a, p.To, p)
  537. if c == nil {
  538. panic(fmt.Sprintf("No control for udp tx %s", p.To))
  539. }
  540. fp := r.unlockedInjectFlow(from, c, p, false)
  541. c.InjectUDPPacket(p) // copies internally; original is ours to release
  542. fp.WasReceived()
  543. p.Release()
  544. }
  545. // selectCasesFor returns the SelectCase array used by routeReflect: one slot for the receiver's TUN TX channel followed
  546. // by one per control's UDP TX channel. Cached for the test lifetime, only rebuilt if the receiver changes.
  547. func (r *R) selectCasesFor(receiver *nebula.Control) ([]reflect.SelectCase, []*nebula.Control) {
  548. r.Lock()
  549. defer r.Unlock()
  550. if r.selRecvCtl == receiver && r.selCases != nil {
  551. return r.selCases, r.selCtls
  552. }
  553. sc := make([]reflect.SelectCase, len(r.controls)+1)
  554. cm := make([]*nebula.Control, len(r.controls)+1)
  555. sc[0] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(receiver.GetTunTxChan())}
  556. cm[0] = receiver
  557. i := 1
  558. for _, c := range r.controls {
  559. sc[i] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(c.GetUDPTxChan())}
  560. cm[i] = c
  561. i++
  562. }
  563. r.selRecvCtl = receiver
  564. r.selCases = sc
  565. r.selCtls = cm
  566. return sc, cm
  567. }
  568. // RouteExitFunc will call the whatDo func with each udp packet from sender.
  569. // whatDo can return:
  570. // - exitNow: the packet will not be routed and this call will return immediately
  571. // - routeAndExit: this call will return immediately after routing the last packet from sender
  572. // - keepRouting: the packet will be routed and whatDo will be called again on the next packet from sender
  573. func (r *R) RouteExitFunc(sender *nebula.Control, whatDo ExitFunc) {
  574. h := &header.H{}
  575. for {
  576. p := sender.GetFromUDP(true)
  577. r.Lock()
  578. if err := h.Parse(p.Data); err != nil {
  579. panic(err)
  580. }
  581. receiver := r.getControl(sender.GetUDPAddr(), p.To, p)
  582. if receiver == nil {
  583. r.Unlock()
  584. panic("Can't RouteExitFunc for host: " + p.To.String())
  585. }
  586. e := whatDo(p, receiver)
  587. switch e {
  588. case ExitNow:
  589. r.Unlock()
  590. p.Release()
  591. return
  592. case RouteAndExit:
  593. fp := r.unlockedInjectFlow(sender, receiver, p, false)
  594. receiver.InjectUDPPacket(p)
  595. fp.WasReceived()
  596. r.Unlock()
  597. p.Release()
  598. return
  599. case KeepRouting:
  600. fp := r.unlockedInjectFlow(sender, receiver, p, false)
  601. receiver.InjectUDPPacket(p)
  602. fp.WasReceived()
  603. default:
  604. panic(fmt.Sprintf("Unknown exitFunc return: %v", e))
  605. }
  606. r.Unlock()
  607. p.Release()
  608. }
  609. }
  610. // RouteUntilAfterMsgType will route for sender until a message type is seen and sent from sender
  611. // If the router doesn't have the nebula controller for that address, we panic
  612. func (r *R) RouteUntilAfterMsgType(sender *nebula.Control, msgType header.MessageType, subType header.MessageSubType) {
  613. h := &header.H{}
  614. r.RouteExitFunc(sender, func(p *udp.Packet, r *nebula.Control) ExitType {
  615. if err := h.Parse(p.Data); err != nil {
  616. panic(err)
  617. }
  618. if h.Type == msgType && h.Subtype == subType {
  619. return RouteAndExit
  620. }
  621. return KeepRouting
  622. })
  623. }
  624. // RouteFor routes everything that shows up for the given duration and then returns. Use it to let a test settle
  625. // deterministically rather than sleeping and hoping: a single FlushAll races a completing handshake, which queues
  626. // more packets right behind it.
  627. func (r *R) RouteFor(d time.Duration) {
  628. r.RouteForAllExitFuncOrTimeout(d, func(*udp.Packet, *nebula.Control) ExitType {
  629. return KeepRouting
  630. })
  631. }
  632. // RouteForAllExitFuncOrTimeout is RouteForAllExitFunc with a deadline, reporting whether whatDo asked to exit
  633. // before time ran out. The unbounded version blocks forever on a quiet network, so this is what a test needs to
  634. // assert that something does NOT happen, or to route for a fixed settling period.
  635. func (r *R) RouteForAllExitFuncOrTimeout(timeout time.Duration, whatDo ExitFunc) bool {
  636. sc := make([]reflect.SelectCase, 0, len(r.controls)+1)
  637. cm := make([]*nebula.Control, 0, len(r.controls))
  638. for _, c := range r.controls {
  639. sc = append(sc, reflect.SelectCase{
  640. Dir: reflect.SelectRecv,
  641. Chan: reflect.ValueOf(c.GetUDPTxChan()),
  642. Send: reflect.Value{},
  643. })
  644. cm = append(cm, c)
  645. }
  646. timer := time.NewTimer(timeout)
  647. defer timer.Stop()
  648. sc = append(sc, reflect.SelectCase{
  649. Dir: reflect.SelectRecv,
  650. Chan: reflect.ValueOf(timer.C),
  651. Send: reflect.Value{},
  652. })
  653. for {
  654. x, rx, _ := reflect.Select(sc)
  655. if x == len(cm) {
  656. return false
  657. }
  658. r.Lock()
  659. p := rx.Interface().(*udp.Packet)
  660. receiver := r.getControl(cm[x].GetUDPAddr(), p.To, p)
  661. if receiver == nil {
  662. r.Unlock()
  663. panic("Can't RouteForAllExitFuncOrTimeout for host: " + p.To.String())
  664. }
  665. e := whatDo(p, receiver)
  666. switch e {
  667. case ExitNow:
  668. r.Unlock()
  669. p.Release()
  670. return true
  671. case RouteAndExit:
  672. fp := r.unlockedInjectFlow(cm[x], receiver, p, false)
  673. receiver.InjectUDPPacket(p)
  674. fp.WasReceived()
  675. r.Unlock()
  676. p.Release()
  677. return true
  678. case KeepRouting:
  679. fp := r.unlockedInjectFlow(cm[x], receiver, p, false)
  680. receiver.InjectUDPPacket(p)
  681. fp.WasReceived()
  682. default:
  683. panic(fmt.Sprintf("Unknown exitFunc return: %v", e))
  684. }
  685. r.Unlock()
  686. p.Release()
  687. }
  688. }
  689. func (r *R) RouteForAllUntilAfterMsgTypeTo(receiver *nebula.Control, msgType header.MessageType, subType header.MessageSubType) {
  690. h := &header.H{}
  691. r.RouteForAllExitFunc(func(p *udp.Packet, r *nebula.Control) ExitType {
  692. if r != receiver {
  693. return KeepRouting
  694. }
  695. if err := h.Parse(p.Data); err != nil {
  696. panic(err)
  697. }
  698. if h.Type == msgType && h.Subtype == subType {
  699. return RouteAndExit
  700. }
  701. return KeepRouting
  702. })
  703. }
  704. func (r *R) InjectUDPPacket(sender, receiver *nebula.Control, packet *udp.Packet) {
  705. r.Lock()
  706. defer r.Unlock()
  707. fp := r.unlockedInjectFlow(sender, receiver, packet, false)
  708. receiver.InjectUDPPacket(packet)
  709. fp.WasReceived()
  710. }
  711. // RouteForUntilAfterToAddr will route for sender and return only after it sees and sends a packet destined for toAddr
  712. // finish can be any of the exitType values except `keepRouting`, the default value is `routeAndExit`
  713. // If the router doesn't have the nebula controller for that address, we panic
  714. func (r *R) RouteForUntilAfterToAddr(sender *nebula.Control, toAddr netip.AddrPort, finish ExitType) {
  715. if finish == KeepRouting {
  716. finish = RouteAndExit
  717. }
  718. r.RouteExitFunc(sender, func(p *udp.Packet, r *nebula.Control) ExitType {
  719. if p.To == toAddr {
  720. return finish
  721. }
  722. return KeepRouting
  723. })
  724. }
  725. // RouteForAllExitFunc will route for every registered controller and calls the whatDo func with each udp packet from
  726. // whatDo can return:
  727. // - exitNow: the packet will not be routed and this call will return immediately
  728. // - routeAndExit: this call will return immediately after routing the last packet from sender
  729. // - keepRouting: the packet will be routed and whatDo will be called again on the next packet from sender
  730. func (r *R) RouteForAllExitFunc(whatDo ExitFunc) {
  731. sc := make([]reflect.SelectCase, len(r.controls))
  732. cm := make([]*nebula.Control, len(r.controls))
  733. i := 0
  734. for _, c := range r.controls {
  735. sc[i] = reflect.SelectCase{
  736. Dir: reflect.SelectRecv,
  737. Chan: reflect.ValueOf(c.GetUDPTxChan()),
  738. Send: reflect.Value{},
  739. }
  740. cm[i] = c
  741. i++
  742. }
  743. for {
  744. x, rx, _ := reflect.Select(sc)
  745. r.Lock()
  746. p := rx.Interface().(*udp.Packet)
  747. receiver := r.getControl(cm[x].GetUDPAddr(), p.To, p)
  748. if receiver == nil {
  749. r.Unlock()
  750. panic("Can't RouteForAllExitFunc for host: " + p.To.String())
  751. }
  752. e := whatDo(p, receiver)
  753. switch e {
  754. case ExitNow:
  755. r.Unlock()
  756. p.Release()
  757. return
  758. case RouteAndExit:
  759. fp := r.unlockedInjectFlow(cm[x], receiver, p, false)
  760. receiver.InjectUDPPacket(p)
  761. fp.WasReceived()
  762. r.Unlock()
  763. p.Release()
  764. return
  765. case KeepRouting:
  766. fp := r.unlockedInjectFlow(cm[x], receiver, p, false)
  767. receiver.InjectUDPPacket(p)
  768. fp.WasReceived()
  769. default:
  770. panic(fmt.Sprintf("Unknown exitFunc return: %v", e))
  771. }
  772. r.Unlock()
  773. p.Release()
  774. }
  775. }
  776. // FlushAll will route for every registered controller, exiting once there are no packets left to route
  777. func (r *R) FlushAll() {
  778. sc := make([]reflect.SelectCase, len(r.controls))
  779. cm := make([]*nebula.Control, len(r.controls))
  780. i := 0
  781. for _, c := range r.controls {
  782. sc[i] = reflect.SelectCase{
  783. Dir: reflect.SelectRecv,
  784. Chan: reflect.ValueOf(c.GetUDPTxChan()),
  785. Send: reflect.Value{},
  786. }
  787. cm[i] = c
  788. i++
  789. }
  790. // Add a default case to exit when nothing is left to send
  791. sc = append(sc, reflect.SelectCase{
  792. Dir: reflect.SelectDefault,
  793. Chan: reflect.Value{},
  794. Send: reflect.Value{},
  795. })
  796. for {
  797. x, rx, ok := reflect.Select(sc)
  798. if !ok {
  799. return
  800. }
  801. r.Lock()
  802. p := rx.Interface().(*udp.Packet)
  803. receiver := r.getControl(cm[x].GetUDPAddr(), p.To, p)
  804. if receiver == nil {
  805. r.Unlock()
  806. panic("Can't FlushAll for host: " + p.To.String())
  807. }
  808. receiver.InjectUDPPacket(p)
  809. r.Unlock()
  810. p.Release()
  811. }
  812. }
  813. // getControl performs or seeds NAT translation and returns the control for toAddr, p from fields may change
  814. // This is an internal router function, the caller must hold the lock
  815. func (r *R) getControl(fromAddr, toAddr netip.AddrPort, p *udp.Packet) *nebula.Control {
  816. if newAddr, ok := r.outNat[outNatKey{from: fromAddr, to: toAddr}]; ok {
  817. p.From = newAddr
  818. }
  819. c, ok := r.inNat[toAddr]
  820. if ok {
  821. r.outNat[outNatKey{from: c.GetUDPAddr(), to: fromAddr}] = toAddr
  822. return c
  823. }
  824. return r.controls[toAddr]
  825. }
  826. func (r *R) formatUdpPacket(p *packet) string {
  827. var packet gopacket.Packet
  828. var srcAddr netip.Addr
  829. packet = gopacket.NewPacket(p.packet.Data, layers.LayerTypeIPv6, gopacket.Lazy)
  830. if packet.ErrorLayer() == nil {
  831. v6 := packet.Layer(layers.LayerTypeIPv6).(*layers.IPv6)
  832. if v6 == nil {
  833. panic("not an ipv6 packet")
  834. }
  835. srcAddr, _ = netip.AddrFromSlice(v6.SrcIP)
  836. } else {
  837. packet = gopacket.NewPacket(p.packet.Data, layers.LayerTypeIPv4, gopacket.Lazy)
  838. v6 := packet.Layer(layers.LayerTypeIPv4).(*layers.IPv4)
  839. if v6 == nil {
  840. panic("not an ipv6 packet")
  841. }
  842. srcAddr, _ = netip.AddrFromSlice(v6.SrcIP)
  843. }
  844. from := "unknown"
  845. if c, ok := r.vpnControls[srcAddr]; ok {
  846. from = c.GetUDPAddr().String()
  847. }
  848. udpLayer := packet.Layer(layers.LayerTypeUDP).(*layers.UDP)
  849. if udpLayer == nil {
  850. panic("not a udp packet")
  851. }
  852. data := packet.ApplicationLayer()
  853. return fmt.Sprintf(
  854. " %s-->>%s: src port: %v<br/>dest port: %v<br/>data: \"%v\"\n",
  855. normalizeName(from),
  856. normalizeName(p.to.GetUDPAddr().String()),
  857. udpLayer.SrcPort,
  858. udpLayer.DstPort,
  859. string(data.Payload()),
  860. )
  861. }