github.com/psiphon-labs/goarista@v0.0.0-20160825065156-d002785f4c67/dscp/listen_test.go (about)

     1  // Copyright (C) 2016  Arista Networks, Inc.
     2  // Use of this source code is governed by the Apache License 2.0
     3  // that can be found in the COPYING file.
     4  
     5  package dscp_test
     6  
     7  import (
     8  	"net"
     9  	"testing"
    10  
    11  	"github.com/aristanetworks/goarista/dscp"
    12  )
    13  
    14  func TestListenTCPWithTOS(t *testing.T) {
    15  	testListenTCPWithTOS(t, "127.0.0.1")
    16  	testListenTCPWithTOS(t, "::1")
    17  }
    18  
    19  func testListenTCPWithTOS(t *testing.T, ip string) {
    20  	// Note: This test doesn't actually verify that the connection uses the
    21  	// desired TOS byte, because that's kinda hard to check, but at least it
    22  	// verifies that we return a usable TCPListener.
    23  	addr := &net.TCPAddr{IP: net.ParseIP(ip), Port: 0}
    24  	listen, err := dscp.ListenTCPWithTOS(addr, 40)
    25  	if err != nil {
    26  		t.Fatal(err)
    27  	}
    28  	defer listen.Close()
    29  
    30  	done := make(chan struct{})
    31  	go func() {
    32  		conn, err := listen.Accept()
    33  		if err != nil {
    34  			t.Fatal(err)
    35  		}
    36  		defer conn.Close()
    37  		buf := []byte{'!'}
    38  		conn.Write(buf)
    39  		n, err := conn.Read(buf)
    40  		if n != 1 || err != nil {
    41  			t.Fatalf("Read returned %d / %s", n, err)
    42  		} else if buf[0] != '!' {
    43  			t.Fatalf("Expected to read '!' but got %q", buf)
    44  		}
    45  		close(done)
    46  	}()
    47  
    48  	conn, err := net.Dial(listen.Addr().Network(), listen.Addr().String())
    49  	if err != nil {
    50  		t.Fatal("Connection failed:", err)
    51  	}
    52  	defer conn.Close()
    53  	buf := make([]byte, 1)
    54  	n, err := conn.Read(buf)
    55  	if n != 1 || err != nil {
    56  		t.Fatalf("Read returned %d / %s", n, err)
    57  	} else if buf[0] != '!' {
    58  		t.Fatalf("Expected to read '!' but got %q", buf)
    59  	}
    60  	conn.Write(buf)
    61  	<-done
    62  }