vm_test.go 3.25 KB
Newer Older
zhangweiwei's avatar
init  
zhangweiwei committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144
// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package bpf_test

import (
	"fmt"
	"testing"

	"golang.org/x/net/bpf"
)

var _ bpf.Instruction = unknown{}

type unknown struct{}

func (unknown) Assemble() (bpf.RawInstruction, error) {
	return bpf.RawInstruction{}, nil
}

func TestVMUnknownInstruction(t *testing.T) {
	vm, done, err := testVM(t, []bpf.Instruction{
		bpf.LoadConstant{
			Dst: bpf.RegA,
			Val: 100,
		},
		// Should terminate the program with an error immediately
		unknown{},
		bpf.RetA{},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer done()

	_, err = vm.Run([]byte{
		0xff, 0xff, 0xff, 0xff,
		0xff, 0xff, 0xff, 0xff,
		0x00, 0x00,
	})
	if errStr(err) != "unknown Instruction at index 1: bpf_test.unknown" {
		t.Fatalf("unexpected error while running program: %v", err)
	}
}

func TestVMNoReturnInstruction(t *testing.T) {
	_, _, err := testVM(t, []bpf.Instruction{
		bpf.LoadConstant{
			Dst: bpf.RegA,
			Val: 1,
		},
	})
	if errStr(err) != "BPF program must end with RetA or RetConstant" {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestVMNoInputInstructions(t *testing.T) {
	_, _, err := testVM(t, []bpf.Instruction{})
	if errStr(err) != "one or more Instructions must be specified" {
		t.Fatalf("unexpected error: %v", err)
	}
}

// ExampleNewVM demonstrates usage of a VM, using an Ethernet frame
// as input and checking its EtherType to determine if it should be accepted.
func ExampleNewVM() {
	// Offset | Length | Comment
	// -------------------------
	//   00   |   06   | Ethernet destination MAC address
	//   06   |   06   | Ethernet source MAC address
	//   12   |   02   | Ethernet EtherType
	const (
		etOff = 12
		etLen = 2

		etARP = 0x0806
	)

	// Set up a VM to filter traffic based on if its EtherType
	// matches the ARP EtherType.
	vm, err := bpf.NewVM([]bpf.Instruction{
		// Load EtherType value from Ethernet header
		bpf.LoadAbsolute{
			Off:  etOff,
			Size: etLen,
		},
		// If EtherType is equal to the ARP EtherType, jump to allow
		// packet to be accepted
		bpf.JumpIf{
			Cond:     bpf.JumpEqual,
			Val:      etARP,
			SkipTrue: 1,
		},
		// EtherType does not match the ARP EtherType
		bpf.RetConstant{
			Val: 0,
		},
		// EtherType matches the ARP EtherType, accept up to 1500
		// bytes of packet
		bpf.RetConstant{
			Val: 1500,
		},
	})
	if err != nil {
		panic(fmt.Sprintf("failed to load BPF program: %v", err))
	}

	// Create an Ethernet frame with the ARP EtherType for testing
	frame := []byte{
		0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		0x00, 0x11, 0x22, 0x33, 0x44, 0x55,
		0x08, 0x06,
		// Payload omitted for brevity
	}

	// Run our VM's BPF program using the Ethernet frame as input
	out, err := vm.Run(frame)
	if err != nil {
		panic(fmt.Sprintf("failed to accept Ethernet frame: %v", err))
	}

	// BPF VM can return a byte count greater than the number of input
	// bytes, so trim the output to match the input byte length
	if out > len(frame) {
		out = len(frame)
	}

	fmt.Printf("out: %d bytes", out)

	// Output:
	// out: 14 bytes
}

// errStr returns the string representation of an error, or
// "<nil>" if it is nil.
func errStr(err error) string {
	if err == nil {
		return "<nil>"
	}

	return err.Error()
}