TUN-6666: Define packet package

This package defines IP and ICMP packet, decoders, encoder and flow

packet/decoder.go

@@ -0,0 +1,184 @@
+package packet
+
+import (
+	"fmt"
+
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/layers"
+	"github.com/pkg/errors"
+	"golang.org/x/net/icmp"
+)
+
+func FindProtocol(p []byte) (layers.IPProtocol, error) {
+	version, err := FindIPVersion(p)
+	if err != nil {
+		return 0, err
+	}
+	switch version {
+	case 4:
+		if len(p) < ipv4HeaderLen {
+			return 0, fmt.Errorf("IPv4 packet should have at least %d bytes, got %d bytes", ipv4HeaderLen, len(p))
+		}
+		// Protocol is in the 10th byte of IPv4 header
+		return layers.IPProtocol(p[9]), nil
+	case 6:
+		if len(p) < ipv6HeaderLen {
+			return 0, fmt.Errorf("IPv6 packet should have at least %d bytes, got %d bytes", ipv6HeaderLen, len(p))
+		}
+		// Next header is in the 7th byte of IPv6 header
+		return layers.IPProtocol(p[6]), nil
+	default:
+		return 0, fmt.Errorf("unknow ip version %d", version)
+	}
+}
+
+func FindIPVersion(p []byte) (uint8, error) {
+	if len(p) == 0 {
+		return 0, fmt.Errorf("packet length is 0")
+	}
+	return p[0] >> 4, nil
+}
+
+// IPDecoder decodes raw packets into IP. It can process packets sequentially without allocating
+// memory for the layers, so it cannot be called concurrently.
+type IPDecoder struct {
+	ipv4   *layers.IPv4
+	ipv6   *layers.IPv6
+	layers uint8
+
+	v4parser *gopacket.DecodingLayerParser
+	v6parser *gopacket.DecodingLayerParser
+}
+
+func NewIPDecoder() *IPDecoder {
+	var (
+		ipv4 layers.IPv4
+		ipv6 layers.IPv6
+	)
+	dlpv4 := gopacket.NewDecodingLayerParser(layers.LayerTypeIPv4)
+	dlpv4.SetDecodingLayerContainer(gopacket.DecodingLayerSparse(nil))
+	dlpv4.AddDecodingLayer(&ipv4)
+	// Stop parsing when it encounter a layer that it doesn't have a parser
+	dlpv4.IgnoreUnsupported = true
+
+	dlpv6 := gopacket.NewDecodingLayerParser(layers.LayerTypeIPv6)
+	dlpv6.SetDecodingLayerContainer(gopacket.DecodingLayerSparse(nil))
+	dlpv6.AddDecodingLayer(&ipv6)
+	dlpv6.IgnoreUnsupported = true
+
+	return &IPDecoder{
+		ipv4:     &ipv4,
+		ipv6:     &ipv6,
+		layers:   1,
+		v4parser: dlpv4,
+		v6parser: dlpv6,
+	}
+}
+
+func (pd *IPDecoder) Decode(packet []byte) (*IP, error) {
+	// Should decode to IP layer
+	decoded, err := pd.decodeByVersion(packet)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, layerType := range decoded {
+		switch layerType {
+		case layers.LayerTypeIPv4:
+			return newIPv4(pd.ipv4)
+		case layers.LayerTypeIPv6:
+			return newIPv6(pd.ipv6)
+		}
+	}
+	return nil, fmt.Errorf("no ip layer is decoded")
+}
+
+func (pd *IPDecoder) decodeByVersion(packet []byte) ([]gopacket.LayerType, error) {
+	version, err := FindIPVersion(packet)
+	if err != nil {
+		return nil, err
+	}
+	decoded := make([]gopacket.LayerType, 0, pd.layers)
+	switch version {
+	case 4:
+		err = pd.v4parser.DecodeLayers(packet, &decoded)
+	case 6:
+		err = pd.v6parser.DecodeLayers(packet, &decoded)
+	default:
+		err = fmt.Errorf("unknow ip version %d", version)
+	}
+	if err != nil {
+		return nil, err
+	}
+	return decoded, nil
+}
+
+// ICMPDecoder decodes raw packets into IP and ICMP. It can process packets sequentially without allocating
+// memory for the layers, so it cannot be called concurrently.
+type ICMPDecoder struct {
+	*IPDecoder
+	icmpv4 *layers.ICMPv4
+	icmpv6 *layers.ICMPv6
+}
+
+func NewICMPDecoder() *ICMPDecoder {
+	ipDecoder := NewIPDecoder()
+
+	var (
+		icmpv4 layers.ICMPv4
+		icmpv6 layers.ICMPv6
+	)
+	ipDecoder.layers++
+	ipDecoder.v4parser.AddDecodingLayer(&icmpv4)
+	ipDecoder.v6parser.AddDecodingLayer(&icmpv6)
+
+	return &ICMPDecoder{
+		IPDecoder: ipDecoder,
+		icmpv4:    &icmpv4,
+		icmpv6:    &icmpv6,
+	}
+}
+
+func (pd *ICMPDecoder) Decode(packet []byte) (*ICMP, error) {
+	// Should decode to IP and optionally ICMP layer
+	decoded, err := pd.decodeByVersion(packet)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, layerType := range decoded {
+		switch layerType {
+		case layers.LayerTypeICMPv4:
+			ipv4, err := newIPv4(pd.ipv4)
+			if err != nil {
+				return nil, err
+			}
+			msg, err := icmp.ParseMessage(int(layers.IPProtocolICMPv4), append(pd.icmpv4.Contents, pd.icmpv4.Payload...))
+			if err != nil {
+				return nil, errors.Wrap(err, "failed to parse ICMPv4 message")
+			}
+			return &ICMP{
+				IP:      ipv4,
+				Message: msg,
+			}, nil
+		case layers.LayerTypeICMPv6:
+			ipv6, err := newIPv6(pd.ipv6)
+			if err != nil {
+				return nil, err
+			}
+			msg, err := icmp.ParseMessage(int(layers.IPProtocolICMPv6), append(pd.icmpv6.Contents, pd.icmpv6.Payload...))
+			if err != nil {
+				return nil, errors.Wrap(err, "failed to parse ICMPv6")
+			}
+			return &ICMP{
+				IP:      ipv6,
+				Message: msg,
+			}, nil
+		}
+	}
+	layers := make([]string, len(decoded))
+	for i, l := range decoded {
+		layers[i] = l.String()
+	}
+	return nil, fmt.Errorf("Expect to decode IP and ICMP layers, got %s", layers)
+}

