dhcpv6: standardize GetInnerMessage

3 files changed, 12 insertions, 5 deletions

diff --git a/dhcpv6/dhcpv6.go b/dhcpv6/dhcpv6.go
index 3c5e193..c8c4292 100644
--- a/dhcpv6/dhcpv6.go
+++ b/dhcpv6/dhcpv6.go
@@ -16,6 +16,8 @@ type DHCPv6 interface {
 	String() string
 	Summary() string
 	IsRelay() bool
+	GetInnerMessage() (*Message, error)
+
 	GetOption(code OptionCode) []Option
 	GetOneOption(code OptionCode) Option
 	AddOption(Option)
diff --git a/dhcpv6/dhcpv6message.go b/dhcpv6/dhcpv6message.go
index c367d36..afd252d 100644
--- a/dhcpv6/dhcpv6message.go
+++ b/dhcpv6/dhcpv6message.go
@@ -206,6 +206,11 @@ func (m Message) Type() MessageType {
 	return m.MessageType
 }
 
+// GetInnerMessage returns the message itself.
+func (m *Message) GetInnerMessage() (*Message, error) {
+	return m, nil
+}
+
 // AddOption adds an option to this message.
 func (m *Message) AddOption(option Option) {
 	m.Options.Add(option)
diff --git a/dhcpv6/dhcpv6relay.go b/dhcpv6/dhcpv6relay.go
index 497ef64..1506dc4 100644
--- a/dhcpv6/dhcpv6relay.go
+++ b/dhcpv6/dhcpv6relay.go
@@ -91,21 +91,21 @@ func (r *RelayMessage) IsRelay() bool {
 }
 
 // GetInnerMessage recurses into a relay message and extract and return the
-// inner Message.  Return nil if none found (e.g. not a relay message).
-func (r *RelayMessage) GetInnerMessage() (DHCPv6, error) {
+// inner Message. Return nil if none found (e.g. not a relay message).
+func (r *RelayMessage) GetInnerMessage() (*Message, error) {
 	var (
 		p   DHCPv6
 		err error
 	)
 	p = r
 	for {
-		if !p.IsRelay() {
-			return p, nil
-		}
 		p, err = DecapsulateRelay(p)
 		if err != nil {
 			return nil, err
 		}
+		if m, ok := p.(*Message); ok {
+			return m, nil
+		}
 	}
 }