extend DecapsulateRelay

4 files changed, 80 insertions, 64 deletions

diff --git a/dhcpv6/dhcpv6.go b/dhcpv6/dhcpv6.go
index d8427ea..f091897 100644
--- a/dhcpv6/dhcpv6.go
+++ b/dhcpv6/dhcpv6.go
@@ -1,6 +1,7 @@
 package dhcpv6
 
 import (
+	"errors"
 	"fmt"
 	"net"
 )
@@ -127,22 +128,34 @@ func delOption(options []Option, code OptionCode) []Option {
 	return newOpts
 }
 
-// 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) {
+// DecapsulateRelay 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) {
 	if !l.IsRelay() {
 		return l, nil
 	}
-	opt := l.GetOneOption(OPTION_RELAY_MSG)
-	if opt == nil {
-		return nil, fmt.Errorf("No OptRelayMsg found")
+	relay := l.(*DHCPv6Relay)
+	hops := int(relay.HopCount())
+	if index < 0 {
+		index = hops + index
 	}
-	relayOpt := opt.(*OptRelayMsg)
-	if relayOpt.RelayMessage() == nil {
-		return nil, fmt.Errorf("Relay message cannot be nil")
+	if index < 0 || index > hops {
+		return nil, errors.New("index out of range")
 	}
-	return relayOpt.RelayMessage(), nil
+	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")
+		}
+	}
+	return l, nil
 }
 
 // EncapsulateRelay creates a DHCPv6Relay message containing the passed DHCPv6
diff --git a/dhcpv6/dhcpv6_test.go b/dhcpv6/dhcpv6_test.go
index f8ec76d..bc56ea3 100644
--- a/dhcpv6/dhcpv6_test.go
+++ b/dhcpv6/dhcpv6_test.go
@@ -1,6 +1,7 @@
 package dhcpv6
 
 import (
+	"net"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -36,6 +37,59 @@ func TestNewMessage(t *testing.T) {
 	require.Empty(t, d.(*DHCPv6Message).options)
 }
 
+func TestDecapsulateRelay(t *testing.T) {
+	m := DHCPv6Message{}
+	r1, err := EncapsulateRelay(&m, RELAY_FORW, net.IPv6linklocalallnodes, net.IPv6interfacelocalallnodes)
+	require.NoError(t, err)
+	r2, err := EncapsulateRelay(r1, RELAY_FORW, net.IPv6loopback, net.IPv6linklocalallnodes)
+	require.NoError(t, err)
+	r3, err := EncapsulateRelay(r2, RELAY_FORW, net.IPv6unspecified, net.IPv6linklocalallrouters)
+	require.NoError(t, err)
+
+	first, err := DecapsulateRelay(r3, 0)
+	require.NoError(t, err)
+	relay, ok := first.(*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)
+
+	second, err := DecapsulateRelay(r3, 1)
+	require.NoError(t, err)
+	relay, ok = second.(*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)
+
+	third, err := DecapsulateRelay(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)
+	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) {
 	d := DHCPv6Message{}
 	// Message
diff --git a/dhcpv6/dhcpv6relay.go b/dhcpv6/dhcpv6relay.go
index 2387212..eaebfc2 100644
--- a/dhcpv6/dhcpv6relay.go
+++ b/dhcpv6/dhcpv6relay.go
@@ -151,39 +151,13 @@ func (d *DHCPv6Relay) GetInnerMessage() (DHCPv6, error) {
 		if !p.IsRelay() {
 			return p, nil
 		}
-		p, err = DecapsulateRelay(p)
+		p, err = DecapsulateRelay(p, 0)
 		if err != nil {
 			return nil, err
 		}
 	}
 }
 
-// GetInnerRelay recurses into a relay message and extract and return the inner
-// DHCPv6Relay. Return nil if none found (e.g. not a relay message).
-func (r *DHCPv6Relay) GetInnerRelay() (DHCPv6, error) {
-	p := r
-	for {
-		d, err := DecapsulateRelay(p)
-		if err != nil {
-			return nil, err
-		}
-		if !d.IsRelay() {
-			return p, nil
-		}
-		p = d.(*DHCPv6Relay)
-	}
-}
-
-// GetInnerPeerAddr returns the peer address in the inner most relay info
-// header, this is typically the IP address of the client making the request.
-func (r *DHCPv6Relay) GetInnerPeerAddr() (net.IP, error) {
-	p, err := r.GetInnerRelay()
-	if err != nil {
-		return nil, err
-	}
-	return p.(*DHCPv6Relay).PeerAddr(), nil
-}
-
 // NewRelayReplFromRelayForw creates a RELAY_REPL packet based on a RELAY_FORW
 // packet and replaces the inner message with the passed DHCPv6 message.
 func NewRelayReplFromRelayForw(relayForw, msg DHCPv6) (DHCPv6, error) {
@@ -212,7 +186,7 @@ func NewRelayReplFromRelayForw(relayForw, msg DHCPv6) (DHCPv6, error) {
 		linkAddr = append(linkAddr, relay.LinkAddr())
 		peerAddr = append(peerAddr, relay.PeerAddr())
 		optiids = append(optiids, relay.GetOneOption(OPTION_INTERFACE_ID))
-		decap, err := DecapsulateRelay(relay)
+		decap, err := DecapsulateRelay(relay, 0)
 		if err != nil {
 			return nil, err
 		}
diff --git a/dhcpv6/dhcpv6relay_test.go b/dhcpv6/dhcpv6relay_test.go
index 3effb35..e08e6a4 100644
--- a/dhcpv6/dhcpv6relay_test.go
+++ b/dhcpv6/dhcpv6relay_test.go
@@ -107,31 +107,6 @@ func TestDHCPv6RelayToBytes(t *testing.T) {
 	}
 }
 
-func TestGetInnerRelay(t *testing.T) {
-	m := DHCPv6Message{}
-	r1, err := EncapsulateRelay(&m, RELAY_FORW, net.IPv6linklocalallnodes, net.IPv6interfacelocalallnodes)
-	require.NoError(t, err)
-	r2, err := EncapsulateRelay(r1, RELAY_FORW, net.IPv6loopback, net.IPv6linklocalallnodes)
-	require.NoError(t, err)
-	r3, err := EncapsulateRelay(r2, RELAY_FORW, net.IPv6unspecified, net.IPv6linklocalallrouters)
-	require.NoError(t, err)
-
-	relay3, ok := r3.(*DHCPv6Relay)
-	require.True(t, ok)
-
-	ir, err := relay3.GetInnerRelay()
-	require.NoError(t, err)
-	relay, ok := ir.(*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)
-
-	innerPeerAddr, err := relay3.GetInnerPeerAddr()
-	require.NoError(t, err)
-	require.Equal(t, innerPeerAddr, net.IPv6interfacelocalallnodes)
-}
-
 func TestNewRelayRepFromRelayForw(t *testing.T) {
 	rf := DHCPv6Relay{}
 	rf.SetMessageType(RELAY_FORW)