packet/decoder_test.go

@@ -0,0 +1,252 @@
+package packet
+
+import (
+	"net"
+	"net/netip"
+	"testing"
+
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/layers"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/net/icmp"
+	"golang.org/x/net/ipv4"
+	"golang.org/x/net/ipv6"
+)
+
+func TestDecodeIP(t *testing.T) {
+	ipDecoder := NewIPDecoder()
+	icmpDecoder := NewICMPDecoder()
+	udps := []UDP{
+		{
+			IP: IP{
+				Src:      netip.MustParseAddr("172.16.0.1"),
+				Dst:      netip.MustParseAddr("10.0.0.1"),
+				Protocol: layers.IPProtocolUDP,
+			},
+			SrcPort: 31678,
+			DstPort: 53,
+		},
+		{
+
+			IP: IP{
+				Src:      netip.MustParseAddr("fd51:2391:523:f4ee::1"),
+				Dst:      netip.MustParseAddr("fd51:2391:697:f4ee::2"),
+				Protocol: layers.IPProtocolUDP,
+			},
+			SrcPort: 52139,
+			DstPort: 1053,
+		},
+	}
+
+	encoder := NewEncoder()
+	for _, udp := range udps {
+		p, err := encoder.Encode(&udp)
+		require.NoError(t, err)
+
+		ipPacket, err := ipDecoder.Decode(p.Data)
+		require.NoError(t, err)
+		assertIPLayer(t, &udp.IP, ipPacket)
+
+		icmpPacket, err := icmpDecoder.Decode(p.Data)
+		require.Error(t, err)
+		require.Nil(t, icmpPacket)
+	}
+}
+
+func TestDecodeICMP(t *testing.T) {
+	ipDecoder := NewIPDecoder()
+	icmpDecoder := NewICMPDecoder()
+	var (
+		ipv4Packet = IP{
+			Src:      netip.MustParseAddr("172.16.0.1"),
+			Dst:      netip.MustParseAddr("10.0.0.1"),
+			Protocol: layers.IPProtocolICMPv4,
+		}
+		ipv6Packet = IP{
+			Src:      netip.MustParseAddr("fd51:2391:523:f4ee::1"),
+			Dst:      netip.MustParseAddr("fd51:2391:697:f4ee::2"),
+			Protocol: layers.IPProtocolICMPv6,
+		}
+		icmpID  = 100
+		icmpSeq = 52819
+	)
+	tests := []struct {
+		testCase string
+		packet   *ICMP
+	}{
+		{
+			testCase: "icmpv4 time exceed",
+			packet: &ICMP{
+				IP: &ipv4Packet,
+				Message: &icmp.Message{
+					Type: ipv4.ICMPTypeTimeExceeded,
+					Code: 0,
+					Body: &icmp.TimeExceeded{
+						Data: []byte("original packet"),
+					},
+				},
+			},
+		},
+		{
+			testCase: "icmpv4 echo",
+			packet: &ICMP{
+				IP: &ipv4Packet,
+				Message: &icmp.Message{
+					Type: ipv4.ICMPTypeEcho,
+					Code: 0,
+					Body: &icmp.Echo{
+						ID:   icmpID,
+						Seq:  icmpSeq,
+						Data: []byte("icmpv4 echo"),
+					},
+				},
+			},
+		},
+		{
+			testCase: "icmpv6 destination unreachable",
+			packet: &ICMP{
+				IP: &ipv6Packet,
+				Message: &icmp.Message{
+					Type: ipv6.ICMPTypeDestinationUnreachable,
+					Code: 4,
+					Body: &icmp.DstUnreach{
+						Data: []byte("original packet"),
+					},
+				},
+			},
+		},
+		{
+			testCase: "icmpv6 echo",
+			packet: &ICMP{
+				IP: &ipv6Packet,
+				Message: &icmp.Message{
+					Type: ipv6.ICMPTypeEchoRequest,
+					Code: 0,
+					Body: &icmp.Echo{
+						ID:   icmpID,
+						Seq:  icmpSeq,
+						Data: []byte("icmpv6 echo"),
+					},
+				},
+			},
+		},
+	}
+
+	encoder := NewEncoder()
+	for _, test := range tests {
+		p, err := encoder.Encode(test.packet)
+		require.NoError(t, err)
+
+		ipPacket, err := ipDecoder.Decode(p.Data)
+		require.NoError(t, err)
+		if ipPacket.Src.Is4() {
+			assertIPLayer(t, &ipv4Packet, ipPacket)
+		} else {
+			assertIPLayer(t, &ipv6Packet, ipPacket)
+		}
+		icmpPacket, err := icmpDecoder.Decode(p.Data)
+		require.NoError(t, err)
+		require.Equal(t, ipPacket, icmpPacket.IP)
+
+		require.Equal(t, test.packet.Type, icmpPacket.Type)
+		require.Equal(t, test.packet.Code, icmpPacket.Code)
+		require.Equal(t, test.packet.Body, icmpPacket.Body)
+		expectedBody, err := test.packet.Body.Marshal(test.packet.Type.Protocol())
+		require.NoError(t, err)
+		decodedBody, err := icmpPacket.Body.Marshal(test.packet.Type.Protocol())
+		require.NoError(t, err)
+		require.Equal(t, expectedBody, decodedBody)
+	}
+}
+
+// TestDecodeBadPackets makes sure decoders don't decode invalid packets
+func TestDecodeBadPackets(t *testing.T) {
+	var (
+		srcIPv4 = net.ParseIP("172.16.0.1")
+		dstIPv4 = net.ParseIP("10.0.0.1")
+	)
+
+	ipLayer := layers.IPv4{
+		Version:  10,
+		SrcIP:    srcIPv4,
+		DstIP:    dstIPv4,
+		Protocol: layers.IPProtocolICMPv4,
+		TTL:      defaultTTL,
+	}
+	icmpLayer := layers.ICMPv4{
+		TypeCode: layers.CreateICMPv4TypeCode(uint8(ipv4.ICMPTypeEcho), 0),
+		Id:       100,
+		Seq:      52819,
+	}
+	wrongIPVersion, err := createPacket(&ipLayer, &icmpLayer, nil, nil)
+	require.NoError(t, err)
+
+	tests := []struct {
+		testCase string
+		packet   []byte
+	}{
+		{
+			testCase: "unknown IP version",
+			packet:   wrongIPVersion,
+		},
+		{
+			testCase: "invalid packet",
+			packet:   []byte("not a packet"),
+		},
+		{
+			testCase: "zero length packet",
+			packet:   []byte{},
+		},
+	}
+
+	ipDecoder := NewIPDecoder()
+	icmpDecoder := NewICMPDecoder()
+	for _, test := range tests {
+		ipPacket, err := ipDecoder.Decode(test.packet)
+		require.Error(t, err)
+		require.Nil(t, ipPacket)
+
+		icmpPacket, err := icmpDecoder.Decode(test.packet)
+		require.Error(t, err)
+		require.Nil(t, icmpPacket)
+	}
+}
+
+func createPacket(ipLayer, secondLayer, thirdLayer gopacket.SerializableLayer, body []byte) ([]byte, error) {
+	payload := gopacket.Payload(body)
+	packet := gopacket.NewSerializeBuffer()
+	var err error
+	if thirdLayer != nil {
+		err = gopacket.SerializeLayers(packet, serializeOpts, ipLayer, secondLayer, thirdLayer, payload)
+	} else {
+		err = gopacket.SerializeLayers(packet, serializeOpts, ipLayer, secondLayer, payload)
+	}
+	if err != nil {
+		return nil, err
+	}
+	return packet.Bytes(), nil
+}
+
+func assertIPLayer(t *testing.T, expected, actual *IP) {
+	require.Equal(t, expected.Src, actual.Src)
+	require.Equal(t, expected.Dst, actual.Dst)
+	require.Equal(t, expected.Protocol, actual.Protocol)
+}
+
+type UDP struct {
+	IP
+	SrcPort, DstPort layers.UDPPort
+}
+
+func (u *UDP) EncodeLayers() ([]gopacket.SerializableLayer, error) {
+	ipLayers, err := u.IP.EncodeLayers()
+	if err != nil {
+		return nil, err
+	}
+	udpLayer := layers.UDP{
+		SrcPort: u.SrcPort,
+		DstPort: u.DstPort,
+	}
+	udpLayer.SetNetworkLayerForChecksum(ipLayers[0].(gopacket.NetworkLayer))
+	return append(ipLayers, &udpLayer), nil
+}

