ja4sentinel/internal/fingerprint/engine_test.go

package fingerprint

import (
	"strings"
	"testing"

	"ja4sentinel/api"

	tlsfingerprint "github.com/psanford/tlsfingerprint"
)

func TestFromClientHello(t *testing.T) {
	tests := []struct {
		name    string
		ch      api.TLSClientHello
		wantErr bool
	}{
		{
			name: "empty payload",
			ch: api.TLSClientHello{
				Payload: []byte{},
			},
			wantErr: true,
		},
		{
			name: "invalid payload",
			ch: api.TLSClientHello{
				Payload: []byte{0x00, 0x01, 0x02},
			},
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			engine := NewEngine()
			_, err := engine.FromClientHello(tt.ch)
			if (err != nil) != tt.wantErr {
				t.Errorf("FromClientHello() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func TestNewEngine(t *testing.T) {
	engine := NewEngine()
	if engine == nil {
		t.Error("NewEngine() returned nil")
	}
}

func TestFromClientHello_ValidPayload(t *testing.T) {
	// Use a minimal valid TLS 1.2 ClientHello with extensions
	// Build a proper ClientHello using the same structure as parser tests
	clientHello := buildMinimalClientHelloForTest()

	ch := api.TLSClientHello{
		SrcIP:   "192.168.1.100",
		SrcPort: 54321,
		DstIP:   "10.0.0.1",
		DstPort: 443,
		Payload: clientHello,
	}

	engine := NewEngine()
	fp, err := engine.FromClientHello(ch)

	if err != nil {
		t.Fatalf("FromClientHello() error = %v", err)
	}
	if fp == nil {
		t.Fatal("FromClientHello() returned nil")
	}

	// Verify JA4 is populated (format: t13d... or t12d...)
	if fp.JA4 == "" {
		t.Error("JA4 should not be empty")
	}

	// JA4Hash is populated for internal use (but not serialized to LogRecord)
	// It contains the hash portions of the JA4 string
	if fp.JA4Hash == "" {
		t.Error("JA4Hash should be populated for internal use")
	}
}

// buildMinimalClientHelloForTest creates a minimal valid TLS 1.2 ClientHello
func buildMinimalClientHelloForTest() []byte {
	// Cipher suites (minimal set)
	cipherSuites := []byte{0x00, 0x04, 0x13, 0x01, 0x13, 0x02, 0xc0, 0x2f}
	// Compression methods (null only)
	compressionMethods := []byte{0x01, 0x00}
	// No extensions
	extensions := []byte{}
	extLen := len(extensions)

	// Build ClientHello handshake body
	handshakeBody := []byte{
		0x03, 0x03,             // Version: TLS 1.2
		// Random (32 bytes)
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00,                   // Session ID length: 0
	}

	// Add cipher suites (with length prefix)
	cipherSuiteLen := len(cipherSuites)
	handshakeBody = append(handshakeBody, byte(cipherSuiteLen>>8), byte(cipherSuiteLen))
	handshakeBody = append(handshakeBody, cipherSuites...)

	// Add compression methods (with length prefix)
	handshakeBody = append(handshakeBody, compressionMethods...)

	// Add extensions (with length prefix)
	handshakeBody = append(handshakeBody, byte(extLen>>8), byte(extLen))
	handshakeBody = append(handshakeBody, extensions...)

	// Now build full handshake with type and length
	handshakeLen := len(handshakeBody)
	handshake := append([]byte{
		0x01,                   // Handshake type: ClientHello
		byte(handshakeLen >> 16), byte(handshakeLen >> 8), byte(handshakeLen), // Handshake length
	}, handshakeBody...)

	// Build TLS record
	recordLen := len(handshake)
	record := make([]byte, 5+recordLen)
	record[0] = 0x16 // Handshake
	record[1] = 0x03 // Version: TLS 1.2
	record[2] = 0x03
	record[3] = byte(recordLen >> 8)
	record[4] = byte(recordLen)
	copy(record[5:], handshake)

	return record
}

// TestExtractJA4Hash tests the extractJA4Hash helper function
func TestExtractJA4Hash(t *testing.T) {
	tests := []struct {
		name   string
		ja4    string
		want   string
	}{
		{
			name:   "standard_ja4_format",
			ja4:    "t13d1516h2_8daaf6152771_02cb136f2775",
			want:   "8daaf6152771_02cb136f2775",
		},
		{
			name:   "ja4_with_single_underscore",
			ja4:    "t12d1234h1_abcdef123456",
			want:   "abcdef123456",
		},
		{
			name:   "ja4_no_underscore_returns_empty",
			ja4:    "t13d1516h2",
			want:   "",
		},
		{
			name:   "empty_ja4_returns_empty",
			ja4:    "",
			want:   "",
		},
		{
			name:   "underscore_at_start",
			ja4:    "_hash1_hash2",
			want:   "hash1_hash2",
		},
		{
			name:   "multiple_underscores_returns_after_first",
			ja4:    "base_part1_part2_part3",
			want:   "part1_part2_part3",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := extractJA4Hash(tt.ja4)
			if got != tt.want {
				t.Errorf("extractJA4Hash(%q) = %q, want %q", tt.ja4, got, tt.want)
			}
		})
	}
}

// TestFromClientHello_NilPayload tests error handling for nil payload
func TestFromClientHello_NilPayload(t *testing.T) {
	engine := NewEngine()
	ch := api.TLSClientHello{
		Payload: nil,
	}

	_, err := engine.FromClientHello(ch)

	if err == nil {
		t.Error("FromClientHello() with nil payload should return error")
	}
	if !strings.HasPrefix(err.Error(), "empty ClientHello payload") {
		t.Errorf("FromClientHello() error = %v, should start with 'empty ClientHello payload'", err)
	}
}

// TestFromClientHello_JA3Hash tests that JA3Hash is correctly populated
func TestFromClientHello_JA3Hash(t *testing.T) {
	clientHello := buildMinimalClientHelloForTest()

	ch := api.TLSClientHello{
		Payload: clientHello,
	}

	engine := NewEngine()
	fp, err := engine.FromClientHello(ch)

	if err != nil {
		t.Fatalf("FromClientHello() error = %v", err)
	}

	// JA3Hash should be populated (MD5 hash of JA3 string)
	if fp.JA3Hash == "" {
		t.Error("JA3Hash should be populated")
	}

	// JA3 should also be populated
	if fp.JA3 == "" {
		t.Error("JA3 should be populated")
	}
}

// TestFromClientHello_EmptyJA4Hash tests behavior when JA4 has no underscore
func TestFromClientHello_EmptyJA4Hash(t *testing.T) {
	// This test verifies that even if JA4 format changes, the code handles it gracefully
	engine := NewEngine()

	// Use a valid ClientHello - the library should produce a proper JA4
	clientHello := buildMinimalClientHelloForTest()

	ch := api.TLSClientHello{
		Payload: clientHello,
	}

	fp, err := engine.FromClientHello(ch)

	if err != nil {
		t.Fatalf("FromClientHello() error = %v", err)
	}

	// JA4 should always be populated
	if fp.JA4 == "" {
		t.Error("JA4 should be populated")
	}

	// JA4Hash may be empty if the JA4 format doesn't include underscores
	// This is acceptable behavior
}

// buildClientHelloWithTruncatedExtension creates a ClientHello where the last
// extension declares more data than actually present.
func buildClientHelloWithTruncatedExtension() []byte {
	// Build a valid SNI extension first
	sniHostname := []byte("example.com")
	sniExt := []byte{
		0x00, 0x00, // Extension type: server_name
	}
	sniData := []byte{0x00}
	sniListLen := 1 + 2 + len(sniHostname) // type(1) + len(2) + hostname
	sniData = append(sniData, byte(sniListLen>>8), byte(sniListLen))
	sniData = append(sniData, 0x00) // hostname type
	sniData = append(sniData, byte(len(sniHostname)>>8), byte(len(sniHostname)))
	sniData = append(sniData, sniHostname...)
	sniExt = append(sniExt, byte(len(sniData)>>8), byte(len(sniData)))
	sniExt = append(sniExt, sniData...)

	// Build a truncated extension: declares 100 bytes but only has 5
	truncatedExt := []byte{
		0x00, 0x15, // Extension type: padding
		0x00, 0x64, // Extension data length: 100 (but we only provide 5)
		0x00, 0x00, 0x00, 0x00, 0x00, // Only 5 bytes of padding
	}

	// Extensions = valid SNI + truncated padding
	extensions := append(sniExt, truncatedExt...)
	// But the extensions length field claims the full size (including the bad extension)
	extLen := len(extensions)

	// Cipher suites
	cipherSuites := []byte{0x00, 0x04, 0x13, 0x01, 0x13, 0x02, 0xc0, 0x2f}
	compressionMethods := []byte{0x01, 0x00}

	handshakeBody := []byte{0x03, 0x03}
	for i := 0; i < 32; i++ {
		handshakeBody = append(handshakeBody, 0x01)
	}
	handshakeBody = append(handshakeBody, 0x00) // session ID length: 0
	handshakeBody = append(handshakeBody, byte(len(cipherSuites)>>8), byte(len(cipherSuites)))
	handshakeBody = append(handshakeBody, cipherSuites...)
	handshakeBody = append(handshakeBody, compressionMethods...)
	handshakeBody = append(handshakeBody, byte(extLen>>8), byte(extLen))
	handshakeBody = append(handshakeBody, extensions...)

	handshakeLen := len(handshakeBody)
	handshake := append([]byte{
		0x01,
		byte(handshakeLen >> 16), byte(handshakeLen >> 8), byte(handshakeLen),
	}, handshakeBody...)

	recordLen := len(handshake)
	record := make([]byte, 5+recordLen)
	record[0] = 0x16
	record[1] = 0x03
	record[2] = 0x03
	record[3] = byte(recordLen >> 8)
	record[4] = byte(recordLen)
	copy(record[5:], handshake)

	return record
}

func TestFromClientHello_TruncatedExtension_StillGeneratesFingerprint(t *testing.T) {
	payload := buildClientHelloWithTruncatedExtension()

	ch := api.TLSClientHello{
		SrcIP:   "4.251.36.192",
		SrcPort: 19346,
		DstIP:   "212.95.72.88",
		DstPort: 443,
		Payload: payload,
		ConnID:  "4.251.36.192:19346->212.95.72.88:443",
	}

	engine := NewEngine()
	fp, err := engine.FromClientHello(ch)

	if err != nil {
		t.Fatalf("FromClientHello() should succeed after sanitization, got error: %v", err)
	}
	if fp == nil {
		t.Fatal("FromClientHello() returned nil fingerprint")
	}
	if fp.JA4 == "" {
		t.Error("JA4 should be populated even with truncated extension")
	}
	if fp.JA3 == "" {
		t.Error("JA3 should be populated even with truncated extension")
	}
}

func TestSanitizeClientHelloExtensions(t *testing.T) {
	t.Run("valid payload returns nil", func(t *testing.T) {
		valid := buildMinimalClientHelloForTest()
		result := sanitizeClientHelloExtensions(valid)
		if result != nil {
			t.Error("should return nil for valid payload (no fix needed)")
		}
	})

	t.Run("truncated extension is fixed", func(t *testing.T) {
		truncated := buildClientHelloWithTruncatedExtension()
		result := sanitizeClientHelloExtensions(truncated)
		if result == nil {
			t.Fatal("should return sanitized payload")
		}
		// The sanitized payload should be parseable by the library
		fp, err := tlsfingerprint.ParseClientHello(result)
		if err != nil {
			t.Fatalf("sanitized payload should parse without error, got: %v", err)
		}
		if fp == nil {
			t.Fatal("sanitized payload should produce a fingerprint")
		}
	})

	t.Run("too short returns nil", func(t *testing.T) {
		if sanitizeClientHelloExtensions([]byte{0x16}) != nil {
			t.Error("should return nil for short payload")
		}
	})

	t.Run("non-TLS returns nil", func(t *testing.T) {
		if sanitizeClientHelloExtensions([]byte{0x15, 0x03, 0x03, 0x00, 0x01, 0x00}) != nil {
			t.Error("should return nil for non-TLS payload")
		}
	})
}