Port dscp back from master

controller/internal/supervisor/interfaces.go

@@ -54,4 +54,7 @@ type Implementor interface {
 
 	// ACLProvider returns the ACL provider used by the implementor
 	ACLProvider() []provider.IptablesProvider
+
+	// CreateCustomRulesChain creates a custom rules chain if it doesnt exist
+	CreateCustomRulesChain() error
 }

controller/internal/supervisor/iptablesctrl/constants_nonwindows.go

@@ -11,4 +11,9 @@ const (
 
 	portSetIpsetType      = ""
 	proxySetPortIpsetType = ""
+
+	customQOSChainNFHook = "POSTROUTING"
+	customQOSChainTable  = "mangle"
+	// CustomQOSChain is the name of the chain where users can install custom QOS rules
+	CustomQOSChain = "POST-CUSTOM-QOS"
 )

controller/internal/supervisor/iptablesctrl/constants_windows.go

@@ -11,4 +11,9 @@ const (
 
 	portSetIpsetType      = "hash:net"
 	proxySetPortIpsetType = "hash:port"
+
+	customQOSChainNFHook = "POSTROUTING"
+	customQOSChainTable  = "mangle"
+	// CustomQOSChain is the name of the chain where users can install custom QOS rules
+	CustomQOSChain = "POST-CUSTOM-QOS"
 )

controller/internal/supervisor/iptablesctrl/instance.go

