github.com/Andyfoo/golang/x/net@v0.0.0-20190901054642-57c1bf301704/ipv4/bpf_test.go (about)

     1  // Copyright 2016 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package ipv4_test
     6  
     7  import (
     8  	"net"
     9  	"runtime"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/Andyfoo/golang/x/net/bpf"
    14  	"github.com/Andyfoo/golang/x/net/ipv4"
    15  )
    16  
    17  func TestBPF(t *testing.T) {
    18  	if runtime.GOOS != "linux" {
    19  		t.Skipf("not supported on %s", runtime.GOOS)
    20  	}
    21  
    22  	l, err := net.ListenPacket("udp4", "127.0.0.1:0")
    23  	if err != nil {
    24  		t.Fatal(err)
    25  	}
    26  	defer l.Close()
    27  
    28  	p := ipv4.NewPacketConn(l)
    29  
    30  	// This filter accepts UDP packets whose first payload byte is
    31  	// even.
    32  	prog, err := bpf.Assemble([]bpf.Instruction{
    33  		// Load the first byte of the payload (skipping UDP header).
    34  		bpf.LoadAbsolute{Off: 8, Size: 1},
    35  		// Select LSB of the byte.
    36  		bpf.ALUOpConstant{Op: bpf.ALUOpAnd, Val: 1},
    37  		// Byte is even?
    38  		bpf.JumpIf{Cond: bpf.JumpEqual, Val: 0, SkipFalse: 1},
    39  		// Accept.
    40  		bpf.RetConstant{Val: 4096},
    41  		// Ignore.
    42  		bpf.RetConstant{Val: 0},
    43  	})
    44  	if err != nil {
    45  		t.Fatalf("compiling BPF: %s", err)
    46  	}
    47  
    48  	if err = p.SetBPF(prog); err != nil {
    49  		t.Fatalf("attaching filter to Conn: %s", err)
    50  	}
    51  
    52  	s, err := net.Dial("udp4", l.LocalAddr().String())
    53  	if err != nil {
    54  		t.Fatal(err)
    55  	}
    56  	defer s.Close()
    57  	go func() {
    58  		for i := byte(0); i < 10; i++ {
    59  			s.Write([]byte{i})
    60  		}
    61  	}()
    62  
    63  	l.SetDeadline(time.Now().Add(2 * time.Second))
    64  	seen := make([]bool, 5)
    65  	for {
    66  		var b [512]byte
    67  		n, _, err := l.ReadFrom(b[:])
    68  		if err != nil {
    69  			t.Fatalf("reading from listener: %s", err)
    70  		}
    71  		if n != 1 {
    72  			t.Fatalf("unexpected packet length, want 1, got %d", n)
    73  		}
    74  		if b[0] >= 10 {
    75  			t.Fatalf("unexpected byte, want 0-9, got %d", b[0])
    76  		}
    77  		if b[0]%2 != 0 {
    78  			t.Fatalf("got odd byte %d, wanted only even bytes", b[0])
    79  		}
    80  		seen[b[0]/2] = true
    81  
    82  		seenAll := true
    83  		for _, v := range seen {
    84  			if !v {
    85  				seenAll = false
    86  				break
    87  			}
    88  		}
    89  		if seenAll {
    90  			break
    91  		}
    92  	}
    93  }