packet/encoder.go

@@ -0,0 +1,39 @@
+package packet
+
+import "github.com/google/gopacket"
+
+var (
+	serializeOpts = gopacket.SerializeOptions{
+		FixLengths:       true,
+		ComputeChecksums: true,
+	}
+)
+
+// RawPacket represents a raw packet or one encoded by Encoder
+type RawPacket struct {
+	Data []byte
+}
+
+type Encoder struct {
+	// buf is reusable because SerializeLayers calls the Clear method before each encoding
+	buf gopacket.SerializeBuffer
+}
+
+func NewEncoder() *Encoder {
+	return &Encoder{
+		buf: gopacket.NewSerializeBuffer(),
+	}
+}
+
+func (e Encoder) Encode(packet Packet) (RawPacket, error) {
+	encodedLayers, err := packet.EncodeLayers()
+	if err != nil {
+		return RawPacket{}, err
+	}
+	if err := gopacket.SerializeLayers(e.buf, serializeOpts, encodedLayers...); err != nil {
+		return RawPacket{}, err
+	}
+	return RawPacket{
+		Data: e.buf.Bytes(),
+	}, nil
+}

packet/flow.go

@@ -0,0 +1,97 @@
+package packet
+
+import (
+	"errors"
+	"net"
+	"net/netip"
+	"sync"
+)
+
+type flowID string
+
+var (
+	ErrFlowNotFound = errors.New("flow not found")
+)
+
+func newFlowID(ip net.IP) flowID {
+	return flowID(ip.String())
+}
+
+type Flow struct {
+	Src       netip.Addr
+	Dst       netip.Addr
+	Responder FlowResponder
+}
+
+func isSameFlow(f1, f2 *Flow) bool {
+	if f1 == nil || f2 == nil {
+		return false
+	}
+	return *f1 == *f2
+}
+
+// FlowResponder sends response packets to the flow
+type FlowResponder interface {
+	// SendPacket returns a packet to the flow. It must not modify the packet,
+	// and after return it must not read the packet
+	SendPacket(pk RawPacket) error
+}
+
+// SrcFlowTracker tracks flow from the perspective of eyeball to origin
+// flowID is the source IP
+type SrcFlowTracker struct {
+	lock  sync.RWMutex
+	flows map[flowID]*Flow
+}
+
+func NewSrcFlowTracker() *SrcFlowTracker {
+	return &SrcFlowTracker{
+		flows: make(map[flowID]*Flow),
+	}
+}
+
+func (sft *SrcFlowTracker) Get(srcIP net.IP) (*Flow, bool) {
+	sft.lock.RLock()
+	defer sft.lock.RUnlock()
+	id := newFlowID(srcIP)
+	flow, ok := sft.flows[id]
+	return flow, ok
+}
+
+// Registers a flow. If shouldReplace = true, replace the current flow
+func (sft *SrcFlowTracker) Register(flow *Flow, shouldReplace bool) (replaced bool) {
+	sft.lock.Lock()
+	defer sft.lock.Unlock()
+	id := flowID(flow.Src.String())
+	currentFlow, ok := sft.flows[id]
+	if !ok {
+		sft.flows[id] = flow
+		return false
+	}
+
+	if shouldReplace && isSameFlow(currentFlow, flow) {
+		sft.flows[id] = flow
+		return true
+	}
+	return false
+}
+
+// Unregisters a flow. If force = true, delete it even if it maps to a different flow
+func (sft *SrcFlowTracker) Unregister(flow *Flow, force bool) (forceDeleted bool) {
+	sft.lock.Lock()
+	defer sft.lock.Unlock()
+	id := flowID(flow.Src.String())
+	currentFlow, ok := sft.flows[id]
+	if !ok {
+		return false
+	}
+	if isSameFlow(currentFlow, flow) {
+		delete(sft.flows, id)
+		return false
+	}
+	if force {
+		delete(sft.flows, id)
+		return true
+	}
+	return false
+}

