Render hostmaps as mermaid graphs in e2e tests (#815)

e2e/router/router.go

@@ -40,7 +40,12 @@ type R struct {
 	// A map of vpn ip to the nebula control it belongs to
 	vpnControls map[iputil.VpnIp]*nebula.Control
 
-	flow []flowEntry
+	ignoreFlows []ignoreFlow
+	flow        []flowEntry
+
+	// A set of additional mermaid graphs to draw in the flow log markdown file
+	// Currently consisting only of hostmap renders
+	additionalGraphs []mermaidGraph
 
 	// All interactions are locked to help serialize behavior
 	sync.Mutex
@@ -50,6 +55,24 @@ type R struct {
 	t            testing.TB
 }
 
+type ignoreFlow struct {
+	tun         NullBool
+	messageType header.MessageType
+	subType     header.MessageSubType
+	//from
+	//to
+}
+
+type mermaidGraph struct {
+	title   string
+	content string
+}
+
+type NullBool struct {
+	HasValue bool
+	IsTrue   bool
+}
+
 type flowEntry struct {
 	note   string
 	packet *packet
@@ -98,6 +121,7 @@ func NewR(t testing.TB, controls ...*nebula.Control) *R {
 		inNat:        make(map[string]*nebula.Control),
 		outNat:       make(map[string]net.UDPAddr),
 		flow:         []flowEntry{},
+		ignoreFlows:  []ignoreFlow{},
 		fn:           filepath.Join("mermaid", fmt.Sprintf("%s.md", t.Name())),
 		t:            t,
 		cancelRender: cancel,
@@ -219,15 +243,55 @@ func (r *R) renderFlow() {
 			}
 
 			fmt.Fprintf(f,
-				"    %s%s%s: %s(%s), counter: %v\n",
+				"    %s%s%s: %s(%s), index %v, counter: %v\n",
 				strings.Replace(p.from.GetUDPAddr(), ":", "#58;", 1),
 				line,
 				strings.Replace(p.to.GetUDPAddr(), ":", "#58;", 1),
-				h.TypeName(), h.SubTypeName(), h.MessageCounter,
+				h.TypeName(), h.SubTypeName(), h.RemoteIndex, h.MessageCounter,
 			)
 		}
 	}
 	fmt.Fprintln(f, "```")
+
+	for _, g := range r.additionalGraphs {
+		fmt.Fprintf(f, "## %s\n", g.title)
+		fmt.Fprintln(f, "```mermaid")
+		fmt.Fprintln(f, g.content)
+		fmt.Fprintln(f, "```")
+	}
+}
+
+// IgnoreFlow tells the router to stop recording future flows that matches the provided criteria.
+// messageType and subType will target nebula underlay packets while tun will target nebula overlay packets
+// NOTE: This is a very broad system, if you set tun to true then no more tun traffic will be rendered
+func (r *R) IgnoreFlow(messageType header.MessageType, subType header.MessageSubType, tun NullBool) {
+	r.Lock()
+	defer r.Unlock()
+	r.ignoreFlows = append(r.ignoreFlows, ignoreFlow{
+		tun,
+		messageType,
+		subType,
+	})
+}
+
+func (r *R) RenderHostmaps(title string, controls ...*nebula.Control) {
+	r.Lock()
+	defer r.Unlock()
+
+	s := renderHostmaps(controls...)
+	if len(r.additionalGraphs) > 0 {
+		lastGraph := r.additionalGraphs[len(r.additionalGraphs)-1]
+		if lastGraph.content == s && lastGraph.title == title {
+			// Ignore this rendering if it matches the last rendering added
+			// This is useful if you want to track rendering changes
+			return
+		}
+	}
+
+	r.additionalGraphs = append(r.additionalGraphs, mermaidGraph{
+		title:   title,
+		content: s,
+	})
 }
 
 // InjectFlow can be used to record packet flow if the test is handling the routing on its own.
@@ -268,6 +332,24 @@ func (r *R) unlockedInjectFlow(from, to *nebula.Control, p *udp.Packet, tun bool
 		return nil
 	}
 
+	if len(r.ignoreFlows) > 0 {
+		var h header.H
+		err := h.Parse(p.Data)
+		if err != nil {
+			panic(err)
+		}
+
+		for _, i := range r.ignoreFlows {
+			if !tun {
+				if i.messageType == h.Type && i.subType == h.Subtype {
+					return nil
+				}
+			} else if i.tun.HasValue && i.tun.IsTrue {
+				return nil
+			}
+		}
+	}
+
 	fp := &packet{
 		from:   from,
 		to:     to,