@@ -3,6 +3,7 @@ package iptablesctrl
 import (
 	"context"
 	"fmt"
+	"strings"
 	"sync"
 
 	"go.aporeto.io/trireme-lib/controller/constants"
@@ -130,6 +131,66 @@ func (i *Instance) CleanUp() error {
 	return nil
 }
 
+// CreateCustomRulesChain creates a custom rules chain if it doesnt exist
+func (i *Instance) CreateCustomRulesChain() error {
+	nonbatchedv4tableprovider, _ := provider.NewGoIPTablesProviderV4([]string{}, CustomQOSChain)
+	nonbatchedv6tableprovider, _ := provider.NewGoIPTablesProviderV6([]string{}, CustomQOSChain)
+	err := nonbatchedv4tableprovider.NewChain(customQOSChainTable, CustomQOSChain)
+	if err != nil {
+		zap.L().Debug("Chain already exists", zap.Error(err))
+
+	}
+	postroutingchainrulesv4, err := nonbatchedv4tableprovider.ListRules(customQOSChainTable, customQOSChainNFHook)
+	if err != nil {
+		return err
+	}
+	checkCustomRulesv4 := func() bool {
+		for _, rule := range postroutingchainrulesv4 {
+			if strings.Contains(rule, CustomQOSChain) {
+				return true
+			}
+		}
+		return false
+	}
+	if !checkCustomRulesv4() {
+		if err := nonbatchedv4tableprovider.Insert(customQOSChainTable, customQOSChainNFHook, 1,
+			"-m", "addrtype",
+			"--src-type", "LOCAL",
+			"-j", CustomQOSChain,
+		); err != nil {
+			zap.L().Debug("Unable to create ipv4 custom rule", zap.Error(err))
+		}
+	}
+
+	err = nonbatchedv6tableprovider.NewChain(customQOSChainTable, CustomQOSChain)
+	if err != nil {
+		zap.L().Debug("Chain already exists", zap.Error(err))
+	}
+	postroutingchainrulesv6, err := nonbatchedv6tableprovider.ListRules(customQOSChainTable, customQOSChainNFHook)
+	if err != nil {
+		return err
+	}
+	checkCustomRulesv6 := func() bool {
+		for _, rule := range postroutingchainrulesv6 {
+			if strings.Contains(rule, CustomQOSChain) {
+				return true
+			}
+		}
+		return false
+	}
+	if !checkCustomRulesv6() {
+		if err := nonbatchedv6tableprovider.Append(customQOSChainTable, customQOSChainNFHook,
+			"-m", "addrtype",
+			"--src-type", "LOCAL",
+			"-j", CustomQOSChain,
+		); err != nil {
+			zap.L().Debug("Unable to create ipv6 custom rule", zap.Error(err))
+		}
+	}
+
+	return nil
+}
+
 // NewInstance creates a new iptables controller instance
 func NewInstance(fqc *fqconfig.FilterQueue, mode constants.ModeType, aclmanager ipsetmanager.ACLManager, ipv6Enabled bool, ebpf ebpf.BPFModule) (*Instance, error) {
 

controller/internal/supervisor/iptablesctrl/ipv4.go

@@ -30,7 +30,7 @@ func init() {
 // GetIPv4Impl creates the instance of ipv4 struct which implements the interface
 // ipImpl
 func GetIPv4Impl() (IPImpl, error) {
-	ipt, err := provider.NewGoIPTablesProviderV4([]string{"mangle"})
+	ipt, err := provider.NewGoIPTablesProviderV4([]string{"mangle"}, CustomQOSChain)
 	if err != nil {
 		return nil, fmt.Errorf("unable to initialize iptables provider: %s", err)
 	}
@@ -109,3 +109,7 @@ func (i *ipv4) RetrieveTable() map[string]map[string][]string {
 func (i *ipv4) ResetRules(subs string) error {
 	return i.ipt.ResetRules(subs)
 }
+
+func (i *ipv4) ListRules(table, chain string) ([]string, error) {
+	return i.ipt.ListRules(table, chain)
+}

controller/internal/supervisor/iptablesctrl/ipv6.go

@@ -134,3 +134,7 @@ func (i *ipv6) ResetRules(subs string) error {
 
 	return i.ipt.ResetRules(subs)
 }
+
+func (i *ipv6) ListRules(table, chain string) ([]string, error) {
+	return i.ipt.ListRules(table, chain)
+}

controller/internal/supervisor/iptablesctrl/ipv6_nonwindows.go

@@ -9,7 +9,7 @@ import (
 // GetIPv6Impl creates the instance of ipv6 struct which implements
 // the interface ipImpl
 func GetIPv6Impl(ipv6Enabled bool) (IPImpl, error) {
-	ipt, err := provider.NewGoIPTablesProviderV6([]string{"mangle"})
+	ipt, err := provider.NewGoIPTablesProviderV6([]string{"mangle"}, CustomQOSChain)
 	if err == nil {
 		// test if the system supports ip6tables
 		if _, err = ipt.ListChains("mangle"); err == nil {

controller/internal/supervisor/iptablesctrl/ipv6_windows.go

@@ -9,7 +9,7 @@ import (
 // GetIPv6Impl creates the instance of ipv6 struct which implements
 // the interface ipImpl
 func GetIPv6Impl(ipv6Enabled bool) (IPImpl, error) {
-	if ipt, err := provider.NewGoIPTablesProviderV6(nil); err == nil {
+	if ipt, err := provider.NewGoIPTablesProviderV6(nil, ""); err == nil {
 		return &ipv6{ipt: ipt, ipv6Enabled: ipv6Enabled}, nil
 	}
 	return &ipv6{ipt: nil, ipv6Enabled: false}, nil

controller/internal/supervisor/supervisor.go

@@ -119,6 +119,9 @@ func (s *Config) Run(ctx context.Context) error {
 	if s.service != nil {
 		s.service.Initialize(s.filterQueue, s.impl.ACLProvider())
 	}
+	if err := s.impl.CreateCustomRulesChain(); err != nil {
+		return err
+	}
 
 	return nil
 }

controller/internal/supervisor/supervisor_test.go

@@ -283,6 +283,7 @@ func TestStart(t *testing.T) {
 		Convey("When I try to start it and the implementor works", func() {
 			impl.EXPECT().Run(gomock.Any()).Return(nil)
 			impl.EXPECT().SetTargetNetworks(&runtime.Configuration{TCPTargetNetworks: []string{"172.17.0.0/16"}}).Return(nil)
+			impl.EXPECT().CreateCustomRulesChain().Return(nil)
 			err := s.Run(context.Background())
 			Convey("I should get no errors", func() {
 				So(err, ShouldBeNil)
@@ -326,6 +327,7 @@ func TestStop(t *testing.T) {
 		Convey("When I try to start it and the implementor works", func() {
 			impl.EXPECT().Run(gomock.Any()).Return(nil)
 			impl.EXPECT().SetTargetNetworks(&runtime.Configuration{TCPTargetNetworks: []string{"172.17.0.0/16"}}).Return(nil)
+			impl.EXPECT().CreateCustomRulesChain().Return(nil)
 			err := s.Run(context.Background())
 			Convey("I should get no errors", func() {
 				So(err, ShouldBeNil)
@@ -362,6 +364,7 @@ func TestEnableIPTablesPacketTracing(t *testing.T) {
 		Convey("When I try to start it and the implementor works", func() {
 			impl.EXPECT().Run(gomock.Any()).Return(nil)
 			impl.EXPECT().SetTargetNetworks(&runtime.Configuration{TCPTargetNetworks: []string{"172.17.0.0/16"}}).Return(nil)
+			impl.EXPECT().CreateCustomRulesChain().Return(nil)
 			err := s.Run(context.Background())
 			Convey("I should get no errors", func() {
 				So(err, ShouldBeNil)

controller/pkg/aclprovider/iptablesprovider.go

@@ -41,6 +41,8 @@ type BaseIPTables interface {
 	DeleteChain(table, chain string) error
 	// NewChain creates a new chain
 	NewChain(table, chain string) error
+	// ListRules lists the rules in the table/chain passed to it
+	ListRules(table, chain string) ([]string, error)
 }
 
 // BatchProvider uses iptables-restore to program ACLs
@@ -52,7 +54,8 @@ type BatchProvider struct {
 	batchTables map[string]bool
 
 	// Allowing for custom commit functions for testing
-	commitFunc func(buf *bytes.Buffer) error
+	commitFunc  func(buf *bytes.Buffer) error
+	customChain string
 	sync.Mutex
 	cmd        string
 	restoreCmd string
@@ -88,7 +91,7 @@ func TestIptablesPinned(bpf string) error {
 
 // NewGoIPTablesProviderV4 returns an IptablesProvider interface based on the go-iptables
 // external package.
-func NewGoIPTablesProviderV4(batchTables []string) (IptablesProvider, error) {
+func NewGoIPTablesProviderV4(batchTables []string, customChain string) (IptablesProvider, error) {
 
 	batchTablesMap := map[string]bool{}
 	for _, t := range batchTables {
@@ -101,6 +104,7 @@ func NewGoIPTablesProviderV4(batchTables []string) (IptablesProvider, error) {
 		batchTables: batchTablesMap,
 		restoreCmd:  restoreCmdV4,
 		saveCmd:     saveCmdV4,
+		customChain: customChain,
 		quote:       true,
 	}
 
@@ -111,7 +115,7 @@ func NewGoIPTablesProviderV4(batchTables []string) (IptablesProvider, error) {
 
 // NewGoIPTablesProviderV6 returns an IptablesProvider interface based on the go-iptables
 // external package.
-func NewGoIPTablesProviderV6(batchTables []string) (IptablesProvider, error) {
+func NewGoIPTablesProviderV6(batchTables []string, customChain string) (IptablesProvider, error) {
 
 	batchTablesMap := map[string]bool{}
 	for _, t := range batchTables {
@@ -124,6 +128,7 @@ func NewGoIPTablesProviderV6(batchTables []string) (IptablesProvider, error) {
 		batchTables: batchTablesMap,
 		restoreCmd:  restoreCmdV6,
 		saveCmd:     saveCmdV6,
+		customChain: customChain,
 		quote:       true,
 	}
 
@@ -469,6 +474,8 @@ func (b *BatchProvider) createDataBuffer() (*bytes.Buffer, error) {
 				}
 			}
 		}
+		customChainRules, _ := b.saveCustomChainRules()
+		fmt.Fprintf(buf, "%s\n", customChainRules.String())
 		if _, err := fmt.Fprintf(buf, "COMMIT\n"); err != nil {
 			return nil, err
 		}
@@ -538,3 +545,49 @@ func (b *BatchProvider) ResetRules(subs string) error {
 
 	return b.commitFunc(buf)
 }
+
+func (b *BatchProvider) saveCustomChainRules() (*bytes.Buffer, error) {
+	var out []byte
+	var err error
+
+	cmd := exec.Command("aporeto-iptables", b.saveCmd)
+	if out, err = cmd.CombinedOutput(); err != nil {
+		zap.L().Error("Failed to get iptables-save command", zap.Error(err),
+			zap.String("Output", string(out)))
+		return nil, err
+	}
+
+	s := string(out)
+	rules := strings.Split(s, "\n")
+
+	var filterRules []string
+
+	for _, rule := range rules {
+		if strings.Contains(rule, b.customChain) {
+			filterRules = append(filterRules, rule)
+		}
+	}
+
+	combineRules := strings.Join(filterRules, "\n")
+	return bytes.NewBufferString(combineRules), nil
+
+}
+
+// ListRules lists the rules in the table/chain passed to it
+func (b *BatchProvider) ListRules(table, chain string) ([]string, error) {
+	var cmd *exec.Cmd
+
+	if chain != "" {
+		cmd = exec.Command("aporeto-iptables", "iptables", "--wait", "-t", table, "-L", chain)
+	} else {
+		cmd = exec.Command("aporeto-iptables", "iptables", "--wait", "-t", table, "-L")
+	}
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		zap.L().Error("Failed to get rules", zap.Error(err), zap.String("table", table), zap.String("chain", chain))
+		return []string{}, err
+	}
+	rules := strings.Split(string(out), "\n")
+	return rules, nil
+
+}

controller/pkg/aclprovider/iptablesprovider_test.go

@@ -341,10 +341,10 @@ func TestDeleteChain(t *testing.T) {
 }
 
 func TestProvider(t *testing.T) {
-	b, err := NewGoIPTablesProviderV4([]string{})
+	b, err := NewGoIPTablesProviderV4([]string{}, "")
 	assert.Equal(t, b != nil, true, "go iptables should not be nil")
 	assert.Equal(t, err == nil, true, "error should be nil")
-	b, err = NewGoIPTablesProviderV6([]string{})
+	b, err = NewGoIPTablesProviderV6([]string{}, "")
 	assert.Equal(t, b != nil, true, "go iptables should not be nil")
 	assert.Equal(t, err == nil, true, "error should be nil")
 }

controller/pkg/aclprovider/iptablesprovider_windows.go

@@ -45,20 +45,22 @@ type BaseIPTables interface {
 	DeleteChain(table, chain string) error
 	// NewChain creates a new chain
 	NewChain(table, chain string) error
+	// ListRules lists the rules in the table/chain passed to it
+	ListRules(table, chain string) ([]string, error)
 }
 
 // BatchProvider uses iptables-restore to program ACLs
 type BatchProvider struct{}
 
 // NewGoIPTablesProviderV4 returns an IptablesProvider interface based on the go-iptables
 // external package.
-func NewGoIPTablesProviderV4(batchTables []string) (IptablesProvider, error) {
+func NewGoIPTablesProviderV4(batchTables []string, customChain string) (IptablesProvider, error) {
 	return &BatchProvider{}, nil
 }
 
 // NewGoIPTablesProviderV6 returns an IptablesProvider interface based on the go-iptables
 // external package.
-func NewGoIPTablesProviderV6(batchTables []string) (IptablesProvider, error) {
+func NewGoIPTablesProviderV6(batchTables []string, customChain string) (IptablesProvider, error) {
 	return &BatchProvider{}, nil
 }
 
@@ -338,3 +340,9 @@ func (b *BatchProvider) ResetRules(subs string) error {
 	// does nothing
 	return nil
 }
+
+// ListRules lists the rules in the table/chain passed to it
+func (b *BatchProvider) ListRules(table, chain string) ([]string, error) {
+	// Unimplemented on windows
+	return []string{}, nil
+}

controller/pkg/aclprovider/iptablesprovidermock.go

@@ -16,6 +16,7 @@ type iptablesProviderMockedMethods struct {
 	commitMock        func() error
 	retrieveTableMock func() map[string]map[string][]string
 	resetMock         func(subs string) error
+	listRulesMock     func(table, chain string) ([]string, error)
 }
 
 // TestIptablesProvider is a test implementation for IptablesProvider
@@ -30,6 +31,7 @@ type TestIptablesProvider interface {
 	MockNewChain(t *testing.T, impl func(table, chain string) error)
 	MockCommit(t *testing.T, impl func() error)
 	MockReset(t *testing.T, impl func(subs string) error)
+	MockListRules(t *testing.T, impl func(table, chain string) ([]string, error))
 }
 
 // A testIptablesProvider is an empty TransactionalManipulator that can be easily mocked.
@@ -44,6 +46,9 @@ func NewTestIptablesProvider() TestIptablesProvider {
 	return &testIptablesProvider{lock: &sync.Mutex{}, mocks: map[*testing.T]*iptablesProviderMockedMethods{}}
 }
 
+func (m *testIptablesProvider) MockListRules(t *testing.T, impl func(table, chain string) ([]string, error)) {
+	m.currentMocks(t).listRulesMock = impl
+}
 func (m *testIptablesProvider) MockAppend(t *testing.T, impl func(table, chain string, rulespec ...string) error) {
 
 	m.currentMocks(t).appendMock = impl
@@ -96,6 +101,13 @@ func (m *testIptablesProvider) Append(table, chain string, rulespec ...string) e
 	return nil
 }
 
+func (m *testIptablesProvider) ListRules(table, chain string) ([]string, error) {
+	if mock := m.currentMocks(m.currentTest); mock != nil && mock.listRulesMock != nil {
+		return mock.listRulesMock(table, chain)
+	}
+	return []string{}, nil
+}
+
 func (m *testIptablesProvider) Insert(table, chain string, pos int, rulespec ...string) error {
 	if mock := m.currentMocks(m.currentTest); mock != nil && mock.insertMock != nil {
 		return mock.insertMock(table, chain, pos, rulespec...)