From ddb6921dab567c5231310f29c7f65e402774fb00 Mon Sep 17 00:00:00 2001
From: Paolo Invernizzi <paolo.invernizzi@suse.com>
Date: Fri, 7 Jun 2024 20:52:39 +0200
Subject: [PATCH] Add hostToNetworkShort unit test

---
 Makefile                          |  2 +-
 cmd/bisturi/main.go               |  4 +--
 {socket => sockets}/raw_socket.go |  6 +++--
 sockets/raw_socket_test.go        | 41 +++++++++++++++++++++++++++++++
 4 files changed, 48 insertions(+), 5 deletions(-)
 rename {socket => sockets}/raw_socket.go (93%)
 create mode 100644 sockets/raw_socket_test.go

diff --git a/Makefile b/Makefile
index 3b874ad..fba36a2 100644
--- a/Makefile
+++ b/Makefile
@@ -9,4 +9,4 @@ run: build
 	@./bin/bisturi
 
 test:
-	@$$GO_EXECUTABLE_PATH test -v ./...
\ No newline at end of file
+	@$$GO_EXECUTABLE_PATH test -v -race ./...
\ No newline at end of file
diff --git a/cmd/bisturi/main.go b/cmd/bisturi/main.go
index ea5a116..4998b14 100644
--- a/cmd/bisturi/main.go
+++ b/cmd/bisturi/main.go
@@ -5,7 +5,7 @@ import (
 	"log"
 	"net"
 
-	"github.com/NamelessOne91/bisturi/socket"
+	"github.com/NamelessOne91/bisturi/sockets"
 )
 
 var iface = flag.String("i", "eth0", "The network interface to listen to")
@@ -21,7 +21,7 @@ func main() {
 	}
 
 	// SYS_SOCKET syscall
-	rs, err := socket.NewRawSocket(*protocol)
+	rs, err := sockets.NewRawSocket(*protocol)
 	if err != nil {
 		log.Fatalf("Failed to open raw socket: %v", err)
 	}
diff --git a/socket/raw_socket.go b/sockets/raw_socket.go
similarity index 93%
rename from socket/raw_socket.go
rename to sockets/raw_socket.go
index f4f6de1..609d8c3 100644
--- a/socket/raw_socket.go
+++ b/sockets/raw_socket.go
@@ -1,4 +1,4 @@
-package socket
+package sockets
 
 import (
 	"errors"
@@ -17,8 +17,10 @@ var protocolEthernetType = map[string]uint16{
 	"arp":  syscall.ETH_P_ARP,
 	"ip":   syscall.ETH_P_IP,
 	"ipv6": syscall.ETH_P_IPV6,
-	"udp":  syscall.ETH_P_IP, // UDP is part of IP, needs special handling if filtered specifically
+	"udp":  syscall.ETH_P_IP, // UDP and TCP are part of IP, need special handling if filtered specifically
 	"udp6": syscall.ETH_P_IPV6,
+	"tcp":  syscall.ETH_P_IP,
+	"tcp6": syscall.ETH_P_IPV6,
 }
 
 var errUnsupportedProtocol = errors.New("unsupported protocol")
diff --git a/sockets/raw_socket_test.go b/sockets/raw_socket_test.go
new file mode 100644
index 0000000..b2f66a2
--- /dev/null
+++ b/sockets/raw_socket_test.go
@@ -0,0 +1,41 @@
+package sockets
+
+import (
+	"testing"
+)
+
+func TestHostToNetworkShort(t *testing.T) {
+	tests := []struct {
+		name     string
+		host     uint16
+		expected uint16
+	}{
+		{
+			name:     "Little Endian to Big Endian 1",
+			host:     0xff00,
+			expected: 0x00ff,
+		},
+		{
+			name:     "Little Endian to Big Endian 2",
+			host:     0xf00f,
+			expected: 0x0ff0,
+		},
+		{
+			name:     "Little Endian to Big Endian 3",
+			host:     0x0ff0,
+			expected: 0xf00f,
+		},
+		{
+			name:     "Little Endian to Big Endian 4",
+			host:     0xfff0,
+			expected: 0xf0ff,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := hostToNetworkShort(tt.host); got != tt.expected {
+				t.Errorf("expected %016b to become %016b: got %016b", tt.host, tt.expected, got)
+			}
+		})
+	}
+}