From 3377bc6d767d859ccad28da26365e3d6fbe98f3e Mon Sep 17 00:00:00 2001 From: Laurence Date: Thu, 9 Apr 2026 13:54:13 +0100 Subject: [PATCH] optimize SNI ClientHello parsing + add regression coverage Move SNI extraction from TLS handshake to direct ClientHello parsing. Keeps forwarding behavior and hostname normalization compatibility. Adds parity + malformed-input tests and broader fixture benchmarks. Benchmark summary (linux/amd64, i7-12700H) | case | new (ns/op) | old (ns/op) | speedup | new allocs | old allocs | |-------------------------|-------------|-------------|---------|------------|------------| | short direct | 8193 | 133580 | 16.3x | 11 | 76 | | short fragmented | 8674 | 136008 | 15.7x | 12 | 77 | | typical direct | 7418 | 137054 | 18.5x | 11 | 76 | | typical fragmented | 8252 | 149462 | 18.1x | 12 | 77 | | mixed_case direct | 7934 | 142533 | 18.0x | 12 | 76 | | mixed_case fragmented | 7690 | 138264 | 18.0x | 13 | 77 | | long direct | 7624 | 136663 | 17.9x | 11 | 76 | | long fragmented | 7733 | 143585 | 18.6x | 12 | 77 | Overall: ~16x-19x faster, ~85% fewer allocs/op. --- proxy/proxy.go | 187 ++++++++++++++---- proxy/proxy_test.go | 468 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 614 insertions(+), 41 deletions(-) diff --git a/proxy/proxy.go b/proxy/proxy.go index 71cf4ed..65eccdc 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -1,9 +1,10 @@ package proxy import ( + "bufio" "bytes" "context" - "crypto/tls" + "encoding/binary" "encoding/json" "fmt" "hash/fnv" @@ -11,6 +12,7 @@ import ( "log" "net" "net/http" + "slices" "strconv" "strings" "sync" @@ -19,6 +21,7 @@ import ( "github.com/fosrl/gerbil/internal/metrics" "github.com/fosrl/gerbil/logger" "github.com/patrickmn/go-cache" + "golang.org/x/crypto/cryptobyte" ) // RouteRecord represents a routing configuration @@ -82,20 +85,6 @@ type activeTunnel struct { conns []net.Conn } -// readOnlyConn is a wrapper for io.Reader that implements net.Conn -type readOnlyConn struct { - reader io.Reader -} - -func (conn readOnlyConn) Read(p []byte) (int, error) { return conn.reader.Read(p) } -func (conn readOnlyConn) Write(p []byte) (int, error) { return 0, io.ErrClosedPipe } -func (conn readOnlyConn) Close() error { return nil } -func (conn readOnlyConn) LocalAddr() net.Addr { return nil } -func (conn readOnlyConn) RemoteAddr() net.Addr { return nil } -func (conn readOnlyConn) SetDeadline(t time.Time) error { return nil } -func (conn readOnlyConn) SetReadDeadline(t time.Time) error { return nil } -func (conn readOnlyConn) SetWriteDeadline(t time.Time) error { return nil } - // parseProxyProtocolHeader parses a PROXY protocol v1 header from the connection func (p *SNIProxy) parseProxyProtocolHeader(conn net.Conn) (*ProxyProtocolInfo, net.Conn, error) { // Check if the connection comes from a trusted upstream @@ -444,6 +433,12 @@ func (p *SNIProxy) Stop() error { return nil } +const ( + tlsRecordTypeHandshake = 22 + tlsHandshakeTypeClientHello = 1 + maxClientHelloBytes = 64 * 1024 +) + // acceptConnections handles incoming connections func (p *SNIProxy) acceptConnections() { for { @@ -463,44 +458,154 @@ func (p *SNIProxy) acceptConnections() { } } -// readClientHello reads and parses the TLS ClientHello message -func (p *SNIProxy) readClientHello(reader io.Reader) (*tls.ClientHelloInfo, error) { - var hello *tls.ClientHelloInfo - err := tls.Server(readOnlyConn{reader: reader}, &tls.Config{ - GetConfigForClient: func(argHello *tls.ClientHelloInfo) (*tls.Config, error) { - hello = new(tls.ClientHelloInfo) - *hello = *argHello - return nil, nil - }, - }).Handshake() - if hello == nil { - return nil, err +func parseClientHelloServerName(clientHello []byte) (string, error) { + s := cryptobyte.String(clientHello) + + var legacyVersion uint16 + if !s.ReadUint16(&legacyVersion) { + return "", fmt.Errorf("client hello missing legacy version") + } + + var random []byte + if !s.ReadBytes(&random, 32) { + return "", fmt.Errorf("client hello missing random") + } + + var sessionID cryptobyte.String + if !s.ReadUint8LengthPrefixed(&sessionID) { + return "", fmt.Errorf("client hello missing session id") + } + + var cipherSuites cryptobyte.String + if !s.ReadUint16LengthPrefixed(&cipherSuites) { + return "", fmt.Errorf("client hello missing cipher suites") + } + + var compressionMethods cryptobyte.String + if !s.ReadUint8LengthPrefixed(&compressionMethods) { + return "", fmt.Errorf("client hello missing compression methods") + } + + var extensions cryptobyte.String + if !s.ReadUint16LengthPrefixed(&extensions) { + return "", fmt.Errorf("client hello missing extensions") + } + if !s.Empty() { + return "", fmt.Errorf("client hello has trailing data") + } + + for !extensions.Empty() { + var extensionType uint16 + var extensionData cryptobyte.String + if !extensions.ReadUint16(&extensionType) || !extensions.ReadUint16LengthPrefixed(&extensionData) { + return "", fmt.Errorf("invalid client hello extension") + } + + if extensionType != 0 { + continue + } + + var serverNameList cryptobyte.String + if !extensionData.ReadUint16LengthPrefixed(&serverNameList) { + return "", fmt.Errorf("invalid server name extension") + } + + serverName := "" + for !serverNameList.Empty() { + var nameType uint8 + var hostName cryptobyte.String + if !serverNameList.ReadUint8(&nameType) || !serverNameList.ReadUint16LengthPrefixed(&hostName) { + return "", fmt.Errorf("invalid server name entry") + } + + // host_name (RFC 6066); keep parsing to allow additional entries. + if nameType == 0 && len(hostName) > 0 && serverName == "" { + serverName = strings.TrimSuffix(strings.ToLower(string(hostName)), ".") + } + } + + if !extensionData.Empty() { + return "", fmt.Errorf("server name extension payload has trailing data") + } + if serverName == "" { + return "", fmt.Errorf("no host_name entry found in server name extension") + } + return serverName, nil } - return hello, nil + + return "", fmt.Errorf("SNI extension not present in ClientHello") } // peekClientHello reads the ClientHello while preserving the data for forwarding -func (p *SNIProxy) peekClientHello(reader io.Reader) (*tls.ClientHelloInfo, io.Reader, error) { - peekedBytes := new(bytes.Buffer) - hello, err := p.readClientHello(io.TeeReader(reader, peekedBytes)) - if err != nil { - return nil, nil, err +func (p *SNIProxy) peekClientHello(reader io.Reader) (string, io.Reader, error) { + bufferedReader := bufio.NewReader(reader) + peekedBytes := make([]byte, 0, 4096) + handshakeBytes := make([]byte, 0, 4096) + expectedClientHelloLen := -1 + + for { + var recordHeader [5]byte + if _, err := io.ReadFull(bufferedReader, recordHeader[:]); err != nil { + return "", nil, err + } + peekedBytes = append(peekedBytes, recordHeader[:]...) + + if recordHeader[0] != tlsRecordTypeHandshake { + return "", nil, fmt.Errorf("unexpected TLS record type %d", recordHeader[0]) + } + + recordLen := int(binary.BigEndian.Uint16(recordHeader[3:5])) + if recordLen == 0 { + return "", nil, fmt.Errorf("empty TLS handshake record") + } + if len(handshakeBytes)+recordLen > maxClientHelloBytes { + return "", nil, fmt.Errorf("client hello handshake bytes exceed %d", maxClientHelloBytes) + } + + recordPayloadOffset := len(peekedBytes) + peekedBytes = slices.Grow(peekedBytes, recordLen) + peekedBytes = peekedBytes[:recordPayloadOffset+recordLen] + if _, err := io.ReadFull(bufferedReader, peekedBytes[recordPayloadOffset:]); err != nil { + return "", nil, err + } + + handshakeBytes = append(handshakeBytes, peekedBytes[recordPayloadOffset:]...) + if expectedClientHelloLen < 0 { + if len(handshakeBytes) < 4 { + continue + } + if handshakeBytes[0] != tlsHandshakeTypeClientHello { + return "", nil, fmt.Errorf("unexpected TLS handshake type %d", handshakeBytes[0]) + } + expectedClientHelloLen = int(handshakeBytes[1])<<16 | int(handshakeBytes[2])<<8 | int(handshakeBytes[3]) + if expectedClientHelloLen == 0 { + return "", nil, fmt.Errorf("empty TLS client hello") + } + if expectedClientHelloLen > maxClientHelloBytes-4 { + return "", nil, fmt.Errorf("client hello payload exceeds %d bytes", maxClientHelloBytes-4) + } + } + + if len(handshakeBytes)-4 < expectedClientHelloLen { + continue + } + + serverName, err := parseClientHelloServerName(handshakeBytes[4 : 4+expectedClientHelloLen]) + if err != nil { + return "", nil, err + } + + return serverName, io.MultiReader(bytes.NewReader(peekedBytes), bufferedReader), nil } - return hello, io.MultiReader(peekedBytes, reader), nil } // extractSNI extracts the SNI hostname from the TLS ClientHello func (p *SNIProxy) extractSNI(conn net.Conn) (string, io.Reader, error) { - clientHello, clientReader, err := p.peekClientHello(conn) + hostname, clientReader, err := p.peekClientHello(conn) if err != nil { return "", nil, fmt.Errorf("failed to peek ClientHello: %w", err) } - - if clientHello.ServerName == "" { - return "", clientReader, fmt.Errorf("no SNI hostname found in ClientHello") - } - - return clientHello.ServerName, clientReader, nil + return hostname, clientReader, nil } // handleConnection processes a single client connection diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go index 747c81d..8d5824a 100644 --- a/proxy/proxy_test.go +++ b/proxy/proxy_test.go @@ -1,8 +1,14 @@ package proxy import ( + "bytes" + "crypto/tls" + "encoding/binary" + "io" "net" + "strings" "testing" + "time" ) func TestBuildProxyProtocolHeader(t *testing.T) { @@ -117,3 +123,465 @@ func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) { t.Errorf("Expected header '%s', got '%s'", expected, header) } } + +func TestPeekClientHelloExtractsSNIAndPreservesData(t *testing.T) { + proxy, err := NewSNIProxy(8443, "", "", "127.0.0.1", 443, nil, true, nil) + if err != nil { + t.Fatalf("Failed to create SNI proxy: %v", err) + } + + clientHello := captureClientHelloRecord(t, "example.com") + stream := append(append([]byte(nil), clientHello...), []byte("remaining bytes")...) + + hostname, reader, err := proxy.peekClientHello(bytes.NewReader(stream)) + if err != nil { + t.Fatalf("peekClientHello returned error: %v", err) + } + if hostname != "example.com" { + t.Fatalf("Expected hostname example.com, got %q", hostname) + } + + forwarded, err := io.ReadAll(reader) + if err != nil { + t.Fatalf("Failed reading forwarded stream: %v", err) + } + if !bytes.Equal(forwarded, stream) { + t.Fatalf("Forwarded stream did not match original input") + } +} + +func TestPeekClientHelloHandlesFragmentedHandshakeRecords(t *testing.T) { + proxy, err := NewSNIProxy(8443, "", "", "127.0.0.1", 443, nil, true, nil) + if err != nil { + t.Fatalf("Failed to create SNI proxy: %v", err) + } + + clientHello := captureClientHelloRecord(t, "fragmented.example.com") + fragmented := fragmentHandshakeRecord(t, clientHello, 37) + + hostname, reader, err := proxy.peekClientHello(bytes.NewReader(fragmented)) + if err != nil { + t.Fatalf("peekClientHello returned error for fragmented handshake: %v", err) + } + if hostname != "fragmented.example.com" { + t.Fatalf("Expected fragmented.example.com, got %q", hostname) + } + + forwarded, err := io.ReadAll(reader) + if err != nil { + t.Fatalf("Failed reading fragmented forwarded stream: %v", err) + } + if !bytes.Equal(forwarded, fragmented) { + t.Fatalf("Forwarded fragmented stream did not match original input") + } +} + +func TestPeekClientHelloWithoutSNI(t *testing.T) { + proxy, err := NewSNIProxy(8443, "", "", "127.0.0.1", 443, nil, true, nil) + if err != nil { + t.Fatalf("Failed to create SNI proxy: %v", err) + } + + clientHello := captureClientHelloRecord(t, "") + _, _, err = proxy.peekClientHello(bytes.NewReader(clientHello)) + if err == nil { + t.Fatal("Expected missing SNI error, got nil") + } + if !strings.Contains(err.Error(), "SNI extension not present") { + t.Fatalf("Expected missing SNI error, got %v", err) + } +} + +func TestPeekClientHelloNormalizesHostname(t *testing.T) { + proxy, err := NewSNIProxy(8443, "", "", "127.0.0.1", 443, nil, true, nil) + if err != nil { + t.Fatalf("Failed to create SNI proxy: %v", err) + } + + clientHello := captureClientHelloRecord(t, "ExAmPle.Com.") + hostname, reader, err := proxy.peekClientHello(bytes.NewReader(clientHello)) + if err != nil { + t.Fatalf("peekClientHello returned error: %v", err) + } + if hostname != "example.com" { + t.Fatalf("Expected canonical hostname example.com, got %q", hostname) + } + if _, err := io.Copy(io.Discard, reader); err != nil { + t.Fatalf("Failed to drain forwarded reader: %v", err) + } +} + +func TestParseClientHelloServerNameRejectsTrailingData(t *testing.T) { + record := captureClientHelloRecord(t, "example.com") + clientHello := extractClientHelloPayload(t, record) + malformed := append(append([]byte(nil), clientHello...), 0x00) + + _, err := parseClientHelloServerName(malformed) + if err == nil { + t.Fatal("Expected parse error for trailing ClientHello data, got nil") + } + if !strings.Contains(err.Error(), "trailing data") { + t.Fatalf("Expected trailing data parse error, got %v", err) + } +} + +func TestParseClientHelloServerNameAllowsAdditionalNameEntries(t *testing.T) { + record := captureClientHelloRecord(t, "example.com") + clientHello := extractClientHelloPayload(t, record) + withExtraName := addUnknownServerNameEntryToClientHello(t, clientHello) + + got, err := parseClientHelloServerName(withExtraName) + if err != nil { + t.Fatalf("Expected parser to accept additional name entries, got error: %v", err) + } + if got != "example.com" { + t.Fatalf("Expected hostname example.com, got %q", got) + } +} + +func TestPeekClientHelloMatchesOldBehavior(t *testing.T) { + proxy, err := NewSNIProxy(8443, "", "", "127.0.0.1", 443, nil, true, nil) + if err != nil { + t.Fatalf("Failed to create SNI proxy: %v", err) + } + + testCases := []struct { + name string + serverName string + }{ + {name: "basic", serverName: "example.com"}, + {name: "mixed_case", serverName: "ExAmPle.Com"}, + {name: "trailing_dot", serverName: "example.com."}, + {name: "single_label", serverName: "localhost"}, + {name: "long_hostname", serverName: "very.long.subdomain.name.for.sni.benchmark.example.com"}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + input := captureClientHelloRecord(t, tc.serverName) + assertNewAndOldPeekMatch(t, proxy, input) + + fragmented := fragmentHandshakeRecord(t, input, 37) + assertNewAndOldPeekMatch(t, proxy, fragmented) + }) + } +} + +func assertNewAndOldPeekMatch(tb testing.TB, proxy *SNIProxy, input []byte) { + tb.Helper() + + newHost, newReader, newErr := proxy.peekClientHello(bytes.NewReader(input)) + oldHost, oldReader, oldErr := oldPeekClientHello(bytes.NewReader(input)) + + if (newErr != nil) != (oldErr != nil) { + tb.Fatalf("error mismatch: newErr=%v oldErr=%v", newErr, oldErr) + } + if newErr != nil && oldErr != nil { + return + } + + // New path normalizes hostnames; compare with same normalization to check behavioral parity. + oldHost = strings.TrimSuffix(strings.ToLower(oldHost), ".") + if newHost != oldHost { + tb.Fatalf("hostname mismatch: new=%q old=%q", newHost, oldHost) + } + + newForwarded, err := io.ReadAll(newReader) + if err != nil { + tb.Fatalf("failed to read new forwarded stream: %v", err) + } + oldForwarded, err := io.ReadAll(oldReader) + if err != nil { + tb.Fatalf("failed to read old forwarded stream: %v", err) + } + + if !bytes.Equal(newForwarded, oldForwarded) { + tb.Fatalf("forwarded stream mismatch between new and old parsers") + } +} + +func captureClientHelloRecord(tb testing.TB, serverName string) []byte { + tb.Helper() + + clientConn, serverConn := net.Pipe() + errCh := make(chan error, 1) + + go func() { + cfg := &tls.Config{InsecureSkipVerify: true} + if serverName != "" { + cfg.ServerName = serverName + } + errCh <- tls.Client(clientConn, cfg).Handshake() + }() + + if err := serverConn.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil { + tb.Fatalf("Failed to set read deadline: %v", err) + } + + var header [5]byte + if _, err := io.ReadFull(serverConn, header[:]); err != nil { + tb.Fatalf("Failed to read TLS record header: %v", err) + } + + recordLen := int(binary.BigEndian.Uint16(header[3:5])) + record := make([]byte, 5+recordLen) + copy(record[:5], header[:]) + if _, err := io.ReadFull(serverConn, record[5:]); err != nil { + tb.Fatalf("Failed to read TLS record payload: %v", err) + } + + _ = clientConn.Close() + _ = serverConn.Close() + <-errCh + + return record +} + +func extractClientHelloPayload(tb testing.TB, record []byte) []byte { + tb.Helper() + + if len(record) < 9 { + tb.Fatalf("TLS record too short for ClientHello: %d", len(record)) + } + if record[0] != tlsRecordTypeHandshake { + tb.Fatalf("Expected handshake record type %d, got %d", tlsRecordTypeHandshake, record[0]) + } + + recordLen := int(binary.BigEndian.Uint16(record[3:5])) + if len(record) != 5+recordLen { + tb.Fatalf("Invalid record length: header=%d bytes, total=%d bytes", recordLen, len(record)) + } + + payload := record[5:] + if payload[0] != tlsHandshakeTypeClientHello { + tb.Fatalf("Expected ClientHello handshake type %d, got %d", tlsHandshakeTypeClientHello, payload[0]) + } + + clientHelloLen := int(payload[1])<<16 | int(payload[2])<<8 | int(payload[3]) + if len(payload) < 4+clientHelloLen { + tb.Fatalf("ClientHello payload truncated: expected=%d actual=%d", clientHelloLen, len(payload)-4) + } + + return append([]byte(nil), payload[4:4+clientHelloLen]...) +} + +func fragmentHandshakeRecord(tb testing.TB, record []byte, firstPayloadLen int) []byte { + tb.Helper() + + if len(record) < 6 { + tb.Fatalf("TLS record too short: %d", len(record)) + } + + payload := record[5:] + if firstPayloadLen <= 0 || firstPayloadLen >= len(payload) { + tb.Fatalf("Invalid fragment length %d for payload length %d", firstPayloadLen, len(payload)) + } + + version := record[1:3] + first := make([]byte, 5+firstPayloadLen) + first[0] = tlsRecordTypeHandshake + copy(first[1:3], version) + binary.BigEndian.PutUint16(first[3:5], uint16(firstPayloadLen)) + copy(first[5:], payload[:firstPayloadLen]) + + remainingLen := len(payload) - firstPayloadLen + second := make([]byte, 5+remainingLen) + second[0] = tlsRecordTypeHandshake + copy(second[1:3], version) + binary.BigEndian.PutUint16(second[3:5], uint16(remainingLen)) + copy(second[5:], payload[firstPayloadLen:]) + + return append(first, second...) +} + +func addUnknownServerNameEntryToClientHello(tb testing.TB, clientHello []byte) []byte { + tb.Helper() + + out := append([]byte(nil), clientHello...) + if len(out) < 39 { + tb.Fatalf("ClientHello too short: %d", len(out)) + } + + offset := 0 + offset += 2 // legacy version + offset += 32 // random + + sessionIDLen := int(out[offset]) + offset++ + offset += sessionIDLen + if offset+2 > len(out) { + tb.Fatal("Invalid ClientHello while reading cipher suites length") + } + + cipherSuitesLen := int(binary.BigEndian.Uint16(out[offset : offset+2])) + offset += 2 + cipherSuitesLen + if offset >= len(out) { + tb.Fatal("Invalid ClientHello while reading compression methods length") + } + + compressionMethodsLen := int(out[offset]) + offset++ + offset += compressionMethodsLen + if offset+2 > len(out) { + tb.Fatal("Invalid ClientHello while reading extensions length") + } + + extensionsLenOffset := offset + extensionsLen := int(binary.BigEndian.Uint16(out[offset : offset+2])) + offset += 2 + extensionsEnd := offset + extensionsLen + if extensionsEnd > len(out) { + tb.Fatal("Invalid ClientHello: extensions overflow payload") + } + + for offset+4 <= extensionsEnd { + extType := binary.BigEndian.Uint16(out[offset : offset+2]) + extDataLen := int(binary.BigEndian.Uint16(out[offset+2 : offset+4])) + extDataOffset := offset + 4 + extDataEnd := extDataOffset + extDataLen + if extDataEnd > extensionsEnd { + tb.Fatal("Invalid extension length in ClientHello") + } + + if extType == 0 { + if extDataLen < 2 { + tb.Fatal("Invalid server_name extension") + } + serverNameListLen := int(binary.BigEndian.Uint16(out[extDataOffset : extDataOffset+2])) + serverNameListOffset := extDataOffset + 2 + serverNameListEnd := serverNameListOffset + serverNameListLen + if serverNameListEnd > extDataEnd { + tb.Fatal("Invalid server_name list length") + } + + // Add a non-host_name entry (type 1) to verify parser tolerates additional names. + extraEntry := []byte{1, 0, 1, 0} + out = append(out[:serverNameListEnd], append(extraEntry, out[serverNameListEnd:]...)...) + + binary.BigEndian.PutUint16(out[extDataOffset:extDataOffset+2], uint16(serverNameListLen+len(extraEntry))) + binary.BigEndian.PutUint16(out[offset+2:offset+4], uint16(extDataLen+len(extraEntry))) + binary.BigEndian.PutUint16(out[extensionsLenOffset:extensionsLenOffset+2], uint16(extensionsLen+len(extraEntry))) + return out + } + + offset = extDataEnd + } + + tb.Fatal("server_name extension not found") + return nil +} + +func BenchmarkSNIExtraction(b *testing.B) { + proxy, err := NewSNIProxy(8443, "", "", "127.0.0.1", 443, nil, true, nil) + if err != nil { + b.Fatalf("Failed to create SNI proxy: %v", err) + } + + fixtures := []struct { + name string + serverName string + }{ + {name: "short", serverName: "a.io"}, + {name: "typical", serverName: "bench.example.com"}, + {name: "mixed_case", serverName: "Bench.Example.Com"}, + {name: "long", serverName: "very.long.subdomain.name.for.sni.benchmark.example.com"}, + } + + for _, fixture := range fixtures { + clientHello := captureClientHelloRecord(b, fixture.serverName) + fragmented := fragmentHandshakeRecord(b, clientHello, 37) + + benchmarks := []struct { + name string + input []byte + fn func(*testing.B, *SNIProxy, []byte) + }{ + {name: "new/direct_parser", input: clientHello, fn: benchmarkNewPeekClientHello}, + {name: "old/tls_handshake", input: clientHello, fn: benchmarkOldPeekClientHello}, + {name: "new/direct_parser_fragmented", input: fragmented, fn: benchmarkNewPeekClientHello}, + {name: "old/tls_handshake_fragmented", input: fragmented, fn: benchmarkOldPeekClientHello}, + } + + for _, bm := range benchmarks { + b.Run(fixture.name+"/"+bm.name, func(b *testing.B) { + b.ReportAllocs() + b.SetBytes(int64(len(bm.input))) + bm.fn(b, proxy, bm.input) + }) + } + } +} + +func benchmarkNewPeekClientHello(b *testing.B, proxy *SNIProxy, input []byte) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + hostname, reader, err := proxy.peekClientHello(bytes.NewReader(input)) + if err != nil { + b.Fatalf("peekClientHello returned error: %v", err) + } + if hostname == "" { + b.Fatal("peekClientHello returned empty hostname") + } + if _, err := io.Copy(io.Discard, reader); err != nil { + b.Fatalf("Failed to drain reader: %v", err) + } + } +} + +func benchmarkOldPeekClientHello(b *testing.B, _ *SNIProxy, input []byte) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + hostname, reader, err := oldPeekClientHello(bytes.NewReader(input)) + if err != nil { + b.Fatalf("oldPeekClientHello returned error: %v", err) + } + if hostname == "" { + b.Fatal("oldPeekClientHello returned empty hostname") + } + if _, err := io.Copy(io.Discard, reader); err != nil { + b.Fatalf("Failed to drain reader: %v", err) + } + } +} + +func oldPeekClientHello(reader io.Reader) (string, io.Reader, error) { + peekedBytes := new(bytes.Buffer) + hello, err := oldReadClientHello(io.TeeReader(reader, peekedBytes)) + if err != nil { + return "", nil, err + } + return hello.ServerName, io.MultiReader(peekedBytes, reader), nil +} + +func oldReadClientHello(reader io.Reader) (*tls.ClientHelloInfo, error) { + var hello *tls.ClientHelloInfo + err := tls.Server(benchmarkReadOnlyConn{reader: reader}, &tls.Config{ + GetConfigForClient: func(argHello *tls.ClientHelloInfo) (*tls.Config, error) { + hello = new(tls.ClientHelloInfo) + *hello = *argHello + return nil, nil + }, + }).Handshake() + if hello == nil { + return nil, err + } + return hello, nil +} + +type benchmarkReadOnlyConn struct { + reader io.Reader +} + +func (conn benchmarkReadOnlyConn) Read(p []byte) (int, error) { return conn.reader.Read(p) } +func (conn benchmarkReadOnlyConn) Write(p []byte) (int, error) { return 0, io.ErrClosedPipe } +func (conn benchmarkReadOnlyConn) Close() error { return nil } +func (conn benchmarkReadOnlyConn) LocalAddr() net.Addr { return benchmarkStaticAddr("local") } +func (conn benchmarkReadOnlyConn) RemoteAddr() net.Addr { return benchmarkStaticAddr("remote") } +func (conn benchmarkReadOnlyConn) SetDeadline(time.Time) error { return nil } +func (conn benchmarkReadOnlyConn) SetReadDeadline(time.Time) error { return nil } +func (conn benchmarkReadOnlyConn) SetWriteDeadline(time.Time) error { return nil } + +type benchmarkStaticAddr string + +func (a benchmarkStaticAddr) Network() string { return "tcp" } +func (a benchmarkStaticAddr) String() string { return string(a) }