add DecapsulateRelayIndex

2 files changed, 40 insertions, 37 deletions

diff --git a/dhcpv6/dhcpv6.go b/dhcpv6/dhcpv6.go
index f091897..c9e42e2 100644
--- a/dhcpv6/dhcpv6.go
+++ b/dhcpv6/dhcpv6.go
@@ -1,7 +1,6 @@
 package dhcpv6
 
 import (
-	"errors"
 	"fmt"
 	"net"
 )
@@ -128,32 +127,50 @@ func delOption(options []Option, code OptionCode) []Option {
 	return newOpts
 }
 
-// DecapsulateRelay extracts the content of a relay message. It takes an
+// DecapsulateRelay extracts the content of a relay message. It does not recurse
+// if there are nested relay messages. Returns the original packet if is not not
+// a relay message
+func DecapsulateRelay(l DHCPv6) (DHCPv6, error) {
+	if !l.IsRelay() {
+		return l, nil
+	}
+	opt := l.GetOneOption(OPTION_RELAY_MSG)
+	if opt == nil {
+		return nil, fmt.Errorf("No OptRelayMsg found")
+	}
+	relayOpt := opt.(*OptRelayMsg)
+	if relayOpt.RelayMessage() == nil {
+		return nil, fmt.Errorf("Relay message cannot be nil")
+	}
+	return relayOpt.RelayMessage(), nil
+}
+
+// DecapsulateRelayIndex extracts the content of a relay message. It takes an
 // integer as index (e.g. if 0 return the outermost relay, 1 returns the
 // second, etc, and -1 returns the last). Returns the original packet if
 // it is not not a relay message.
-func DecapsulateRelay(l DHCPv6, index int) (DHCPv6, error) {
+func DecapsulateRelayIndex(l DHCPv6, index int) (DHCPv6, error) {
 	if !l.IsRelay() {
 		return l, nil
 	}
-	relay := l.(*DHCPv6Relay)
-	hops := int(relay.HopCount())
-	if index < 0 {
-		index = hops + index
-	}
-	if index < 0 || index > hops {
-		return nil, errors.New("index out of range")
+	if index == -1 {
+		for {
+			d, err := DecapsulateRelay(l)
+			if err != nil {
+				return nil, err
+			}
+			if !d.IsRelay() {
+				return l, nil
+			}
+			l = d
+		}
 	}
 	for i := 0; i <= index; i++ {
-		opt := l.GetOneOption(OPTION_RELAY_MSG)
-		if opt == nil {
-			return nil, fmt.Errorf("No OptRelayMsg found")
-		}
-		relayOpt := opt.(*OptRelayMsg)
-		l = relayOpt.RelayMessage()
-		if l == nil {
-			return nil, fmt.Errorf("Relay message cannot be nil")
+		d, err := DecapsulateRelay(l)
+		if err != nil {
+			return nil, err
 		}
+		l = d
 	}
 	return l, nil
 }
diff --git a/dhcpv6/dhcpv6_test.go b/dhcpv6/dhcpv6_test.go
index bc56ea3..6b7d2df 100644
--- a/dhcpv6/dhcpv6_test.go
+++ b/dhcpv6/dhcpv6_test.go
@@ -37,7 +37,7 @@ func TestNewMessage(t *testing.T) {
 	require.Empty(t, d.(*DHCPv6Message).options)
 }
 
-func TestDecapsulateRelay(t *testing.T) {
+func TestDecapsulateRelayIndex(t *testing.T) {
 	m := DHCPv6Message{}
 	r1, err := EncapsulateRelay(&m, RELAY_FORW, net.IPv6linklocalallnodes, net.IPv6interfacelocalallnodes)
 	require.NoError(t, err)
@@ -46,7 +46,7 @@ func TestDecapsulateRelay(t *testing.T) {
 	r3, err := EncapsulateRelay(r2, RELAY_FORW, net.IPv6unspecified, net.IPv6linklocalallrouters)
 	require.NoError(t, err)
 
-	first, err := DecapsulateRelay(r3, 0)
+	first, err := DecapsulateRelayIndex(r3, 0)
 	require.NoError(t, err)
 	relay, ok := first.(*DHCPv6Relay)
 	require.True(t, ok)
@@ -54,7 +54,7 @@ func TestDecapsulateRelay(t *testing.T) {
 	require.Equal(t, relay.LinkAddr(), net.IPv6loopback)
 	require.Equal(t, relay.PeerAddr(), net.IPv6linklocalallnodes)
 
-	second, err := DecapsulateRelay(r3, 1)
+	second, err := DecapsulateRelayIndex(r3, 1)
 	require.NoError(t, err)
 	relay, ok = second.(*DHCPv6Relay)
 	require.True(t, ok)
@@ -62,32 +62,18 @@ func TestDecapsulateRelay(t *testing.T) {
 	require.Equal(t, relay.LinkAddr(), net.IPv6linklocalallnodes)
 	require.Equal(t, relay.PeerAddr(), net.IPv6interfacelocalallnodes)
 
-	third, err := DecapsulateRelay(r3, 2)
+	third, err := DecapsulateRelayIndex(r3, 2)
 	require.NoError(t, err)
 	_, ok = third.(*DHCPv6Message)
 	require.True(t, ok)
 
-	_, err = DecapsulateRelay(r3, 3)
-	require.Error(t, err)
-
-	rfirst, err := DecapsulateRelay(r3, -1)
+	rfirst, err := DecapsulateRelayIndex(r3, -1)
 	require.NoError(t, err)
 	relay, ok = rfirst.(*DHCPv6Relay)
 	require.True(t, ok)
 	require.Equal(t, relay.HopCount(), uint8(0))
 	require.Equal(t, relay.LinkAddr(), net.IPv6linklocalallnodes)
 	require.Equal(t, relay.PeerAddr(), net.IPv6interfacelocalallnodes)
-
-	rsecond, err := DecapsulateRelay(r3, -2)
-	require.NoError(t, err)
-	relay, ok = rsecond.(*DHCPv6Relay)
-	require.True(t, ok)
-	require.Equal(t, relay.HopCount(), uint8(1))
-	require.Equal(t, relay.LinkAddr(), net.IPv6loopback)
-	require.Equal(t, relay.PeerAddr(), net.IPv6linklocalallnodes)
-
-	_, err = DecapsulateRelay(r3, -3)
-	require.Error(t, err)
 }
 
 func TestSettersAndGetters(t *testing.T) {