Merge pull request #7 from tsuna/master

dscp: Add a package to help create listening sockets with DSCP.

dscp/listen.go

@@ -0,0 +1,17 @@
+// Copyright (C) 2016  Arista Networks, Inc.
+// Use of this source code is governed by the Apache License 2.0
+// that can be found in the COPYING file.
+
+// Package dscp provides helper functions to apply DSCP / ECN / CoS flags to sockets.
+package dscp
+
+import (
+	"net"
+)
+
+// ListenTCPWithTOS is similar to net.ListenTCP but with the socket configured
+// to the use the given ToS (Type of Service), to specify DSCP / ECN / class
+// of service flags to use for incoming connections.
+func ListenTCPWithTOS(address *net.TCPAddr, tos byte) (*net.TCPListener, error) {
+	return listenTCPWithTOS(address, tos)
+}

dscp/listen_test.go

@@ -0,0 +1,62 @@
+// Copyright (C) 2016  Arista Networks, Inc.
+// Use of this source code is governed by the Apache License 2.0
+// that can be found in the COPYING file.
+
+package dscp_test
+
+import (
+	"net"
+	"testing"
+
+	"github.com/aristanetworks/goarista/dscp"
+)
+
+func TestListenTCPWithTOS(t *testing.T) {
+	testListenTCPWithTOS(t, "127.0.0.1")
+	testListenTCPWithTOS(t, "::1")
+}
+
+func testListenTCPWithTOS(t *testing.T, ip string) {
+	// Note: This test doesn't actually verify that the connection uses the
+	// desired TOS byte, because that's kinda hard to check, but at least it
+	// verifies that we return a usable TCPListener.
+	addr := &net.TCPAddr{IP: net.ParseIP(ip), Port: 0}
+	listen, err := dscp.ListenTCPWithTOS(addr, 40)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer listen.Close()
+
+	done := make(chan struct{})
+	go func() {
+		conn, err := listen.Accept()
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer conn.Close()
+		buf := []byte{'!'}
+		conn.Write(buf)
+		n, err := conn.Read(buf)
+		if n != 1 || err != nil {
+			t.Fatalf("Read returned %d / %s", n, err)
+		} else if buf[0] != '!' {
+			t.Fatalf("Expected to read '!' but got %q", buf)
+		}
+		close(done)
+	}()
+
+	conn, err := net.Dial(listen.Addr().Network(), listen.Addr().String())
+	if err != nil {
+		t.Fatal("Connection failed:", err)
+	}
+	defer conn.Close()
+	buf := make([]byte, 1)
+	n, err := conn.Read(buf)
+	if n != 1 || err != nil {
+		t.Fatalf("Read returned %d / %s", n, err)
+	} else if buf[0] != '!' {
+		t.Fatalf("Expected to read '!' but got %q", buf)
+	}
+	conn.Write(buf)
+	<-done
+}

dscp/listen_unix.go

@@ -0,0 +1,37 @@
+// Copyright (C) 2016  Arista Networks, Inc.
+// Use of this source code is governed by the Apache License 2.0
+// that can be found in the COPYING file.
+
+package dscp
+
+import (
+	"net"
+	"os"
+	"reflect"
+	"syscall"
+)
+
+func listenTCPWithTOS(address *net.TCPAddr, tos byte) (*net.TCPListener, error) {
+	lsnr, err := net.ListenTCP("tcp", address)
+	if err != nil {
+		return nil, err
+	}
+	// This works for the UNIX implementation of netFD, i.e. not on Windows and Plan9.
+	// This kludge is needed until https://github.com/golang/go/issues/9661 is fixed.
+	fd := int(reflect.ValueOf(lsnr).Elem().FieldByName("fd").Elem().FieldByName("sysfd").Int())
+	var proto, optname int
+	if address.IP.To4() != nil {
+		proto = syscall.IPPROTO_IP
+		optname = syscall.IP_TOS
+	} else {
+		proto = syscall.IPPROTO_IPV6
+		optname = syscall.IPV6_TCLASS
+	}
+	err = syscall.SetsockoptInt(fd, proto, optname, int(tos))
+	if err != nil {
+		lsnr.Close()
+		return nil, os.NewSyscallError("setsockopt", err)
+	}
+
+	return lsnr, nil
+}