2 files changed, 30 insertions, 17 deletions

diff --git a/dhcpv6/dhcpv6relay.go b/dhcpv6/dhcpv6relay.go
index d9555cb..5c30bad 100644
--- a/dhcpv6/dhcpv6relay.go
+++ b/dhcpv6/dhcpv6relay.go
@@ -158,23 +158,26 @@ func (d *DHCPv6Relay) GetInnerMessage() (DHCPv6, error) {
 	}
 }
 
-// NewRelayReplFromRelayForw creates a RELAY_REPL packet based on a RELAY_FORW
-// packet and replaces the inner message with the passed DHCPv6 message.
+// NewRelayReplFromRelayForw creates a MessageTypeRelayReply based on a
+// MessageTypeRelayForward and replaces the inner message with the passed
+// DHCPv6 message. It copies the OptionInterfaceID and OptionRemoteID if the
+// options are present in the Relay packet.
 func NewRelayReplFromRelayForw(relayForw, msg DHCPv6) (DHCPv6, error) {
 	var (
 		err                error
 		linkAddr, peerAddr []net.IP
-		optiids            []Option
+		optiid             []Option
+		optrid             []Option
 	)
 	if relayForw == nil {
-		return nil, errors.New("RELAY_FORW cannot be nil")
+		return nil, errors.New("Relay message cannot be nil")
 	}
 	relay, ok := relayForw.(*DHCPv6Relay)
 	if !ok {
 		return nil, errors.New("Not a DHCPv6Relay")
 	}
 	if relay.Type() != MessageTypeRelayForward {
-		return nil, errors.New("The passed packet is not of type RELAY_FORW")
+		return nil, errors.New("The passed packet is not of type MessageTypeRelayForward")
 	}
 	if msg == nil {
 		return nil, errors.New("The passed message cannot be nil")
@@ -185,7 +188,8 @@ func NewRelayReplFromRelayForw(relayForw, msg DHCPv6) (DHCPv6, error) {
 	for {
 		linkAddr = append(linkAddr, relay.LinkAddr())
 		peerAddr = append(peerAddr, relay.PeerAddr())
-		optiids = append(optiids, relay.GetOneOption(OptionInterfaceID))
+		optiid = append(optiid, relay.GetOneOption(OptionInterfaceID))
+		optrid = append(optrid, relay.GetOneOption(OptionRemoteID))
 		decap, err := DecapsulateRelay(relay)
 		if err != nil {
 			return nil, err
@@ -198,12 +202,15 @@ func NewRelayReplFromRelayForw(relayForw, msg DHCPv6) (DHCPv6, error) {
 	}
 	for i := len(linkAddr) - 1; i >= 0; i-- {
 		msg, err = EncapsulateRelay(msg, MessageTypeRelayReply, linkAddr[i], peerAddr[i])
-		if opt := optiids[i]; opt != nil {
-			msg.AddOption(opt)
-		}
 		if err != nil {
 			return nil, err
 		}
+		if opt := optiid[i]; opt != nil {
+			msg.AddOption(opt)
+		}
+		if opt := optrid[i]; opt != nil {
+			msg.AddOption(opt)
+		}
 	}
 	return msg, nil
 }
diff --git a/dhcpv6/dhcpv6relay_test.go b/dhcpv6/dhcpv6relay_test.go
index afb4086..fe1b840 100644
--- a/dhcpv6/dhcpv6relay_test.go
+++ b/dhcpv6/dhcpv6relay_test.go
@@ -108,19 +108,23 @@ func TestDHCPv6RelayToBytes(t *testing.T) {
 }
 
 func TestNewRelayRepFromRelayForw(t *testing.T) {
+	// create a new relay forward
 	rf := DHCPv6Relay{}
 	rf.SetMessageType(MessageTypeRelayForward)
 	rf.SetPeerAddr(net.IPv6linklocalallrouters)
 	rf.SetLinkAddr(net.IPv6interfacelocalallnodes)
-	oro := OptRelayMsg{}
-	s := DHCPv6Message{}
-	s.SetMessage(MessageTypeSolicit)
-	cid := OptClientId{}
-	s.AddOption(&cid)
-	oro.SetRelayMessage(&s)
-	rf.AddOption(&oro)
+	rf.AddOption(&OptInterfaceId{})
+	rf.AddOption(&OptRemoteId{})
 
-	a, err := NewAdvertiseFromSolicit(&s)
+	// create the inner message
+	s, err := NewMessage()
+	require.NoError(t, err)
+	s.AddOption(&OptClientId{})
+	orm := OptRelayMsg{}
+	orm.SetRelayMessage(s)
+	rf.AddOption(&orm)
+
+	a, err := NewAdvertiseFromSolicit(s)
 	require.NoError(t, err)
 	rr, err := NewRelayReplFromRelayForw(&rf, a)
 	require.NoError(t, err)
@@ -129,6 +133,8 @@ func TestNewRelayRepFromRelayForw(t *testing.T) {
 	require.Equal(t, relay.HopCount(), rf.HopCount())
 	require.Equal(t, relay.PeerAddr(), rf.PeerAddr())
 	require.Equal(t, relay.LinkAddr(), rf.LinkAddr())
+	require.NotNil(t, rr.GetOneOption(OptionInterfaceID))
+	require.NotNil(t, rr.GetOneOption(OptionRemoteID))
 	m, err := relay.GetInnerMessage()
 	require.NoError(t, err)
 	require.Equal(t, m, a)