github.com/hack0072008/kafka-go@v1.0.1/dialer_test.go (about)

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