github.com/QuangHoangHao/kafka-go@v0.4.36/dialer_test.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"crypto/tls"
     6  	"crypto/x509"
     7  	"errors"
     8  	"fmt"
     9  	"io"
    10  	"net"
    11  	"reflect"
    12  	"sort"
    13  	"testing"
    14  	"time"
    15  )
    16  
    17  func TestDialer(t *testing.T) {
    18  	tests := []struct {
    19  		scenario string
    20  		function func(*testing.T, context.Context, *Dialer)
    21  	}{
    22  		{
    23  			scenario: "looking up partitions returns the list of available partitions for a topic",
    24  			function: testDialerLookupPartitions,
    25  		},
    26  	}
    27  
    28  	for _, test := range tests {
    29  		testFunc := test.function
    30  		t.Run(test.scenario, func(t *testing.T) {
    31  			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
    32  			defer cancel()
    33  
    34  			testFunc(t, ctx, &Dialer{})
    35  		})
    36  	}
    37  }
    38  
    39  func testDialerLookupPartitions(t *testing.T, ctx context.Context, d *Dialer) {
    40  	client, topic, shutdown := newLocalClientAndTopic()
    41  	defer shutdown()
    42  
    43  	// Write a message to ensure the partition gets created.
    44  	w := &Writer{
    45  		Addr:      TCP("localhost:9092"),
    46  		Topic:     topic,
    47  		Transport: client.Transport,
    48  	}
    49  	w.WriteMessages(ctx, Message{})
    50  	w.Close()
    51  
    52  	partitions, err := d.LookupPartitions(ctx, "tcp", "localhost:9092", topic)
    53  	if err != nil {
    54  		t.Error(err)
    55  		return
    56  	}
    57  
    58  	sort.Slice(partitions, func(i int, j int) bool {
    59  		return partitions[i].ID < partitions[j].ID
    60  	})
    61  
    62  	want := []Partition{
    63  		{
    64  			Topic:    topic,
    65  			Leader:   Broker{Host: "localhost", Port: 9092, ID: 1},
    66  			Replicas: []Broker{{Host: "localhost", Port: 9092, ID: 1}},
    67  			Isr:      []Broker{{Host: "localhost", Port: 9092, ID: 1}},
    68  			ID:       0,
    69  		},
    70  	}
    71  	if !reflect.DeepEqual(partitions, want) {
    72  		t.Errorf("bad partitions:\ngot:  %+v\nwant: %+v", partitions, want)
    73  	}
    74  }
    75  
    76  func tlsConfig(t *testing.T) *tls.Config {
    77  	const (
    78  		certPEM = `-----BEGIN CERTIFICATE-----
    79  MIID2zCCAsOgAwIBAgIJAMSqbewCgw4xMA0GCSqGSIb3DQEBCwUAMGAxCzAJBgNV
    80  BAYTAlVTMRMwEQYDVQQIDApDYWxpZm9ybmlhMRYwFAYDVQQHDA1TYW4gRnJhbmNp
    81  c2NvMRAwDgYDVQQKDAdTZWdtZW50MRIwEAYDVQQDDAlsb2NhbGhvc3QwHhcNMTcx
    82  MjIzMTU1NzAxWhcNMjcxMjIxMTU1NzAxWjBgMQswCQYDVQQGEwJVUzETMBEGA1UE
    83  CAwKQ2FsaWZvcm5pYTEWMBQGA1UEBwwNU2FuIEZyYW5jaXNjbzEQMA4GA1UECgwH
    84  U2VnbWVudDESMBAGA1UEAwwJbG9jYWxob3N0MIIBIjANBgkqhkiG9w0BAQEFAAOC
    85  AQ8AMIIBCgKCAQEAtda9OWKYNtINe/BKAoB+/zLg2qbaTeHN7L722Ug7YoY6zMVB
    86  aQEHrUmshw/TOrT7GLN/6e6rFN74UuNg72C1tsflZvxqkGdrup3I3jxMh2ApAxLi
    87  zem/M6Eke2OAqt+SzRPqc5GXH/nrWVd3wqg48DZOAR0jVTY2e0fWy+Er/cPJI1lc
    88  L6ZMIRJikHTXkaiFj2Jct1iWvgizx5HZJBxXJn2Awix5nvc+zmXM0ZhoedbJRoBC
    89  dGkRXd3xv2F4lqgVHtP3Ydjc/wYoPiGudSAkhyl9tnkHjvIjA/LeRNshWHbCIaQX
    90  yemnXIcyyf+W+7EK0gXio7uiP+QSoM5v/oeVMQIDAQABo4GXMIGUMHoGA1UdIwRz
    91  MHGhZKRiMGAxCzAJBgNVBAYTAlVTMRMwEQYDVQQIDApDYWxpZm9ybmlhMRYwFAYD
    92  VQQHDA1TYW4gRnJhbmNpc2NvMRAwDgYDVQQKDAdTZWdtZW50MRIwEAYDVQQDDAls
    93  b2NhbGhvc3SCCQCBYUuEuypDMTAJBgNVHRMEAjAAMAsGA1UdDwQEAwIE8DANBgkq
    94  hkiG9w0BAQsFAAOCAQEATk6IlVsXtNp4C1yeegaM+jE8qgKJfNm1sV27zKx8HPiO
    95  F7LvTGYIG7zd+bf3pDSwRxfBhsLEwmN9TUN1d6Aa9zeu95qOnR76POfHILgttu2w
    96  IzegO8I7BycnLjU9o/l9gCpusnN95tIYQhfD08ygUpYTQRuI0cmZ/Dp3xb0S9f5N
    97  miYTuUoStYSA4RWbDWo+Is9YWPu7rwieziOZ96oguGz3mtqvkjxVAQH1xZr3bKHr
    98  HU9LpQh0i6oTK0UCqnDwlhJl1c7A3UooxFpc3NGxyjogzTfI/gnBKfPo7eeswwsV
    99  77rjIkhBW49L35KOo1uyblgK1vTT7VPtzJnuDq3ORg==
   100  -----END CERTIFICATE-----`
   101  
   102  		keyPEM = `-----BEGIN RSA PRIVATE KEY-----
   103  MIIEowIBAAKCAQEAtda9OWKYNtINe/BKAoB+/zLg2qbaTeHN7L722Ug7YoY6zMVB
   104  aQEHrUmshw/TOrT7GLN/6e6rFN74UuNg72C1tsflZvxqkGdrup3I3jxMh2ApAxLi
   105  zem/M6Eke2OAqt+SzRPqc5GXH/nrWVd3wqg48DZOAR0jVTY2e0fWy+Er/cPJI1lc
   106  L6ZMIRJikHTXkaiFj2Jct1iWvgizx5HZJBxXJn2Awix5nvc+zmXM0ZhoedbJRoBC
   107  dGkRXd3xv2F4lqgVHtP3Ydjc/wYoPiGudSAkhyl9tnkHjvIjA/LeRNshWHbCIaQX
   108  yemnXIcyyf+W+7EK0gXio7uiP+QSoM5v/oeVMQIDAQABAoIBAQCa6roHW8JGYipu
   109  vsau3v5TOOtsHN67n3arDf6MGwfM5oLN1ffmF6SMs8myv36781hBMRv3FwjWHSf+
   110  pgz9o6zsbd05Ii8/m3yiXq609zZT107ZeYuU1mG5AL5uCNWjvhn5cdA6aX0RFwC0
   111  +tnjEyJ/NCS8ujBR9n/wA8IxrEKoTGcxRb6qFPPKWYoBevu34td1Szf0kH8AKjtQ
   112  rdPK0Of/ZEiAUxNMLTBEOmC0ZabxJV/YGWcUU4DpmEDZSgQSr4yLT4BFUwF2VC8t
   113  8VXn5dBP3RMo4h7JlteulcKYsMQZXD6KvUwY2LaEpFM/b14r+TZTUQGhwS+Ha11m
   114  xa4eNwFhAoGBANshGlpR9cUUq8vNex0Wb63P9BTRTXwg1yEJVMSua+DlaaqaX/hS
   115  hOxl3K4y2V5OCK31C+SOAqqbrGtMXVym5c5pX8YyC11HupFJwdFLUEc74uF3CtWY
   116  GMMvEvItCK5ZvYvS5I2CQGcp1fhEMle/Uz+hFi1eeWepMqgHbVx5vkdtAoGBANRv
   117  XYQsTAGSkhcHB++/ASDskAew5EoHfwtJzSX0BZC6DCACF/U4dCKzBVndOrELOPXs
   118  2CZXCG4ptWzNgt6YTlMX9U7nLei5pPjoivIJsMudnc22DrDS7C94rCk++M3JeLOM
   119  KSN0ou9+1iEdE7rQdMgZMryaY71OBonCIDsWgJZVAoGAB+k0CFq5IrpSUXNDpJMw
   120  yPee+jlsMLUGzzyFAOzDHEVsASq9mDtybQ5oXymay1rJ2W3lVgUCd6JTITSKklO8
   121  LC2FtaQM4Ps78w7Unne3mDrDQByKGZf6HOHQL0oM7C51N10Pv0Qaix7piKL9pklT
   122  +hIYuN6WR3XGTGaoPhRvGCkCgYBqaQ5y8q1v7Dd5iXAUS50JHPZYo+b2niKpSOKW
   123  LFHNWSRRtDrD/u9Nolb/2K1ZmcGCjo0HR3lVlVbnlVoEnk49mTaru2lntfZJKFLR
   124  QsFofR9at+NL95uPe+bhEkYW7uCjL4Y72GT1ipdAJwyG+3xD7ztW9g8X+EmWH8N9
   125  VZw7sQKBgGxp820jbjWhG1O9RnYLwflcZzUlSkhWJDg9tKJXBjD+hFX98Okuf0gu
   126  DUpdbxbJHSi0xAjOjLVswNws4pVwzgtZVK8R7k8j3Z5TtYTJTSQLfgVowuyEdAaI
   127  C8OxVJ/At/IJGnWSIz8z+/YCUf7p4jd2LJgmZVVzXeDsOFcH62gu
   128  -----END RSA PRIVATE KEY-----`
   129  
   130  		caPEM = `-----BEGIN CERTIFICATE-----
   131  MIIDPDCCAiQCCQCBYUuEuypDMTANBgkqhkiG9w0BAQsFADBgMQswCQYDVQQGEwJV
   132  UzETMBEGA1UECAwKQ2FsaWZvcm5pYTEWMBQGA1UEBwwNU2FuIEZyYW5jaXNjbzEQ
   133  MA4GA1UECgwHU2VnbWVudDESMBAGA1UEAwwJbG9jYWxob3N0MB4XDTE3MTIyMzE1
   134  NTMxOVoXDTI3MTIyMTE1NTMxOVowYDELMAkGA1UEBhMCVVMxEzARBgNVBAgMCkNh
   135  bGlmb3JuaWExFjAUBgNVBAcMDVNhbiBGcmFuY2lzY28xEDAOBgNVBAoMB1NlZ21l
   136  bnQxEjAQBgNVBAMMCWxvY2FsaG9zdDCCASIwDQYJKoZIhvcNAQEBBQADggEPADCC
   137  AQoCggEBAJwB+Yp6MyUepgtaRDxVjpMI2RmlAaV1qApMWu60LWGKJs4KWoIoLl6p
   138  oSEqnWrpMmb38pyGP99X1+t3uZjiK9L8nFhuKZ581tsTKLxaSl+YVg7JbH5LVCS6
   139  opsfB5ON1gJxf1HA9YyMqKHkBFh8/hdOGR0T6Bll9TPO1NQB/UqMy/tKr3sA3KZm
   140  XVDbRKSuUAQWz5J9/hLPmVMU41F/uD7mvyDY+x8GymInZjUXG4e0oq2RJgU6SYZ8
   141  mkscM6qhKY3mL487w/kHVFtFlMkOhvI7LIh3zVvWwgGSAoAv9yai9BDZNFSk0cEb
   142  bb/IK7BQW9sNI3lcnGirdbnjV94X9/sCAwEAATANBgkqhkiG9w0BAQsFAAOCAQEA
   143  MJLeGdYO3dpsPx2R39Bw0qa5cUh42huPf8n7rp4a4Ca5jJjcAlCYV8HzqOzpiKYy
   144  ZNuHy8LnNVYYh5Qoh8EO45bplMV1wnHfi6hW6DY5j3SQdcxkoVsW5R7rBF7a7SDg
   145  6uChVRPHgsnALUUc7Wvvd3sAs/NKHzHu86mgD3EefkdqWAaCapzcqT9mo9KXkWJM
   146  DhSJS+/iIaroc8umDnbPfhhgnlMf0/D4q0TjiLSSqyLzVifxnv9yHz56TrhHG/QP
   147  E/8+FEGCHYKM4JLr5smGlzv72Kfx9E1CkG6TgFNIHjipVv1AtYDvaNMdPF2533+F
   148  wE3YmpC3Q0g9r44nEbz4Bw==
   149  -----END CERTIFICATE-----`
   150  	)
   151  
   152  	// Define TLS configuration
   153  	certificate, err := tls.X509KeyPair([]byte(certPEM), []byte(keyPEM))
   154  	if err != nil {
   155  		t.Error(err)
   156  		t.FailNow()
   157  	}
   158  
   159  	caCertPool := x509.NewCertPool()
   160  	if ok := caCertPool.AppendCertsFromPEM([]byte(caPEM)); !ok {
   161  		t.Error(err)
   162  		t.FailNow()
   163  	}
   164  
   165  	return &tls.Config{
   166  		Certificates:       []tls.Certificate{certificate},
   167  		RootCAs:            caCertPool,
   168  		InsecureSkipVerify: true,
   169  	}
   170  }
   171  
   172  func TestDialerTLS(t *testing.T) {
   173  	client, topic, shutdown := newLocalClientAndTopic()
   174  	defer shutdown()
   175  
   176  	// Write a message to ensure the partition gets created.
   177  	w := &Writer{
   178  		Addr:      TCP("localhost:9092"),
   179  		Topic:     topic,
   180  		Transport: client.Transport,
   181  	}
   182  	w.WriteMessages(context.Background(), Message{})
   183  	w.Close()
   184  
   185  	// Create an SSL proxy using the tls.Config that connects to the
   186  	// docker-composed kafka
   187  	config := tlsConfig(t)
   188  	l, err := tls.Listen("tcp", "127.0.0.1:", config)
   189  	if err != nil {
   190  		t.Error(err)
   191  		return
   192  	}
   193  	defer l.Close()
   194  
   195  	go func() {
   196  		for {
   197  			conn, err := l.Accept()
   198  			if err != nil {
   199  				return // intentionally ignored
   200  			}
   201  
   202  			go func(in net.Conn) {
   203  				out, err := net.Dial("tcp", "localhost:9092")
   204  				if err != nil {
   205  					t.Error(err)
   206  					return
   207  				}
   208  				defer out.Close()
   209  
   210  				go io.Copy(in, out)
   211  				io.Copy(out, in)
   212  			}(conn)
   213  		}
   214  	}()
   215  
   216  	// Use the tls.Config and connect to the SSL proxy
   217  	d := &Dialer{
   218  		TLS: config,
   219  	}
   220  	partitions, err := d.LookupPartitions(context.Background(), "tcp", l.Addr().String(), topic)
   221  	if err != nil {
   222  		t.Error(err)
   223  		return
   224  	}
   225  
   226  	// Verify returned partition data is what we expect
   227  	sort.Slice(partitions, func(i int, j int) bool {
   228  		return partitions[i].ID < partitions[j].ID
   229  	})
   230  
   231  	want := []Partition{
   232  		{
   233  			Topic:    topic,
   234  			Leader:   Broker{Host: "localhost", Port: 9092, ID: 1},
   235  			Replicas: []Broker{{Host: "localhost", Port: 9092, ID: 1}},
   236  			Isr:      []Broker{{Host: "localhost", Port: 9092, ID: 1}},
   237  			ID:       0,
   238  		},
   239  	}
   240  	if !reflect.DeepEqual(partitions, want) {
   241  		t.Errorf("bad partitions:\ngot:  %+v\nwant: %+v", partitions, want)
   242  	}
   243  }
   244  
   245  type MockConn struct {
   246  	net.Conn
   247  	done       chan struct{}
   248  	partitions []Partition
   249  }
   250  
   251  func (m *MockConn) Read(b []byte) (n int, err error) {
   252  	select {
   253  	case <-time.After(time.Minute):
   254  	case <-m.done:
   255  		return 0, context.Canceled
   256  	}
   257  
   258  	return 0, io.EOF
   259  }
   260  
   261  func (m *MockConn) Write(b []byte) (n int, err error) {
   262  	select {
   263  	case <-time.After(time.Minute):
   264  	case <-m.done:
   265  		return 0, context.Canceled
   266  	}
   267  
   268  	return 0, io.EOF
   269  }
   270  
   271  func (m *MockConn) Close() error {
   272  	select {
   273  	case <-m.done:
   274  	default:
   275  		close(m.done)
   276  	}
   277  	return nil
   278  }
   279  
   280  func (m *MockConn) ReadPartitions(topics ...string) (partitions []Partition, err error) {
   281  	return m.partitions, err
   282  }
   283  
   284  func TestDialerConnectTLSHonorsContext(t *testing.T) {
   285  	config := tlsConfig(t)
   286  	d := &Dialer{
   287  		TLS: config,
   288  	}
   289  
   290  	conn := &MockConn{
   291  		done: make(chan struct{}),
   292  	}
   293  
   294  	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*25)
   295  	defer cancel()
   296  
   297  	_, err := d.connectTLS(ctx, conn, d.TLS)
   298  	if !errors.Is(err, context.DeadlineExceeded) {
   299  		t.Errorf("expected err to be %v; got %v", context.DeadlineExceeded, err)
   300  		t.FailNow()
   301  	}
   302  }
   303  
   304  func TestDialerResolver(t *testing.T) {
   305  	ctx := context.TODO()
   306  
   307  	tests := []struct {
   308  		scenario string
   309  		address  string
   310  		resolver map[string][]string
   311  	}{
   312  		{
   313  			scenario: "resolve domain to ip",
   314  			address:  "example.com",
   315  			resolver: map[string][]string{
   316  				"example.com": {"127.0.0.1"},
   317  			},
   318  		},
   319  		{
   320  			scenario: "resolve domain to ip and port",
   321  			address:  "example.com",
   322  			resolver: map[string][]string{
   323  				"example.com": {"127.0.0.1:9092"},
   324  			},
   325  		},
   326  		{
   327  			scenario: "resolve domain with port to ip",
   328  			address:  "example.com:9092",
   329  			resolver: map[string][]string{
   330  				"example.com": {"127.0.0.1:9092"},
   331  			},
   332  		},
   333  		{
   334  			scenario: "resolve domain with port to ip with different port",
   335  			address:  "example.com:9092",
   336  			resolver: map[string][]string{
   337  				"example.com": {"127.0.0.1:80"},
   338  			},
   339  		},
   340  		{
   341  			scenario: "resolve domain with port to ip",
   342  			address:  "example.com:9092",
   343  			resolver: map[string][]string{
   344  				"example.com": {"127.0.0.1"},
   345  			},
   346  		},
   347  	}
   348  
   349  	for _, test := range tests {
   350  		t.Run(test.scenario, func(t *testing.T) {
   351  			topic := makeTopic()
   352  			createTopic(t, topic, 1)
   353  			defer deleteTopic(t, topic)
   354  
   355  			d := Dialer{
   356  				Resolver: &mockResolver{addrs: test.resolver},
   357  			}
   358  
   359  			// Write a message to ensure the partition gets created.
   360  			w := NewWriter(WriterConfig{
   361  				Brokers: []string{"localhost:9092"},
   362  				Topic:   topic,
   363  				Dialer:  &d,
   364  			})
   365  			w.WriteMessages(context.Background(), Message{})
   366  			w.Close()
   367  
   368  			partitions, err := d.LookupPartitions(ctx, "tcp", test.address, topic)
   369  			if err != nil {
   370  				t.Error(err)
   371  				return
   372  			}
   373  
   374  			sort.Slice(partitions, func(i int, j int) bool {
   375  				return partitions[i].ID < partitions[j].ID
   376  			})
   377  
   378  			want := []Partition{
   379  				{
   380  					Topic:    topic,
   381  					Leader:   Broker{Host: "localhost", Port: 9092, ID: 1},
   382  					Replicas: []Broker{{Host: "localhost", Port: 9092, ID: 1}},
   383  					Isr:      []Broker{{Host: "localhost", Port: 9092, ID: 1}},
   384  					ID:       0,
   385  				},
   386  			}
   387  			if !reflect.DeepEqual(partitions, want) {
   388  				t.Errorf("bad partitions:\ngot:  %+v\nwant: %+v", partitions, want)
   389  			}
   390  		})
   391  	}
   392  }
   393  
   394  type mockResolver struct {
   395  	addrs map[string][]string
   396  }
   397  
   398  func (mr *mockResolver) LookupHost(ctx context.Context, host string) ([]string, error) {
   399  	if addrs, ok := mr.addrs[host]; !ok {
   400  		return nil, fmt.Errorf("unrecognized host %s", host)
   401  	} else {
   402  		return addrs, nil
   403  	}
   404  }