Optimize netdb Services database

pkg/pipeline/transform/netdb/netdb.go

@@ -0,0 +1,154 @@
+/*
+ * Copyright (C) 2022 IBM, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * > Note: this code is a revised and enhanced version of the netdb.go file
+ * > from https://github.com/dominikh/go-netdb/ (MIT License)
+ */
+
+package netdb
+
+import (
+	"fmt"
+	"io"
+	"strconv"
+	"strings"
+
+	"github.com/sirupsen/logrus"
+)
+
+var slog = logrus.WithField("component", "netdb.ServiceNames")
+
+type numKey struct {
+	port           int
+	protocolNumber int
+}
+
+type nameKey struct {
+	port         int
+	protocolName string
+}
+
+type ServiceNames struct {
+	protoNums map[int]struct{}
+	// key: protocol name, value: protocol number
+	protoNames  map[string]int
+	byPort      map[int]string
+	byProtoNum  map[numKey]string
+	byProtoName map[nameKey]string
+}
+
+// LoadServicesDB receives readers to the /etc/protocols and /etc/services formatted content
+// and returns a database that allow querying service names from ports and protocol information
+func LoadServicesDB(protocols, services io.Reader) (*ServiceNames, error) {
+	log := slog.WithField("method", "LoadServicesDB")
+	db := ServiceNames{
+		protoNums:   map[int]struct{}{},
+		protoNames:  map[string]int{},
+		byPort:      map[int]string{},
+		byProtoNum:  map[numKey]string{},
+		byProtoName: map[nameKey]string{},
+	}
+	// Load protocols
+	protoData, err := io.ReadAll(protocols)
+	if err != nil {
+		return nil, fmt.Errorf("reading protocols data: %w", err)
+	}
+
+	// key: proto name, value: aliases
+	protoAliases := map[string][]string{}
+
+	for i, line := range strings.Split(string(protoData), "\n") {
+		line = strings.TrimSpace(line)
+		split := strings.SplitN(line, "#", 2)
+		fields := strings.Fields(split[0])
+		if len(fields) < 2 {
+			continue
+		}
+
+		num, err := strconv.ParseInt(fields[1], 10, 32)
+		if err != nil {
+			log.WithFields(logrus.Fields{
+				logrus.ErrorKey: err,
+				"lineNum":       i,
+				"line":          line,
+			}).Debug("wrong protocol number. Ignoring entry")
+			continue
+		}
+
+		db.protoNums[int(num)] = struct{}{}
+		db.protoNames[fields[0]] = int(num)
+		for _, alias := range fields[2:] {
+			db.protoNames[alias] = int(num)
+		}
+		protoAliases[fields[0]] = fields[2:]
+	}
+
+	// Load services
+	svcData, err := io.ReadAll(services)
+	if err != nil {
+		return nil, fmt.Errorf("reading services data: %w", err)
+	}
+
+	for i, line := range strings.Split(string(svcData), "\n") {
+		line = strings.TrimSpace(line)
+		split := strings.SplitN(line, "#", 2)
+		fields := strings.Fields(split[0])
+		if len(fields) < 2 {
+			continue
+		}
+
+		svcName := fields[0]
+		portproto := strings.SplitN(fields[1], "/", 2)
+		protoName := portproto[1]
+		port, err := strconv.ParseInt(portproto[0], 10, 32)
+		if err != nil {
+			log.WithFields(logrus.Fields{
+				logrus.ErrorKey: err,
+				"lineNum":       i,
+				"line":          line,
+			}).Debug("wrong service port number. Ignoring entry")
+			continue
+		}
+		db.byPort[int(port)] = svcName
+		if protoNum, ok := db.protoNames[protoName]; ok {
+			db.byProtoNum[numKey{port: int(port), protocolNumber: protoNum}] = svcName
+		}
+		db.byProtoName[nameKey{port: int(port), protocolName: protoName}] = svcName
+		for _, alias := range protoAliases[protoName] {
+			db.byProtoName[nameKey{port: int(port), protocolName: alias}] = svcName
+		}
+	}
+	return &db, nil
+}
+
+// ByPortAndProtocolName returns the service name given a port and a protocol name (or
+// its alias). If the protocol does not exist, returns the name of any service matching
+// the port number.
+func (db *ServiceNames) ByPortAndProtocolName(port int, nameOrAlias string) string {
+	if _, ok := db.protoNames[nameOrAlias]; ok {
+		return db.byProtoName[nameKey{port: port, protocolName: nameOrAlias}]
+	}
+	return db.byPort[port]
+}
+
+// ByPortAndProtocolNumber returns the service name given a port and a protocol number.
+// If the protocol does not exist, returns the name of any service matching
+// the port number.
+func (db *ServiceNames) ByPortAndProtocolNumber(port, protoNum int) string {
+	if _, ok := db.protoNums[protoNum]; ok {
+		return db.byProtoNum[numKey{port: port, protocolNumber: protoNum}]
+	}
+	return db.byPort[port]
+}

