1 files changed, 201 insertions, 11 deletions
diff --git a/test/packetimpact/testbench/layers.go b/test/packetimpact/testbench/layers.go
index 2cbbbb318..817f5c261 100644
--- a/test/packetimpact/testbench/layers.go
+++ b/test/packetimpact/testbench/layers.go
@@ -36,14 +36,14 @@ import (
 type Layer interface {
 	fmt.Stringer
 
-	// toBytes converts the Layer into bytes. In places where the Layer's field
+	// ToBytes converts the Layer into bytes. In places where the Layer's field
 	// isn't nil, the value that is pointed to is used. When the field is nil, a
 	// reasonable default for the Layer is used. For example, "64" for IPv4 TTL
 	// and a calculated checksum for TCP or IP. Some layers require information
 	// from the previous or next layers in order to compute a default, such as
 	// TCP's checksum or Ethernet's type, so each Layer has a doubly-linked list
 	// to the layer's neighbors.
-	toBytes() ([]byte, error)
+	ToBytes() ([]byte, error)
 
 	// match checks if the current Layer matches the provided Layer. If either
 	// Layer has a nil in a given field, that field is considered matching.
@@ -174,7 +174,8 @@ func (l *Ether) String() string {
 	return stringLayer(l)
 }
 
-func (l *Ether) toBytes() ([]byte, error) {
+// ToBytes implements Layer.ToBytes.
+func (l *Ether) ToBytes() ([]byte, error) {
 	b := make([]byte, header.EthernetMinimumSize)
 	h := header.Ethernet(b)
 	fields := &header.EthernetFields{}
@@ -190,8 +191,9 @@ func (l *Ether) toBytes() ([]byte, error) {
 		switch n := l.next().(type) {
 		case *IPv4:
 			fields.Type = header.IPv4ProtocolNumber
+		case *IPv6:
+			fields.Type = header.IPv6ProtocolNumber
 		default:
-			// TODO(b/150301488): Support more protocols, like IPv6.
 			return nil, fmt.Errorf("ethernet header's next layer is unrecognized: %#v", n)
 		}
 	}
@@ -246,6 +248,8 @@ func parseEther(b []byte) (Layer, layerParser) {
 	switch h.Type() {
 	case header.IPv4ProtocolNumber:
 		nextParser = parseIPv4
+	case header.IPv6ProtocolNumber:
+		nextParser = parseIPv6
 	default:
 		// Assume that the rest is a payload.
 		nextParser = parsePayload
@@ -286,7 +290,8 @@ func (l *IPv4) String() string {
 	return stringLayer(l)
 }
 
-func (l *IPv4) toBytes() ([]byte, error) {
+// ToBytes implements Layer.ToBytes.
+func (l *IPv4) ToBytes() ([]byte, error) {
 	b := make([]byte, header.IPv4MinimumSize)
 	h := header.IPv4(b)
 	fields := &header.IPv4Fields{
@@ -421,6 +426,186 @@ func (l *IPv4) merge(other Layer) error {
 	return mergeLayer(l, other)
 }
 
+// IPv6 can construct and match an IPv6 encapsulation.
+type IPv6 struct {
+	LayerBase
+	TrafficClass  *uint8
+	FlowLabel     *uint32
+	PayloadLength *uint16
+	NextHeader    *uint8
+	HopLimit      *uint8
+	SrcAddr       *tcpip.Address
+	DstAddr       *tcpip.Address
+}
+
+func (l *IPv6) String() string {
+	return stringLayer(l)
+}
+
+// ToBytes implements Layer.ToBytes.
+func (l *IPv6) ToBytes() ([]byte, error) {
+	b := make([]byte, header.IPv6MinimumSize)
+	h := header.IPv6(b)
+	fields := &header.IPv6Fields{
+		HopLimit: 64,
+	}
+	if l.TrafficClass != nil {
+		fields.TrafficClass = *l.TrafficClass
+	}
+	if l.FlowLabel != nil {
+		fields.FlowLabel = *l.FlowLabel
+	}
+	if l.PayloadLength != nil {
+		fields.PayloadLength = *l.PayloadLength
+	} else {
+		for current := l.next(); current != nil; current = current.next() {
+			fields.PayloadLength += uint16(current.length())
+		}
+	}
+	if l.NextHeader != nil {
+		fields.NextHeader = *l.NextHeader
+	} else {
+		switch n := l.next().(type) {
+		case *TCP:
+			fields.NextHeader = uint8(header.TCPProtocolNumber)
+		case *UDP:
+			fields.NextHeader = uint8(header.UDPProtocolNumber)
+		case *ICMPv6:
+			fields.NextHeader = uint8(header.ICMPv6ProtocolNumber)
+		default:
+			// TODO(b/150301488): Support more protocols as needed.
+			return nil, fmt.Errorf("ToBytes can't deduce the IPv6 header's next protocol: %#v", n)
+		}
+	}
+	if l.HopLimit != nil {
+		fields.HopLimit = *l.HopLimit
+	}
+	if l.SrcAddr != nil {
+		fields.SrcAddr = *l.SrcAddr
+	}
+	if l.DstAddr != nil {
+		fields.DstAddr = *l.DstAddr
+	}
+	h.Encode(fields)
+	return h, nil
+}
+
+// parseIPv6 parses the bytes assuming that they start with an ipv6 header and
+// continues parsing further encapsulations.
+func parseIPv6(b []byte) (Layer, layerParser) {
+	h := header.IPv6(b)
+	tos, flowLabel := h.TOS()
+	ipv6 := IPv6{
+		TrafficClass:  &tos,
+		FlowLabel:     &flowLabel,
+		PayloadLength: Uint16(h.PayloadLength()),
+		NextHeader:    Uint8(h.NextHeader()),
+		HopLimit:      Uint8(h.HopLimit()),
+		SrcAddr:       Address(h.SourceAddress()),
+		DstAddr:       Address(h.DestinationAddress()),
+	}
+	var nextParser layerParser
+	switch h.TransportProtocol() {
+	case header.TCPProtocolNumber:
+		nextParser = parseTCP
+	case header.UDPProtocolNumber:
+		nextParser = parseUDP
+	case header.ICMPv6ProtocolNumber:
+		nextParser = parseICMPv6
+	default:
+		// Assume that the rest is a payload.
+		nextParser = parsePayload
+	}
+	return &ipv6, nextParser
+}
+
+func (l *IPv6) match(other Layer) bool {
+	return equalLayer(l, other)
+}
+
+func (l *IPv6) length() int {
+	return header.IPv6MinimumSize
+}
+
+// merge overrides the values in l with the values from other but only in fields
+// where the value is not nil.
+func (l *IPv6) merge(other Layer) error {
+	return mergeLayer(l, other)
+}
+
+// ICMPv6 can construct and match an ICMPv6 encapsulation.
+type ICMPv6 struct {
+	LayerBase
+	Type       *header.ICMPv6Type
+	Code       *byte
+	Checksum   *uint16
+	NDPPayload []byte
+}
+
+func (l *ICMPv6) String() string {
+	// TODO(eyalsoha): Do something smarter here when *l.Type is ParameterProblem?
+	// We could parse the contents of the Payload as if it were an IPv6 packet.
+	return stringLayer(l)
+}
+
+// ToBytes implements Layer.ToBytes.
+func (l *ICMPv6) ToBytes() ([]byte, error) {
+	b := make([]byte, header.ICMPv6HeaderSize+len(l.NDPPayload))
+	h := header.ICMPv6(b)
+	if l.Type != nil {
+		h.SetType(*l.Type)
+	}
+	if l.Code != nil {
+		h.SetCode(*l.Code)
+	}
+	copy(h.NDPPayload(), l.NDPPayload)
+	if l.Checksum != nil {
+		h.SetChecksum(*l.Checksum)
+	} else {
+		ipv6 := l.prev().(*IPv6)
+		h.SetChecksum(header.ICMPv6Checksum(h, *ipv6.SrcAddr, *ipv6.DstAddr, buffer.VectorisedView{}))
+	}
+	return h, nil
+}
+
+// ICMPv6Type is a helper routine that allocates a new ICMPv6Type value to store
+// v and returns a pointer to it.
+func ICMPv6Type(v header.ICMPv6Type) *header.ICMPv6Type {
+	return &v
+}
+
+// Byte is a helper routine that allocates a new byte value to store
+// v and returns a pointer to it.
+func Byte(v byte) *byte {
+	return &v
+}
+
+// parseICMPv6 parses the bytes assuming that they start with an ICMPv6 header.
+func parseICMPv6(b []byte) (Layer, layerParser) {
+	h := header.ICMPv6(b)
+	icmpv6 := ICMPv6{
+		Type:       ICMPv6Type(h.Type()),
+		Code:       Byte(h.Code()),
+		Checksum:   Uint16(h.Checksum()),
+		NDPPayload: h.NDPPayload(),
+	}
+	return &icmpv6, nil
+}
+
+func (l *ICMPv6) match(other Layer) bool {
+	return equalLayer(l, other)
+}
+
+func (l *ICMPv6) length() int {
+	return header.ICMPv6HeaderSize + len(l.NDPPayload)
+}
+
+// merge overrides the values in l with the values from other but only in fields
+// where the value is not nil.
+func (l *ICMPv6) merge(other Layer) error {
+	return mergeLayer(l, other)
+}
+
 // TCP can construct and match a TCP encapsulation.
 type TCP struct {
 	LayerBase
@@ -439,7 +624,8 @@ func (l *TCP) String() string {
 	return stringLayer(l)
 }
 
-func (l *TCP) toBytes() ([]byte, error) {
+// ToBytes implements Layer.ToBytes.
+func (l *TCP) ToBytes() ([]byte, error) {
 	b := make([]byte, header.TCPMinimumSize)
 	h := header.TCP(b)
 	if l.SrcPort != nil {
@@ -504,7 +690,7 @@ func layerChecksum(l Layer, protoNumber tcpip.TransportProtocolNumber) (uint16,
 	}
 	var payloadBytes buffer.VectorisedView
 	for current := l.next(); current != nil; current = current.next() {
-		payload, err := current.toBytes()
+		payload, err := current.ToBytes()
 		if err != nil {
 			return 0, fmt.Errorf("can't get bytes for next header: %s", payload)
 		}
@@ -578,7 +764,8 @@ func (l *UDP) String() string {
 	return stringLayer(l)
 }
 
-func (l *UDP) toBytes() ([]byte, error) {
+// ToBytes implements Layer.ToBytes.
+func (l *UDP) ToBytes() ([]byte, error) {
 	b := make([]byte, header.UDPMinimumSize)
 	h := header.UDP(b)
 	if l.SrcPort != nil {
@@ -661,7 +848,8 @@ func parsePayload(b []byte) (Layer, layerParser) {
 	return &payload, nil
 }
 
-func (l *Payload) toBytes() ([]byte, error) {
+// ToBytes implements Layer.ToBytes.
+func (l *Payload) ToBytes() ([]byte, error) {
 	return l.Bytes, nil
 }
 
@@ -697,11 +885,13 @@ func (ls *Layers) linkLayers() {
 	}
 }
 
-func (ls *Layers) toBytes() ([]byte, error) {
+// ToBytes converts the Layers into bytes. It creates a linked list of the Layer
+// structs and then concatentates the output of ToBytes on each Layer.
+func (ls *Layers) ToBytes() ([]byte, error) {
 	ls.linkLayers()
 	outBytes := []byte{}
 	for _, l := range *ls {
-		layerBytes, err := l.toBytes()
+		layerBytes, err := l.ToBytes()
 		if err != nil {
 			return nil, err
 		}