Merge pull request #83 from pmazzini/master

add DecapsulateRelayIndex

3 files changed, 74 insertions, 24 deletions

diff --git a/dhcpv6/dhcpv6.go b/dhcpv6/dhcpv6.go
index d8427ea..b424101 100644
--- a/dhcpv6/dhcpv6.go
+++ b/dhcpv6/dhcpv6.go
@@ -145,6 +145,38 @@ func DecapsulateRelay(l DHCPv6) (DHCPv6, error) {
 	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 DecapsulateRelayIndex(l DHCPv6, index int) (DHCPv6, error) {
+	if !l.IsRelay() {
+		return l, nil
+	}
+	if index < -1 {
+		return nil, fmt.Errorf("Invalid index: %d", index)
+	} else 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++ {
+		d, err := DecapsulateRelay(l)
+		if err != nil {
+			return nil, err
+		}
+		l = d
+	}
+	return l, nil
+}
+
 // EncapsulateRelay creates a DHCPv6Relay message containing the passed DHCPv6
 // message as payload. The passed message type must be  either RELAY_FORW or
 // RELAY_REPL
diff --git a/dhcpv6/dhcpv6_test.go b/dhcpv6/dhcpv6_test.go
index 3db4cee..3e2b6bc 100644
--- a/dhcpv6/dhcpv6_test.go
+++ b/dhcpv6/dhcpv6_test.go
@@ -38,6 +38,48 @@ func TestNewMessage(t *testing.T) {
 	require.Empty(t, d.(*DHCPv6Message).options)
 }
 
+func TestDecapsulateRelayIndex(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 := DecapsulateRelayIndex(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 := DecapsulateRelayIndex(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 := DecapsulateRelayIndex(r3, 2)
+	require.NoError(t, err)
+	_, ok = third.(*DHCPv6Message)
+	require.True(t, ok)
+
+	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)
+
+	_, err = DecapsulateRelayIndex(r3, -2)
+	require.Error(t, err)
+}
+
 func TestSettersAndGetters(t *testing.T) {
 	d := DHCPv6Message{}
 	// Message
diff --git a/dhcpv6/dhcpv6relay.go b/dhcpv6/dhcpv6relay.go
index dbdb623..1bde354 100644
--- a/dhcpv6/dhcpv6relay.go
+++ b/dhcpv6/dhcpv6relay.go
@@ -158,30 +158,6 @@ func (d *DHCPv6Relay) GetInnerMessage() (DHCPv6, error) {
 	}
 }
 
-// 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) {
-	var (
-		p   DHCPv6
-		err error
-	)
-	p = r
-	hops := r.HopCount()
-	addr := r.PeerAddr()
-	for i := uint8(0); i < hops; i++ {
-		p, err = DecapsulateRelay(p)
-		if err != nil {
-			return nil, err
-		}
-		if p.IsRelay() {
-			addr = p.(*DHCPv6Relay).PeerAddr()
-		} else {
-			return nil, fmt.Errorf("Wrong Hop count")
-		}
-	}
-	return addr, 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) {