fixed tests, added marshalbinary to ethernet ans arp, fixed graceful shutdown, small improvements

Author: shadowy-pycoder

layers/arp.go

@@ -25,7 +25,7 @@ func (ao *ARPOperation) String() string {
 	return fmt.Sprintf("%s (%d)", ao.Desc, ao.Val)
 }
 
-// The Address Resolution Protocol (ARP) is a communication protocol
+// ARPPacket represents The Address Resolution Protocol (ARP) that is a communication protocol
 // used for discovering the link layer address, such as a MAC address,
 // associated with a given internet layer address, typically an IPv4 address.
 // Defined in RFC 826.
@@ -47,7 +47,7 @@ type ARPPacket struct {
 	TargetIP  netip.Addr // Internetwork address of the intended receiver.
 }
 
-func New(
+func NewARPPacket(
 	op Operation,
 	senderMAC net.HardwareAddr,
 	senderIP netip.Addr,
@@ -106,7 +106,7 @@ func (ap *ARPPacket) Summary() string {
 	case OperationRequest:
 		message = fmt.Sprintf("ARP Packet: (%s) Who has %s? Tell %s", ap.Op.Desc, ap.TargetIP, ap.SenderIP)
 	case OperationReply:
-		message = fmt.Sprintf("ARP Packet: (%s) %s at %s", ap.Op.Desc, ap.SenderIP, ap.SenderMAC)
+		message = fmt.Sprintf("ARP Packet: (%s) %s is at %s", ap.Op.Desc, ap.SenderIP, ap.SenderMAC)
 	default:
 		message = fmt.Sprintf("ARP Packet: (%s)", ap.Op.Desc)
 	}
@@ -115,23 +115,24 @@ func (ap *ARPPacket) Summary() string {
 
 // MarshalBinary implements encoding.BinaryMarshaler interface
 func (ap *ARPPacket) MarshalBinary() ([]byte, error) {
-	buf := make([]byte, 2+2+1+1+2+(ap.Plen*2)+(ap.Hlen*2))
-	binary.BigEndian.PutUint16(buf[0:2], ap.HardwareType)
-	binary.BigEndian.PutUint16(buf[2:4], ap.ProtocolType)
-	buf[4] = ap.Hlen
-	buf[5] = ap.Plen
-	binary.BigEndian.PutUint16(buf[6:8], uint16(ap.Op.Val))
+	b := make([]byte, 2+2+1+1+2+(ap.Plen*2)+(ap.Hlen*2))
+	binary.BigEndian.PutUint16(b[0:2], ap.HardwareType)
+	binary.BigEndian.PutUint16(b[2:4], ap.ProtocolType)
+	b[4] = ap.Hlen
+	b[5] = ap.Plen
+	binary.BigEndian.PutUint16(b[6:8], uint16(ap.Op.Val))
 	hoffset := 8 + ap.Hlen
 	poffset := hoffset + ap.Plen
-	copy(buf[8:hoffset], ap.SenderMAC)
-	copy(buf[hoffset:poffset], ap.SenderIP.AsSlice())
-	copy(buf[poffset:poffset+ap.Hlen], ap.TargetMAC)
-	copy(buf[poffset+ap.Hlen:poffset+ap.Hlen+ap.Plen], ap.TargetIP.AsSlice())
-	return buf, nil
+	copy(b[8:hoffset], ap.SenderMAC)
+	copy(b[hoffset:poffset], ap.SenderIP.AsSlice())
+	copy(b[poffset:poffset+ap.Hlen], ap.TargetMAC)
+	copy(b[poffset+ap.Hlen:poffset+ap.Hlen+ap.Plen], ap.TargetIP.AsSlice())
+	return b, nil
 }
 
-func (ap *ARPPacket) ToBytes() ([]byte, error) {
-	return ap.MarshalBinary()
+func (ap *ARPPacket) ToBytes() []byte {
+	b, _ := ap.MarshalBinary()
+	return b
 }
 
 // UnmarshalBinary implements encoding.BinaryUnmarshaler interface

layers/arp_test.go

@@ -60,8 +60,7 @@ func TestParseARP(t *testing.T) {
 		ProtocolTypeDesc: "IPv4",
 		Hlen:             0x06,
 		Plen:             0x04,
-		Op:               0x0001,
-		OpDesc:           "request",
+		Op:               &ARPOperation{Val: 1, Desc: "request"},
 		SenderMAC:        net.HardwareAddr{0x7b, 0x13, 0x0b, 0x87, 0xea, 0x51},
 		SenderIP:         netip.AddrFrom4([4]byte{0x7F, 0x00, 0x00, 0x01}),
 		TargetMAC:        net.HardwareAddr{0x43, 0x40, 0x8d, 0x28, 0xca, 0x0b},

layers/ethernet.go

@@ -9,60 +9,111 @@ import (
 
 const headerSizeEthernet = 14
 
+type EtherType uint16
+
+const (
+	EtherTypeIPv4 EtherType = 0x0800
+	EtherTypeARP  EtherType = 0x0806
+	EtherTypeIPv6 EtherType = 0x86dd
+)
+
+type EthernetType struct {
+	Val  EtherType
+	Desc string // Protocol description
+}
+
+func (et *EthernetType) String() string {
+	return fmt.Sprintf("%s (%#04x)", et.Desc, et.Val)
+}
+
+// EthernetFrame represents Ethernet Frame
 // An Ethernet frame is a data link layer protocol data unit.
 type EthernetFrame struct {
-	DstMAC        net.HardwareAddr // MAC address of the destination device.
-	SrcMAC        net.HardwareAddr // MAC address of the source device.
-	EtherType     uint16           // The protocol of the upper layer.
-	EtherTypeDesc string           // Protocol description
-	payload       []byte
+	DstMAC    net.HardwareAddr // MAC address of the destination device.
+	SrcMAC    net.HardwareAddr // MAC address of the source device.
+	EtherType *EthernetType    // The protocol of the upper layer.
+	Payload   []byte
+}
+
+func NewEthernetFrame(dstMAC, srcMAC net.HardwareAddr, et EtherType, payload []byte) (*EthernetFrame, error) {
+	if len(dstMAC) != 6 || len(srcMAC) != 6 {
+		return nil, fmt.Errorf("malformed hardware address")
+	}
+	return &EthernetFrame{
+		DstMAC:    dstMAC,
+		SrcMAC:    srcMAC,
+		EtherType: &EthernetType{Val: et, Desc: ethertypedesc(et)},
+		Payload:   payload,
+	}, nil
 }
 
 func (ef *EthernetFrame) String() string {
 	return fmt.Sprintf(`%s
 - DstMAC: %s
 - SrcMAC: %s
-- EtherType: %s (%#04x)
-- Payload: %d bytes 
+- EtherType: %s
+- Payload: %d bytes
 %s`,
 		ef.Summary(),
 		ef.DstMAC,
 		ef.SrcMAC,
-		ef.EtherTypeDesc,
 		ef.EtherType,
-		len(ef.payload),
-		hex.Dump(ef.payload))
+		len(ef.Payload),
+		hex.Dump(ef.ToBytes()))
 }
 
 func (ef *EthernetFrame) Summary() string {
 	return fmt.Sprintf("Ethernet Frame: Src MAC: %s -> Dst MAC: %s", ef.SrcMAC, ef.DstMAC)
 }
 
-// Parse parses the given byte data into an Ethernet frame.
-func (ef *EthernetFrame) Parse(data []byte) error {
+func (ef *EthernetFrame) MarshalBinary() ([]byte, error) {
+	b := make([]byte, 6+6+2+len(ef.Payload))
+	copy(b[0:6], ef.DstMAC)
+	copy(b[6:12], ef.SrcMAC)
+	binary.BigEndian.PutUint16(b[12:14], uint16(ef.EtherType.Val))
+	copy(b[14:], ef.Payload)
+	return b, nil
+}
+
+func (ef *EthernetFrame) ToBytes() []byte {
+	b, _ := ef.MarshalBinary()
+	return b
+}
+
+func (ef *EthernetFrame) UnmarshalBinary(data []byte) error {
 	if len(data) < headerSizeEthernet {
 		return fmt.Errorf("did not read a complete Ethernet frame, only %d bytes read", len(data))
 	}
 	ef.DstMAC = net.HardwareAddr(data[0:6])
 	ef.SrcMAC = net.HardwareAddr(data[6:12])
-	ef.EtherType = binary.BigEndian.Uint16(data[12:14])
-	ef.payload = data[headerSizeEthernet:]
-	ef.EtherTypeDesc, _ = ef.NextLayer()
+	et := EtherType(binary.BigEndian.Uint16(data[12:14]))
+	etdesc := ethertypedesc(et)
+	ef.EtherType = &EthernetType{Val: et, Desc: etdesc}
+	ef.Payload = data[headerSizeEthernet:]
 	return nil
 }
 
+// Parse parses the given byte data into an Ethernet frame.
+func (ef *EthernetFrame) Parse(data []byte) error {
+	return ef.UnmarshalBinary(data)
+}
+
 // NextLayer returns the name and payload of the next layer protocol based on the EtherType field of the EthernetFrame.
 func (ef *EthernetFrame) NextLayer() (string, []byte) {
-	var layer string
-	switch ef.EtherType {
-	case 0x0800:
-		layer = "IPv4"
-	case 0x0806:
-		layer = "ARP"
-	case 0x86dd:
-		layer = "IPv6"
+	return ef.EtherType.Desc, ef.Payload
+}
+
+func ethertypedesc(et EtherType) string {
+	var etdesc string
+	switch et {
+	case EtherTypeIPv4:
+		etdesc = "IPv4"
+	case EtherTypeARP:
+		etdesc = "ARP"
+	case EtherTypeIPv6:
+		etdesc = "IPv6"
 	default:
-		layer = ""
+		etdesc = ""
 	}
-	return layer, ef.payload
+	return etdesc
 }

layers/ethernet_test.go

@@ -22,11 +22,10 @@ func BenchmarkParseEthernet(b *testing.B) {
 
 func TestParseEthernet(t *testing.T) {
 	expected := &EthernetFrame{
-		DstMAC:        net.HardwareAddr{0x7b, 0x13, 0x0b, 0x87, 0xea, 0x51},
-		SrcMAC:        net.HardwareAddr{0x43, 0x40, 0x8d, 0x28, 0xca, 0x0b},
-		EtherType:     0x0800,
-		EtherTypeDesc: "IPv4",
-		payload:       []byte{},
+		DstMAC:    net.HardwareAddr{0x7b, 0x13, 0x0b, 0x87, 0xea, 0x51},
+		SrcMAC:    net.HardwareAddr{0x43, 0x40, 0x8d, 0x28, 0xca, 0x0b},
+		EtherType: &EthernetType{Val: 0x0800, Desc: "IPv4"},
+		Payload:   []byte{},
 	}
 	eth := &EthernetFrame{}
 	packet, close := testPacket(t, "ethernet")

layers/http.go

@@ -42,17 +42,23 @@ func (h *HTTPMessage) String() string {
 func (h *HTTPMessage) Summary() string {
 	m := fmt.Sprintf("HTTP Message: %s", contdata)
 	if h.Request != nil {
-		m = fmt.Sprintf("HTTP Request: %s %s%s%s %s Content-Length: %d",
-			h.Request.Method, h.Request.Host, h.Request.URL.Path, h.Request.URL.RawQuery, h.Request.Proto, h.Request.ContentLength)
+		m = fmt.Sprintf(
+			"HTTP Request: %s %s%s%s %s Content-Length: %d",
+			h.Request.Method,
+			h.Request.Host,
+			h.Request.URL.Path,
+			h.Request.URL.RawQuery,
+			h.Request.Proto,
+			h.Request.ContentLength,
+		)
 	} else if h.Response != nil {
 		m = fmt.Sprintf("HTTP Response: %s %s Content-Length: %d",
 			h.Response.Proto, h.Response.Status, h.Response.ContentLength)
 	}
-	return fmt.Sprintf("%s", m)
+	return m
 }
 
 func (h *HTTPMessage) Parse(data []byte) error {
-
 	if !bytes.Contains(data, protohttp10) && !bytes.Contains(data, protohttp11) {
 		h.Request = nil
 		h.Response = nil
@@ -119,7 +125,8 @@ func (h *HTTPMessage) MarshalJSON() ([]byte, error) {
 			Proto:         h.Response.Proto,
 			Status:        h.Response.Status,
 			ContentLength: int(h.Response.ContentLength),
-			Header:        h.Response.Header}})
+			Header:        h.Response.Header,
+		}})
 	}
 	return nil, fmt.Errorf("both request and response are empty")
 }

layers/http_test.go

@@ -1,3 +1,4 @@
+// Package layers
 package layers
 
 import (