packet/packet.go

@@ -0,0 +1,115 @@
+package packet
+
+import (
+	"fmt"
+	"net/netip"
+
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/layers"
+	"golang.org/x/net/icmp"
+)
+
+const (
+	defaultTTL    uint8 = 64
+	ipv4HeaderLen       = 20
+	ipv6HeaderLen       = 40
+)
+
+// Packet represents an IP packet or a packet that is encapsulated by IP
+type Packet interface {
+	// IPLayer returns the IP of the packet
+	IPLayer() *IP
+	// EncodeLayers returns the layers that make up this packet. They can be passed to an Encoder to serialize into RawPacket
+	EncodeLayers() ([]gopacket.SerializableLayer, error)
+}
+
+// IP represents a generic IP packet. It can be embedded in more specific IP protocols
+type IP struct {
+	Src      netip.Addr
+	Dst      netip.Addr
+	Protocol layers.IPProtocol
+}
+
+func newIPv4(ipLayer *layers.IPv4) (*IP, error) {
+	src, ok := netip.AddrFromSlice(ipLayer.SrcIP)
+	if !ok {
+		return nil, fmt.Errorf("cannot convert source IP %s to netip.Addr", ipLayer.SrcIP)
+	}
+	dst, ok := netip.AddrFromSlice(ipLayer.DstIP)
+	if !ok {
+		return nil, fmt.Errorf("cannot convert source IP %s to netip.Addr", ipLayer.DstIP)
+	}
+	return &IP{
+		Src:      src,
+		Dst:      dst,
+		Protocol: ipLayer.Protocol,
+	}, nil
+}
+
+func newIPv6(ipLayer *layers.IPv6) (*IP, error) {
+	src, ok := netip.AddrFromSlice(ipLayer.SrcIP)
+	if !ok {
+		return nil, fmt.Errorf("cannot convert source IP %s to netip.Addr", ipLayer.SrcIP)
+	}
+	dst, ok := netip.AddrFromSlice(ipLayer.DstIP)
+	if !ok {
+		return nil, fmt.Errorf("cannot convert source IP %s to netip.Addr", ipLayer.DstIP)
+	}
+	return &IP{
+		Src:      src,
+		Dst:      dst,
+		Protocol: ipLayer.NextHeader,
+	}, nil
+}
+
+func (ip *IP) IPLayer() *IP {
+	return ip
+}
+
+func (ip *IP) isIPv4() bool {
+	return ip.Src.Is4()
+}
+
+func (ip *IP) EncodeLayers() ([]gopacket.SerializableLayer, error) {
+	if ip.isIPv4() {
+		return []gopacket.SerializableLayer{
+			&layers.IPv4{
+				Version:  4,
+				SrcIP:    ip.Src.AsSlice(),
+				DstIP:    ip.Dst.AsSlice(),
+				Protocol: layers.IPProtocol(ip.Protocol),
+				TTL:      defaultTTL,
+			},
+		}, nil
+	} else {
+		return []gopacket.SerializableLayer{
+			&layers.IPv6{
+				Version:    6,
+				SrcIP:      ip.Src.AsSlice(),
+				DstIP:      ip.Dst.AsSlice(),
+				NextHeader: layers.IPProtocol(ip.Protocol),
+				HopLimit:   defaultTTL,
+			},
+		}, nil
+	}
+}
+
+// ICMP represents is an IP packet + ICMP message
+type ICMP struct {
+	*IP
+	*icmp.Message
+}
+
+func (i *ICMP) EncodeLayers() ([]gopacket.SerializableLayer, error) {
+	ipLayers, err := i.IP.EncodeLayers()
+	if err != nil {
+		return nil, err
+	}
+
+	msg, err := i.Marshal(nil)
+	if err != nil {
+		return nil, err
+	}
+	icmpLayer := gopacket.Payload(msg)
+	return append(ipLayers, icmpLayer), nil
+}

packet/session.go

@@ -0,0 +1,8 @@
+package packet
+
+import "github.com/google/uuid"
+
+type Session struct {
+	ID      uuid.UUID
+	Payload []byte
+}