IPv6 Support for afthelper.NextHopAddrsForPrefix

afthelper/afthelper.go

@@ -18,6 +18,7 @@ package afthelper
 
 import (
 	"fmt"
+	"net/netip"
 
 	"github.com/openconfig/gribigo/aft"
 )
@@ -42,23 +43,41 @@ func NextHopAddrsForPrefix(rib map[string]*aft.RIB, netinst, prefix string) (map
 		return nil, fmt.Errorf("network instance %s does not exist", netinst)
 	}
 
-	v4 := niAFT.GetAfts().GetIpv4Entry(prefix)
-	if v4 == nil {
-		return nil, fmt.Errorf("cannot find IPv4 prefix in AFT")
+	pfx, err := netip.ParsePrefix(prefix)
+	if err != nil {
+		return nil, fmt.Errorf("invalid prefix: %v", err)
+	}
+
+	var otherNI string
+	var nhgID uint64
+	if pfx.Addr().Is4() {
+		v4 := niAFT.GetAfts().GetIpv4Entry(prefix)
+		if v4 == nil {
+			return nil, fmt.Errorf("cannot find IPv4 prefix in AFT")
+		}
+		otherNI = v4.GetNextHopGroupNetworkInstance()
+		nhgID = v4.GetNextHopGroup()
+	} else {
+		v6 := niAFT.GetAfts().GetIpv6Entry(prefix)
+		if v6 == nil {
+			return nil, fmt.Errorf("cannot find IPv6 prefix in AFT")
+		}
+		otherNI = v6.GetNextHopGroupNetworkInstance()
+		nhgID = v6.GetNextHopGroup()
 	}
 
 	nhNI := netinst
-	if otherNI := v4.GetNextHopGroupNetworkInstance(); otherNI != "" {
+	if otherNI != "" {
 		nhNI = otherNI
 	}
 
 	if _, ok := rib[nhNI]; !ok {
 		return nil, fmt.Errorf("got invalid network instance, %s", nhNI)
 	}
 
-	nhg := rib[nhNI].GetAfts().GetNextHopGroup(v4.GetNextHopGroup())
+	nhg := rib[nhNI].GetAfts().GetNextHopGroup(nhgID)
 	if nhg == nil {
-		return nil, fmt.Errorf("got unknown NHG %d in NI %s", v4.GetNextHopGroup(), nhNI)
+		return nil, fmt.Errorf("got unknown NHG %d in NI %s", nhgID, nhNI)
 	}
 
 	// sum is a map of index -> weight.

afthelper/afthelper_test.go

@@ -81,6 +81,33 @@ func TestNextHopAddrsForPrefix(t *testing.T) {
 				NetworkInstance: defName,
 			},
 		},
+	}, {
+		desc: "ipv6 with two NH, one v4 and one v6",
+		inRIB: map[string]*aft.RIB{
+			defName: func() *aft.RIB {
+				r := &aft.RIB{}
+				r.GetOrCreateAfts().GetOrCreateIpv6Entry("2001:aaaa::/64").NextHopGroup = ygot.Uint64(1)
+				r.GetOrCreateAfts().GetOrCreateNextHopGroup(1).GetOrCreateNextHop(1).Weight = ygot.Uint64(1)
+				r.GetOrCreateAfts().GetOrCreateNextHopGroup(1).GetOrCreateNextHop(2).Weight = ygot.Uint64(1)
+				r.GetOrCreateAfts().GetOrCreateNextHop(1).IpAddress = ygot.String("1000:10:10::10")
+				r.GetOrCreateAfts().GetOrCreateNextHop(2).IpAddress = ygot.String("2.2.2.2")
+				return r
+			}(),
+		},
+		inNetInst: defName,
+		inPrefix:  "2001:aaaa::/64",
+		want: map[string]*NextHopSummary{
+			"1000:10:10::10": {
+				Address:         "1000:10:10::10",
+				Weight:          1,
+				NetworkInstance: defName,
+			},
+			"2.2.2.2": {
+				Address:         "2.2.2.2",
+				Weight:          1,
+				NetworkInstance: defName,
+			},
+		},
 	}, {
 		desc:      "can't find network instance",
 		inRIB:     map[string]*aft.RIB{},
@@ -119,6 +146,39 @@ func TestNextHopAddrsForPrefix(t *testing.T) {
 				NetworkInstance: "VRF-1",
 			},
 		},