pkg/pipeline/transform/netdb/netdb_test.go

@@ -0,0 +1,82 @@
+package netdb
+
+import (
+	"os"
+	"path"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func testingServicesDB() (*ServiceNames, error) {
+	etcProtos, err := os.Open(path.Join("testdata", "etcProtocols.txt"))
+	if err != nil {
+		return nil, err
+	}
+	defer etcProtos.Close()
+	etcSvcs, err := os.Open(path.Join("testdata", "etcServices.txt"))
+	if err != nil {
+		return nil, err
+	}
+	defer etcSvcs.Close()
+
+	return LoadServicesDB(etcProtos, etcSvcs)
+}
+
+func TestServicesDB(t *testing.T) {
+	db, err := testingServicesDB()
+	require.NoError(t, err)
+
+	assert.Equal(t, "netbios-dgm", db.ByPortAndProtocolNumber(138, 6))
+	assert.Equal(t, "netbios-dgm", db.ByPortAndProtocolName(138, "tcp"))
+	// verify it also finds service name by protocol alias
+	assert.Equal(t, "netbios-dgm", db.ByPortAndProtocolName(138, "TCP"))
+
+	// verify multiple protocols can be associated to the same port
+	assert.Equal(t, "ms-sql-s", db.ByPortAndProtocolNumber(1433, 6))
+	assert.Equal(t, "ms-sql-s", db.ByPortAndProtocolName(1433, "tcp"))
+	assert.Equal(t, "ms-sql-s", db.ByPortAndProtocolName(1433, "TCP"))
+	assert.Equal(t, "ms-sql-s", db.ByPortAndProtocolNumber(1433, 17))
+	assert.Equal(t, "ms-sql-s", db.ByPortAndProtocolName(1433, "udp"))
+	assert.Equal(t, "ms-sql-s", db.ByPortAndProtocolName(1433, "UDP"))
+
+	// verify it does search only by port number, if the protocol does not exist
+	assert.Equal(t, "ms-sql-s", db.ByPortAndProtocolNumber(1433, 99999))
+	assert.Equal(t, "ms-sql-s", db.ByPortAndProtocolName(1433, "tralara"))
+
+	// verify it returns nothing if the protocol exist but it's not associated to that port
+	assert.Empty(t, db.ByPortAndProtocolNumber(1433, 18))
+	assert.Empty(t, db.ByPortAndProtocolName(1433, "mux"))
+	assert.Empty(t, db.ByPortAndProtocolName(1433, "MUX"))
+}
+
+func BenchmarkGetProtoByNumber(b *testing.B) {
+	b.StopTimer()
+	db, err := testingServicesDB()
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		db.ByPortAndProtocolNumber(80, 6)
+		db.ByPortAndProtocolNumber(443, 17)
+		db.ByPortAndProtocolNumber(3306, 17)
+		db.ByPortAndProtocolNumber(27017, 6)
+	}
+}
+
+func BenchmarkGetProtoByName(b *testing.B) {
+	b.StopTimer()
+	db, err := testingServicesDB()
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		db.ByPortAndProtocolName(80, "tcp")
+		db.ByPortAndProtocolName(443, "udp")
+		db.ByPortAndProtocolName(3306, "UDP")
+		db.ByPortAndProtocolName(27017, "TCP")
+	}
+}