Merge pull request #67 from chantra/class_replace

Class replace/change

class_linux.go

@@ -9,24 +9,38 @@ import (
 // ClassDel will delete a class from the system.
 // Equivalent to: `tc class del $class`
 func ClassDel(class Class) error {
-	req := nl.NewNetlinkRequest(syscall.RTM_DELTCLASS, syscall.NLM_F_ACK)
-	base := class.Attrs()
-	msg := &nl.TcMsg{
-		Family:  nl.FAMILY_ALL,
-		Ifindex: int32(base.LinkIndex),
-		Handle:  base.Handle,
-		Parent:  base.Parent,
-	}
-	req.AddData(msg)
+	return classModify(syscall.RTM_DELTCLASS, 0, class)
+}
 
-	_, err := req.Execute(syscall.NETLINK_ROUTE, 0)
-	return err
+// ClassChange will change a class in place
+// Equivalent to: `tc class change $class`
+// The parent and handle MUST NOT be changed.
+
+func ClassChange(class Class) error {
+	return classModify(syscall.RTM_NEWTCLASS, 0, class)
+}
+
+// ClassReplace will replace a class to the system.
+// quivalent to: `tc class replace $class`
+// The handle MAY be changed.
+// If a class already exist with this parent/handle pair, the class is changed.
+// If a class does not already exist with this parent/handle, a new class is created.
+func ClassReplace(class Class) error {
+	return classModify(syscall.RTM_NEWTCLASS, syscall.NLM_F_CREATE, class)
 }
 
 // ClassAdd will add a class to the system.
 // Equivalent to: `tc class add $class`
 func ClassAdd(class Class) error {
-	req := nl.NewNetlinkRequest(syscall.RTM_NEWTCLASS, syscall.NLM_F_CREATE|syscall.NLM_F_EXCL|syscall.NLM_F_ACK)
+	return classModify(
+		syscall.RTM_NEWTCLASS,
+		syscall.NLM_F_CREATE|syscall.NLM_F_EXCL,
+		class,
+	)
+}
+
+func classModify(cmd, flags int, class Class) error {
+	req := nl.NewNetlinkRequest(cmd, flags|syscall.NLM_F_ACK)
 	base := class.Attrs()
 	msg := &nl.TcMsg{
 		Family:  nl.FAMILY_ALL,
@@ -35,6 +49,17 @@ func ClassAdd(class Class) error {
 		Parent:  base.Parent,
 	}
 	req.AddData(msg)
+
+	if cmd != syscall.RTM_DELTCLASS {
+		if err := classPayload(req, class); err != nil {
+			return err
+		}
+	}
+	_, err := req.Execute(syscall.NETLINK_ROUTE, 0)
+	return err
+}
+
+func classPayload(req *nl.NetlinkRequest, class Class) error {
 	req.AddData(nl.NewRtAttr(nl.TCA_KIND, nl.ZeroTerminated(class.Type())))
 
 	options := nl.NewRtAttr(nl.TCA_OPTIONS, nil)
@@ -51,13 +76,12 @@ func ClassAdd(class Class) error {
 		nl.NewRtAttrChild(options, nl.TCA_HTB_PARMS, opt.Serialize())
 	}
 	req.AddData(options)
-	_, err := req.Execute(syscall.NETLINK_ROUTE, 0)
-	return err
+	return nil
 }
 
 // ClassList gets a list of classes in the system.
 // Equivalent to: `tc class show`.
-// Generally retunrs nothing if link and parent are not specified.
+// Generally returns nothing if link and parent are not specified.
 func ClassList(link Link, parent uint32) ([]Class, error) {
 	req := nl.NewNetlinkRequest(syscall.RTM_GETTCLASS, syscall.NLM_F_DUMP)
 	msg := &nl.TcMsg{

class_test.go

@@ -55,7 +55,7 @@ func TestClassAddDel(t *testing.T) {
 	if err := ClassAdd(class); err != nil {
 		t.Fatal(err)
 	}
-	classes, err := ClassList(link, MakeHandle(0xffff, 2))
+	classes, err := ClassList(link, MakeHandle(0xffff, 0))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -159,3 +159,248 @@ func TestClassAddDel(t *testing.T) {
 		t.Fatal("Failed to remove qdisc")
 	}
 }
+
+func TestHtbClassAddHtbClassChangeDel(t *testing.T) {
+	/**
+	This test first set up a interface ans set up a Htb qdisc
+	A HTB class is attach to it and a Netem qdisc is attached to that class
+	Next, we test changing the HTB class in place and confirming the Netem is
+	still attached. We also check that invoting ClassChange with a non-existing
+	class will fail.
+	Finally, we test ClassReplace. We confirm it correctly behave like
+	ClassChange when the parent/handle pair exists and that it will create a
+	new class if the handle is modified.
+	*/
+	tearDown := setUpNetlinkTest(t)
+	defer tearDown()
+	if err := LinkAdd(&Ifb{LinkAttrs{Name: "foo"}}); err != nil {
+		t.Fatal(err)
+	}
+	link, err := LinkByName("foo")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := LinkSetUp(link); err != nil {
+		t.Fatal(err)
+	}
+	attrs := QdiscAttrs{
+		LinkIndex: link.Attrs().Index,
+		Handle:    MakeHandle(0xffff, 0),
+		Parent:    HANDLE_ROOT,
+	}
+	qdisc := NewHtb(attrs)
+	if err := QdiscAdd(qdisc); err != nil {
+		t.Fatal(err)
+	}
+	qdiscs, err := QdiscList(link)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(qdiscs) != 1 {
+		t.Fatal("Failed to add qdisc")
+	}
+	_, ok := qdiscs[0].(*Htb)
+	if !ok {
+		t.Fatal("Qdisc is the wrong type")
+	}
+
+	classattrs := ClassAttrs{
+		LinkIndex: link.Attrs().Index,
+		Parent:    MakeHandle(0xffff, 0),
+		Handle:    MakeHandle(0xffff, 2),
+	}
+
+	htbclassattrs := HtbClassAttrs{
+		Rate:    1234000,
+		Cbuffer: 1690,
+	}
+	class := NewHtbClass(classattrs, htbclassattrs)
+	if err := ClassAdd(class); err != nil {
+		t.Fatal(err)
+	}
+	classes, err := ClassList(link, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(classes) != 1 {
+		t.Fatal("Failed to add class")
+	}
+
+	htb, ok := classes[0].(*HtbClass)
+	if !ok {
+		t.Fatal("Class is the wrong type")
+	}
+
+	qattrs := QdiscAttrs{
+		LinkIndex: link.Attrs().Index,
+		Handle:    MakeHandle(0x2, 0),
+		Parent:    MakeHandle(0xffff, 2),
+	}
+	nattrs := NetemQdiscAttrs{
+		Latency:     20000,
+		Loss:        23.4,
+		Duplicate:   14.3,
+		LossCorr:    8.34,
+		Jitter:      1000,
+		DelayCorr:   12.3,
+		ReorderProb: 23.4,
+		CorruptProb: 10.0,
+		CorruptCorr: 10,
+	}
+	qdiscnetem := NewNetem(qattrs, nattrs)
+	if err := QdiscAdd(qdiscnetem); err != nil {
+		t.Fatal(err)
+	}
+
+	qdiscs, err = QdiscList(link)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(qdiscs) != 2 {
+		t.Fatal("Failed to add qdisc")
+	}
+
+	_, ok = qdiscs[1].(*Netem)
+	if !ok {
+		t.Fatal("Qdisc is the wrong type")
+	}
+
+	// Change
+	// For change to work, the handle and parent cannot be changed.
+
+	// First, test it fails if we change the Handle.
+	old_handle := classattrs.Handle
+	classattrs.Handle = MakeHandle(0xffff, 3)
+	class = NewHtbClass(classattrs, htbclassattrs)
+	if err := ClassChange(class); err == nil {
+		t.Fatal("ClassChange should not work when using a different handle.")
+	}
+	// It should work with the same handle
+	classattrs.Handle = old_handle
+	htbclassattrs.Rate = 4321000
+	class = NewHtbClass(classattrs, htbclassattrs)
+	if err := ClassChange(class); err != nil {
+		t.Fatal(err)
+	}
+
+	classes, err = ClassList(link, MakeHandle(0xffff, 0))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(classes) != 1 {
+		t.Fatalf(
+			"1 class expected, %d found",
+			len(classes),
+		)
+	}
+
+	htb, ok = classes[0].(*HtbClass)
+	if !ok {
+		t.Fatal("Class is the wrong type")
+	}
+	// Verify that the rate value has changed.
+	if htb.Rate != class.Rate {
+		t.Fatal("Rate did not get changed while changing the class.")
+	}
+
+	// Check that we still have the netem child qdisc
+	qdiscs, err = QdiscList(link)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(qdiscs) != 2 {
+		t.Fatalf("2 qdisc expected, %d found", len(qdiscs))
+	}
+	_, ok = qdiscs[0].(*Htb)
+	if !ok {
+		t.Fatal("Qdisc is the wrong type")
+	}
+
+	_, ok = qdiscs[1].(*Netem)
+	if !ok {
+		t.Fatal("Qdisc is the wrong type")
+	}
+
+	// Replace
+	// First replace by keeping the same handle, class will be changed.
+	// Then, replace by providing a new handle, n new class will be created.
+
+	// Replace acting as Change
+	class = NewHtbClass(classattrs, htbclassattrs)
+	if err := ClassReplace(class); err != nil {
+		t.Fatal("Failed to replace class that is existing.")
+	}
+
+	classes, err = ClassList(link, MakeHandle(0xffff, 0))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(classes) != 1 {
+		t.Fatalf(
+			"1 class expected, %d found",
+			len(classes),
+		)
+	}
+
+	htb, ok = classes[0].(*HtbClass)
+	if !ok {
+		t.Fatal("Class is the wrong type")
+	}
+	// Verify that the rate value has changed.
+	if htb.Rate != class.Rate {
+		t.Fatal("Rate did not get changed while changing the class.")
+	}
+
+	// It should work with the same handle
+	classattrs.Handle = MakeHandle(0xffff, 3)
+	class = NewHtbClass(classattrs, htbclassattrs)
+	if err := ClassReplace(class); err != nil {
+		t.Fatal(err)
+	}
+
+	classes, err = ClassList(link, MakeHandle(0xffff, 0))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(classes) != 2 {
+		t.Fatalf(
+			"2 classes expected, %d found",
+			len(classes),
+		)
+	}
+
+	htb, ok = classes[1].(*HtbClass)
+	if !ok {
+		t.Fatal("Class is the wrong type")
+	}
+	// Verify that the rate value has changed.
+	if htb.Rate != class.Rate {
+		t.Fatal("Rate did not get changed while changing the class.")
+	}
+
+	// Deletion
+	for _, class := range classes {
+		if err := ClassDel(class); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	classes, err = ClassList(link, MakeHandle(0xffff, 0))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(classes) != 0 {
+		t.Fatal("Failed to remove class")
+	}
+	if err := QdiscDel(qdisc); err != nil {
+		t.Fatal(err)
+	}
+	qdiscs, err = QdiscList(link)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(qdiscs) != 0 {
+		t.Fatal("Failed to remove qdisc")
+	}
+}