Encrypt TCP tunnel payload with AES-GCM (#6)

pkg/encryption/frame_test.go

@@ -0,0 +1,136 @@
+package encryption
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"strings"
+	"testing"
+)
+
+func testKey(t *testing.T) []byte {
+	t.Helper()
+	return DeriveKey("test-encryption-key-for-tcp-frames")
+}
+
+func TestEncryptedFrameRoundTrip(t *testing.T) {
+	key := testKey(t)
+	original := []byte("Hello, framed TCP tunnel!")
+
+	var buf bytes.Buffer
+	if err := WriteEncryptedFrame(&buf, original, key); err != nil {
+		t.Fatalf("WriteEncryptedFrame failed: %v", err)
+	}
+
+	got, err := ReadEncryptedFrame(&buf, key)
+	if err != nil {
+		t.Fatalf("ReadEncryptedFrame failed: %v", err)
+	}
+	if !bytes.Equal(got, original) {
+		t.Errorf("round trip mismatch: got %q want %q", got, original)
+	}
+}
+
+func TestEncryptedFrameMultipleChunks(t *testing.T) {
+	key := testKey(t)
+	chunks := [][]byte{
+		[]byte("chunk-one"),
+		[]byte("chunk-two-is-longer"),
+		[]byte{0x00, 0x01, 0xff},
+	}
+
+	var buf bytes.Buffer
+	for _, c := range chunks {
+		if err := WriteEncryptedFrame(&buf, c, key); err != nil {
+			t.Fatalf("WriteEncryptedFrame failed: %v", err)
+		}
+	}
+
+	for i, want := range chunks {
+		got, err := ReadEncryptedFrame(&buf, key)
+		if err != nil {
+			t.Fatalf("ReadEncryptedFrame chunk %d failed: %v", i, err)
+		}
+		if !bytes.Equal(got, want) {
+			t.Errorf("chunk %d mismatch: got %q want %q", i, got, want)
+		}
+	}
+}
+
+func TestEncryptedFrameRejectsZeroLength(t *testing.T) {
+	key := testKey(t)
+	var buf bytes.Buffer
+	if err := binary.Write(&buf, binary.BigEndian, uint32(0)); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := ReadEncryptedFrame(&buf, key); err == nil {
+		t.Fatal("expected error for zero-length frame")
+	}
+}
+
+func TestEncryptedFrameRejectsOversizeLength(t *testing.T) {
+	key := testKey(t)
+	var buf bytes.Buffer
+	if err := binary.Write(&buf, binary.BigEndian, uint32(MaxFrameSize+1)); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := ReadEncryptedFrame(&buf, key); err == nil {
+		t.Fatal("expected error for oversize frame length")
+	} else if !strings.Contains(err.Error(), "too large") {
+		t.Errorf("expected too-large error, got %v", err)
+	}
+}
+
+func TestEncryptedFrameWrongKey(t *testing.T) {
+	key1 := DeriveKey("tcp-frame-key-1")
+	key2 := DeriveKey("tcp-frame-key-2")
+
+	var buf bytes.Buffer
+	if err := WriteEncryptedFrame(&buf, []byte("secret"), key1); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := ReadEncryptedFrame(&buf, key2); err == nil {
+		t.Fatal("decrypting with the wrong key should fail")
+	}
+}
+
+func TestHTTPPlaintextNotVisibleOnFramedTunnel(t *testing.T) {
+	key := testKey(t)
+	req := []byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")
+
+	var wire bytes.Buffer
+	if err := WriteEncryptedFrame(&wire, req, key); err != nil {
+		t.Fatalf("WriteEncryptedFrame failed: %v", err)
+	}
+	sniffed := wire.Bytes()
+
+	if bytes.Contains(sniffed, []byte("GET /")) {
+		t.Fatal("sniffer saw HTTP request line on the tunnel")
+	}
+	if bytes.Contains(sniffed, []byte("Host:")) {
+		t.Fatal("sniffer saw HTTP Host header on the tunnel")
+	}
+	if bytes.Contains(sniffed, []byte("example.com")) {
+		t.Fatal("sniffer saw HTTP hostname on the tunnel")
+	}
+
+	// Length prefix is 4 bytes; ciphertext must be longer than plaintext.
+	if len(sniffed) < 4+len(req) {
+		t.Fatalf("framed ciphertext too short: %d", len(sniffed))
+	}
+
+	got, err := ReadEncryptedFrame(bytes.NewReader(sniffed), key)
+	if err != nil {
+		t.Fatalf("decrypt failed: %v", err)
+	}
+	if !bytes.Equal(got, req) {
+		t.Errorf("decrypted HTTP mismatch: got %q", got)
+	}
+}
+
+func TestReadEncryptedFrameEOF(t *testing.T) {
+	key := testKey(t)
+	if _, err := ReadEncryptedFrame(bytes.NewReader(nil), key); err != io.EOF {
+		t.Errorf("expected io.EOF, got %v", err)
+	}
+}