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

internal/server/forward_test.go

@@ -0,0 +1,118 @@
+package server
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"net"
+	"testing"
+	"time"
+
+	"teleport/pkg/config"
+	"teleport/pkg/encryption"
+)
+
+func testServer(t *testing.T) *TeleportServer {
+	t.Helper()
+	return NewTeleportServer(&config.Config{
+		EncryptionKey: "a0e3dd20a761b118ca234160dd8b87230a001e332a97c9cfe3b8b9c99efaae03",
+	})
+}
+
+func TestForwardDataHidesHTTPPlaintextOnTunnel(t *testing.T) {
+	ts := testServer(t)
+	defer ts.Stop()
+
+	targetA, targetB := net.Pipe()
+	tunnelA, tunnelB := net.Pipe()
+	defer targetA.Close()
+	defer targetB.Close()
+	defer tunnelA.Close()
+	defer tunnelB.Close()
+
+	// target -> tunnel (encrypt)
+	go ts.forwardData(targetA, tunnelA, true)
+
+	resp := []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\nhello")
+	writeErr := make(chan error, 1)
+	go func() {
+		_, err := targetB.Write(resp)
+		writeErr <- err
+	}()
+
+	tunnelB.SetReadDeadline(time.Now().Add(5 * time.Second))
+	var length uint32
+	if err := binary.Read(tunnelB, binary.BigEndian, &length); err != nil {
+		t.Fatalf("failed to read frame length from tunnel: %v", err)
+	}
+	ciphertext := make([]byte, length)
+	if _, err := io.ReadFull(tunnelB, ciphertext); err != nil {
+		t.Fatalf("failed to read ciphertext from tunnel: %v", err)
+	}
+
+	if bytes.Contains(ciphertext, []byte("HTTP/1.1")) ||
+		bytes.Contains(ciphertext, []byte("Content-Type")) ||
+		bytes.Contains(ciphertext, []byte("hello")) {
+		t.Fatal("HTTP plaintext visible on server->client tunnel")
+	}
+
+	var frame bytes.Buffer
+	if err := binary.Write(&frame, binary.BigEndian, length); err != nil {
+		t.Fatal(err)
+	}
+	frame.Write(ciphertext)
+	plain, err := encryption.ReadEncryptedFrame(&frame, ts.derivedKey)
+	if err != nil {
+		t.Fatalf("decrypt of tunneled frame failed: %v", err)
+	}
+	if !bytes.Equal(plain, resp) {
+		t.Errorf("decrypted payload mismatch: got %q want %q", plain, resp)
+	}
+
+	select {
+	case err := <-writeErr:
+		if err != nil {
+			t.Fatalf("write to target conn failed: %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for target write")
+	}
+}
+
+func TestForwardDataDecryptsTunnelToTarget(t *testing.T) {
+	ts := testServer(t)
+	defer ts.Stop()
+
+	tunnelA, tunnelB := net.Pipe()
+	targetA, targetB := net.Pipe()
+	defer tunnelA.Close()
+	defer tunnelB.Close()
+	defer targetA.Close()
+	defer targetB.Close()
+
+	go ts.forwardData(tunnelA, targetA, false)
+
+	req := []byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")
+	writeErr := make(chan error, 1)
+	go func() {
+		writeErr <- encryption.WriteEncryptedFrame(tunnelB, req, ts.derivedKey)
+	}()
+
+	targetB.SetReadDeadline(time.Now().Add(5 * time.Second))
+	got := make([]byte, len(req))
+	if _, err := io.ReadFull(targetB, got); err != nil {
+		t.Fatalf("failed to read decrypted payload: %v", err)
+	}
+	if !bytes.Equal(got, req) {
+		t.Errorf("got %q want %q", got, req)
+	}
+
+	select {
+	case err := <-writeErr:
+		if err != nil {
+			t.Fatalf("WriteEncryptedFrame failed: %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for tunnel write")
+	}
+}