+	}, {
+		desc: "ipv6 with two NH, one v4 and one v6 in different network instance",
+		inRIB: map[string]*aft.RIB{
+			defName: func() *aft.RIB {
+				r := &aft.RIB{}
+				v6 := r.GetOrCreateAfts().GetOrCreateIpv6Entry("2001:aaaa::/64")
+				v6.NextHopGroup = ygot.Uint64(1)
+				v6.NextHopGroupNetworkInstance = ygot.String("VRF-1")
+				return r
+			}(),
+			"VRF-1": func() *aft.RIB {
+				r := &aft.RIB{}
+				r.GetOrCreateAfts().GetOrCreateNextHopGroup(1).GetOrCreateNextHop(1).Weight = ygot.Uint64(1)
+				r.GetOrCreateAfts().GetOrCreateNextHopGroup(1).GetOrCreateNextHop(2).Weight = ygot.Uint64(1)
+				r.GetOrCreateAfts().GetOrCreateNextHop(1).IpAddress = ygot.String("1000:10:10::10")
+				r.GetOrCreateAfts().GetOrCreateNextHop(2).IpAddress = ygot.String("2.2.2.2")
+				return r
+			}(),
+		},
+		inNetInst: defName,
+		inPrefix:  "2001:aaaa::/64",
+		want: map[string]*NextHopSummary{
+			"1000:10:10::10": {
+				Address:         "1000:10:10::10",
+				Weight:          1,
+				NetworkInstance: "VRF-1",
+			},
+			"2.2.2.2": {
+				Address:         "2.2.2.2",
+				Weight:          1,
+				NetworkInstance: "VRF-1",
+			},
+		},
 	}}
 
 	for _, tt := range tests {

rib/rib_test.go

@@ -4184,6 +4184,29 @@ func TestResolvedEntryHook(t *testing.T) {
 					},
 				},
 			},
+		}, {
+			Entry: &spb.AFTOperation_NextHop{
+				NextHop: &aftpb.Afts_NextHopKey{
+					Index: 10,
+					NextHop: &aftpb.Afts_NextHop{
+						IpAddress: &wpb.StringValue{Value: "1000:10:10::10"},
+					},
+				},
+			},
+		}, {
+			Entry: &spb.AFTOperation_NextHopGroup{
+				NextHopGroup: &aftpb.Afts_NextHopGroupKey{
+					Id: 10,
+					NextHopGroup: &aftpb.Afts_NextHopGroup{
+						NextHop: []*aftpb.Afts_NextHopGroup_NextHopKey{{
+							Index: 10,
+							NextHop: &aftpb.Afts_NextHopGroup_NextHop{
+								Weight: &wpb.UintValue{Value: 32},
+							},
+						}},
+					},
+				},
+			},
 		}}
 
 		for i, op := range ops {
@@ -4259,7 +4282,7 @@ func TestResolvedEntryHook(t *testing.T) {
 			return nil
 		},
 	}, {
-		desc:  "aft helper",
+		desc:  "aft helper v4",
 		inRIB: baseRIB(),
 		inOperation: &spb.AFTOperation{
 			Id: 42,
@@ -4305,6 +4328,53 @@ func TestResolvedEntryHook(t *testing.T) {
 			}
 			return nil
 		},
+	}, {
+		desc:  "aft helper v6",
+		inRIB: baseRIB(),
+		inOperation: &spb.AFTOperation{
+			Id: 42,
+			Op: spb.AFTOperation_ADD,
+			Entry: &spb.AFTOperation_Ipv6{
+				Ipv6: &aftpb.Afts_Ipv6EntryKey{
+					Prefix: "2001:aaaa::/64",
+					Ipv6Entry: &aftpb.Afts_Ipv6Entry{
+						NextHopGroup: &wpb.UintValue{Value: 10},
+					},
+				},
+			},
+		},
+		inHook: func(ribs map[string]*aft.RIB, op constants.OpType, netinst string, aft constants.AFT, prefix any, _ ...ResolvedDetails) {
+			p, ok := prefix.(string)
+			if aft != constants.IPv6 || !ok {
+				gotCh <- fmt.Errorf("invalid prefix received, mismatched AFT (%s) or prefix data type (%v:%T)", aft, prefix, prefix)
+				return
+			}
+			summ, err := afthelper.NextHopAddrsForPrefix(ribs, netinst, p)
+			if err != nil {
+				gotCh <- err
+				return
+			}
+			gotCh <- summ
+		},
+		checkFn: func() error {
+			got := <-gotCh
+			switch t := got.(type) {
+			case error:
+				return fmt.Errorf("got error, %v", t)
+			case map[string]*afthelper.NextHopSummary:
+				want := map[string]*afthelper.NextHopSummary{
+					"1000:10:10::10": {
+						Weight:          32,
+						Address:         "1000:10:10::10",
+						NetworkInstance: "DEFAULT",
+					},
+				}
+				if diff := cmp.Diff(got, want); diff != "" {
+					return fmt.Errorf("got diff, %s", diff)
+				}
+			}
+			return nil
+		},
 	}, {
 		desc:  "add ipv6 entry",
 		inRIB: baseRIB(),