conn_test.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929
  1. // Copyright (c) 2012 The gocql Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. // +build all unit
  5. package gocql
  6. import (
  7. "context"
  8. "crypto/tls"
  9. "crypto/x509"
  10. "fmt"
  11. "io"
  12. "io/ioutil"
  13. "net"
  14. "strings"
  15. "sync"
  16. "sync/atomic"
  17. "testing"
  18. "time"
  19. )
  20. const (
  21. defaultProto = protoVersion2
  22. )
  23. func TestApprove(t *testing.T) {
  24. tests := map[bool]bool{
  25. approve("org.apache.cassandra.auth.PasswordAuthenticator"): true,
  26. approve("com.instaclustr.cassandra.auth.SharedSecretAuthenticator"): true,
  27. approve("com.datastax.bdp.cassandra.auth.DseAuthenticator"): true,
  28. approve("com.apache.cassandra.auth.FakeAuthenticator"): false,
  29. }
  30. for k, v := range tests {
  31. if k != v {
  32. t.Fatalf("expected '%v', got '%v'", k, v)
  33. }
  34. }
  35. }
  36. func TestJoinHostPort(t *testing.T) {
  37. tests := map[string]string{
  38. "127.0.0.1:0": JoinHostPort("127.0.0.1", 0),
  39. "127.0.0.1:1": JoinHostPort("127.0.0.1:1", 9142),
  40. "[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:0": JoinHostPort("2001:0db8:85a3:0000:0000:8a2e:0370:7334", 0),
  41. "[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:1": JoinHostPort("[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:1", 9142),
  42. }
  43. for k, v := range tests {
  44. if k != v {
  45. t.Fatalf("expected '%v', got '%v'", k, v)
  46. }
  47. }
  48. }
  49. func testCluster(addr string, proto protoVersion) *ClusterConfig {
  50. cluster := NewCluster(addr)
  51. cluster.ProtoVersion = int(proto)
  52. cluster.disableControlConn = true
  53. return cluster
  54. }
  55. func TestSimple(t *testing.T) {
  56. srv := NewTestServer(t, defaultProto, context.Background())
  57. defer srv.Stop()
  58. cluster := testCluster(srv.Address, defaultProto)
  59. db, err := cluster.CreateSession()
  60. if err != nil {
  61. t.Fatalf("0x%x: NewCluster: %v", defaultProto, err)
  62. }
  63. if err := db.Query("void").Exec(); err != nil {
  64. t.Fatalf("0x%x: %v", defaultProto, err)
  65. }
  66. }
  67. func TestSSLSimple(t *testing.T) {
  68. srv := NewSSLTestServer(t, defaultProto, context.Background())
  69. defer srv.Stop()
  70. db, err := createTestSslCluster(srv.Address, defaultProto, true).CreateSession()
  71. if err != nil {
  72. t.Fatalf("0x%x: NewCluster: %v", defaultProto, err)
  73. }
  74. if err := db.Query("void").Exec(); err != nil {
  75. t.Fatalf("0x%x: %v", defaultProto, err)
  76. }
  77. }
  78. func TestSSLSimpleNoClientCert(t *testing.T) {
  79. srv := NewSSLTestServer(t, defaultProto, context.Background())
  80. defer srv.Stop()
  81. db, err := createTestSslCluster(srv.Address, defaultProto, false).CreateSession()
  82. if err != nil {
  83. t.Fatalf("0x%x: NewCluster: %v", defaultProto, err)
  84. }
  85. if err := db.Query("void").Exec(); err != nil {
  86. t.Fatalf("0x%x: %v", defaultProto, err)
  87. }
  88. }
  89. func createTestSslCluster(addr string, proto protoVersion, useClientCert bool) *ClusterConfig {
  90. cluster := testCluster(addr, proto)
  91. sslOpts := &SslOptions{
  92. CaPath: "testdata/pki/ca.crt",
  93. EnableHostVerification: false,
  94. }
  95. if useClientCert {
  96. sslOpts.CertPath = "testdata/pki/gocql.crt"
  97. sslOpts.KeyPath = "testdata/pki/gocql.key"
  98. }
  99. cluster.SslOpts = sslOpts
  100. return cluster
  101. }
  102. func TestClosed(t *testing.T) {
  103. t.Skip("Skipping the execution of TestClosed for now to try to concentrate on more important test failures on Travis")
  104. srv := NewTestServer(t, defaultProto, context.Background())
  105. defer srv.Stop()
  106. session, err := newTestSession(srv.Address, defaultProto)
  107. if err != nil {
  108. t.Fatalf("0x%x: NewCluster: %v", defaultProto, err)
  109. }
  110. session.Close()
  111. if err := session.Query("void").Exec(); err != ErrSessionClosed {
  112. t.Fatalf("0x%x: expected %#v, got %#v", defaultProto, ErrSessionClosed, err)
  113. }
  114. }
  115. func newTestSession(addr string, proto protoVersion) (*Session, error) {
  116. return testCluster(addr, proto).CreateSession()
  117. }
  118. func TestDNSLookupConnected(t *testing.T) {
  119. log := &testLogger{}
  120. Logger = log
  121. defer func() {
  122. Logger = &defaultLogger{}
  123. }()
  124. srv := NewTestServer(t, defaultProto, context.Background())
  125. defer srv.Stop()
  126. cluster := NewCluster("cassandra1.invalid", srv.Address, "cassandra2.invalid")
  127. cluster.ProtoVersion = int(defaultProto)
  128. cluster.disableControlConn = true
  129. // CreateSession() should attempt to resolve the DNS name "cassandraX.invalid"
  130. // and fail, but continue to connect via srv.Address
  131. _, err := cluster.CreateSession()
  132. if err != nil {
  133. t.Fatal("CreateSession() should have connected")
  134. }
  135. if !strings.Contains(log.String(), "gocql: dns error") {
  136. t.Fatalf("Expected to receive dns error log message - got '%s' instead", log.String())
  137. }
  138. }
  139. func TestDNSLookupError(t *testing.T) {
  140. log := &testLogger{}
  141. Logger = log
  142. defer func() {
  143. Logger = &defaultLogger{}
  144. }()
  145. srv := NewTestServer(t, defaultProto, context.Background())
  146. defer srv.Stop()
  147. cluster := NewCluster("cassandra1.invalid", "cassandra2.invalid")
  148. cluster.ProtoVersion = int(defaultProto)
  149. cluster.disableControlConn = true
  150. // CreateSession() should attempt to resolve each DNS name "cassandraX.invalid"
  151. // and fail since it could not resolve any dns entries
  152. _, err := cluster.CreateSession()
  153. if err == nil {
  154. t.Fatal("CreateSession() should have returned an error")
  155. }
  156. if !strings.Contains(log.String(), "gocql: dns error") {
  157. t.Fatalf("Expected to receive dns error log message - got '%s' instead", log.String())
  158. }
  159. if err.Error() != "gocql: unable to create session: failed to resolve any of the provided hostnames" {
  160. t.Fatalf("Expected CreateSession() to fail with message - got '%s' instead", err.Error())
  161. }
  162. }
  163. func TestStartupTimeout(t *testing.T) {
  164. ctx, cancel := context.WithCancel(context.Background())
  165. log := &testLogger{}
  166. Logger = log
  167. defer func() {
  168. Logger = &defaultLogger{}
  169. }()
  170. srv := NewTestServer(t, defaultProto, ctx)
  171. defer srv.Stop()
  172. // Tell the server to never respond to Startup frame
  173. atomic.StoreInt32(&srv.TimeoutOnStartup, 1)
  174. startTime := time.Now()
  175. cluster := NewCluster(srv.Address)
  176. cluster.ProtoVersion = int(defaultProto)
  177. cluster.disableControlConn = true
  178. // Set very long query connection timeout
  179. // so we know CreateSession() is using the ConnectTimeout
  180. cluster.Timeout = time.Second * 5
  181. // Create session should timeout during connect attempt
  182. _, err := cluster.CreateSession()
  183. if err == nil {
  184. t.Fatal("CreateSession() should have returned a timeout error")
  185. }
  186. elapsed := time.Since(startTime)
  187. if elapsed > time.Second*5 {
  188. t.Fatal("ConnectTimeout is not respected")
  189. }
  190. if !strings.Contains(err.Error(), "no connections were made when creating the session") {
  191. t.Fatalf("Expected to receive no connections error - got '%s'", err)
  192. }
  193. if !strings.Contains(log.String(), "no response to connection startup within timeout") {
  194. t.Fatalf("Expected to receive timeout log message - got '%s'", log.String())
  195. }
  196. cancel()
  197. }
  198. func TestTimeout(t *testing.T) {
  199. ctx, cancel := context.WithCancel(context.Background())
  200. srv := NewTestServer(t, defaultProto, ctx)
  201. defer srv.Stop()
  202. db, err := newTestSession(srv.Address, defaultProto)
  203. if err != nil {
  204. t.Fatalf("NewCluster: %v", err)
  205. }
  206. defer db.Close()
  207. var wg sync.WaitGroup
  208. wg.Add(1)
  209. go func() {
  210. defer wg.Done()
  211. select {
  212. case <-time.After(5 * time.Second):
  213. t.Errorf("no timeout")
  214. case <-ctx.Done():
  215. }
  216. }()
  217. if err := db.Query("kill").WithContext(ctx).Exec(); err == nil {
  218. t.Fatal("expected error got nil")
  219. }
  220. cancel()
  221. wg.Wait()
  222. }
  223. // TestQueryRetry will test to make sure that gocql will execute
  224. // the exact amount of retry queries designated by the user.
  225. func TestQueryRetry(t *testing.T) {
  226. ctx, cancel := context.WithCancel(context.Background())
  227. defer cancel()
  228. srv := NewTestServer(t, defaultProto, ctx)
  229. defer srv.Stop()
  230. db, err := newTestSession(srv.Address, defaultProto)
  231. if err != nil {
  232. t.Fatalf("NewCluster: %v", err)
  233. }
  234. defer db.Close()
  235. go func() {
  236. select {
  237. case <-ctx.Done():
  238. return
  239. case <-time.After(5 * time.Second):
  240. t.Errorf("no timeout")
  241. }
  242. }()
  243. rt := &SimpleRetryPolicy{NumRetries: 1}
  244. qry := db.Query("kill").RetryPolicy(rt)
  245. if err := qry.Exec(); err == nil {
  246. t.Fatalf("expected error")
  247. }
  248. requests := atomic.LoadInt64(&srv.nKillReq)
  249. attempts := qry.Attempts()
  250. if requests != int64(attempts) {
  251. t.Fatalf("expected requests %v to match query attempts %v", requests, attempts)
  252. }
  253. // the query will only be attempted once, but is being retried
  254. if requests != int64(rt.NumRetries) {
  255. t.Fatalf("failed to retry the query %v time(s). Query executed %v times", rt.NumRetries, requests-1)
  256. }
  257. }
  258. func TestStreams_Protocol1(t *testing.T) {
  259. srv := NewTestServer(t, protoVersion1, context.Background())
  260. defer srv.Stop()
  261. // TODO: these are more like session tests and should instead operate
  262. // on a single Conn
  263. cluster := testCluster(srv.Address, protoVersion1)
  264. cluster.NumConns = 1
  265. cluster.ProtoVersion = 1
  266. db, err := cluster.CreateSession()
  267. if err != nil {
  268. t.Fatal(err)
  269. }
  270. defer db.Close()
  271. var wg sync.WaitGroup
  272. for i := 1; i < 128; i++ {
  273. // here were just validating that if we send NumStream request we get
  274. // a response for every stream and the lengths for the queries are set
  275. // correctly.
  276. wg.Add(1)
  277. go func() {
  278. defer wg.Done()
  279. if err := db.Query("void").Exec(); err != nil {
  280. t.Error(err)
  281. }
  282. }()
  283. }
  284. wg.Wait()
  285. }
  286. func TestStreams_Protocol3(t *testing.T) {
  287. srv := NewTestServer(t, protoVersion3, context.Background())
  288. defer srv.Stop()
  289. // TODO: these are more like session tests and should instead operate
  290. // on a single Conn
  291. cluster := testCluster(srv.Address, protoVersion3)
  292. cluster.NumConns = 1
  293. cluster.ProtoVersion = 3
  294. db, err := cluster.CreateSession()
  295. if err != nil {
  296. t.Fatal(err)
  297. }
  298. defer db.Close()
  299. for i := 1; i < 32768; i++ {
  300. // the test server processes each conn synchronously
  301. // here were just validating that if we send NumStream request we get
  302. // a response for every stream and the lengths for the queries are set
  303. // correctly.
  304. if err = db.Query("void").Exec(); err != nil {
  305. t.Fatal(err)
  306. }
  307. }
  308. }
  309. func BenchmarkProtocolV3(b *testing.B) {
  310. srv := NewTestServer(b, protoVersion3, context.Background())
  311. defer srv.Stop()
  312. // TODO: these are more like session tests and should instead operate
  313. // on a single Conn
  314. cluster := NewCluster(srv.Address)
  315. cluster.NumConns = 1
  316. cluster.ProtoVersion = 3
  317. db, err := cluster.CreateSession()
  318. if err != nil {
  319. b.Fatal(err)
  320. }
  321. defer db.Close()
  322. b.ResetTimer()
  323. b.ReportAllocs()
  324. for i := 0; i < b.N; i++ {
  325. if err = db.Query("void").Exec(); err != nil {
  326. b.Fatal(err)
  327. }
  328. }
  329. }
  330. // This tests that the policy connection pool handles SSL correctly
  331. func TestPolicyConnPoolSSL(t *testing.T) {
  332. srv := NewSSLTestServer(t, defaultProto, context.Background())
  333. defer srv.Stop()
  334. cluster := createTestSslCluster(srv.Address, defaultProto, true)
  335. cluster.PoolConfig.HostSelectionPolicy = RoundRobinHostPolicy()
  336. db, err := cluster.CreateSession()
  337. if err != nil {
  338. t.Fatalf("failed to create new session: %v", err)
  339. }
  340. if err := db.Query("void").Exec(); err != nil {
  341. t.Fatalf("query failed due to error: %v", err)
  342. }
  343. db.Close()
  344. // wait for the pool to drain
  345. time.Sleep(100 * time.Millisecond)
  346. size := db.pool.Size()
  347. if size != 0 {
  348. t.Fatalf("connection pool did not drain, still contains %d connections", size)
  349. }
  350. }
  351. func TestQueryTimeout(t *testing.T) {
  352. srv := NewTestServer(t, defaultProto, context.Background())
  353. defer srv.Stop()
  354. cluster := testCluster(srv.Address, defaultProto)
  355. // Set the timeout arbitrarily low so that the query hits the timeout in a
  356. // timely manner.
  357. cluster.Timeout = 1 * time.Millisecond
  358. db, err := cluster.CreateSession()
  359. if err != nil {
  360. t.Fatalf("NewCluster: %v", err)
  361. }
  362. defer db.Close()
  363. ch := make(chan error, 1)
  364. go func() {
  365. err := db.Query("timeout").Exec()
  366. if err != nil {
  367. ch <- err
  368. return
  369. }
  370. t.Errorf("err was nil, expected to get a timeout after %v", db.cfg.Timeout)
  371. }()
  372. select {
  373. case err := <-ch:
  374. if err != ErrTimeoutNoResponse {
  375. t.Fatalf("expected to get %v for timeout got %v", ErrTimeoutNoResponse, err)
  376. }
  377. case <-time.After(10*time.Millisecond + db.cfg.Timeout):
  378. // ensure that the query goroutines have been scheduled
  379. t.Fatalf("query did not timeout after %v", db.cfg.Timeout)
  380. }
  381. }
  382. func BenchmarkSingleConn(b *testing.B) {
  383. srv := NewTestServer(b, 3, context.Background())
  384. defer srv.Stop()
  385. cluster := testCluster(srv.Address, 3)
  386. // Set the timeout arbitrarily low so that the query hits the timeout in a
  387. // timely manner.
  388. cluster.Timeout = 500 * time.Millisecond
  389. cluster.NumConns = 1
  390. db, err := cluster.CreateSession()
  391. if err != nil {
  392. b.Fatalf("NewCluster: %v", err)
  393. }
  394. defer db.Close()
  395. b.ResetTimer()
  396. b.RunParallel(func(pb *testing.PB) {
  397. for pb.Next() {
  398. err := db.Query("void").Exec()
  399. if err != nil {
  400. b.Error(err)
  401. return
  402. }
  403. }
  404. })
  405. }
  406. func TestQueryTimeoutReuseStream(t *testing.T) {
  407. t.Skip("no longer tests anything")
  408. // TODO(zariel): move this to conn test, we really just want to check what
  409. // happens when a conn is
  410. srv := NewTestServer(t, defaultProto, context.Background())
  411. defer srv.Stop()
  412. cluster := testCluster(srv.Address, defaultProto)
  413. // Set the timeout arbitrarily low so that the query hits the timeout in a
  414. // timely manner.
  415. cluster.Timeout = 1 * time.Millisecond
  416. cluster.NumConns = 1
  417. db, err := cluster.CreateSession()
  418. if err != nil {
  419. t.Fatalf("NewCluster: %v", err)
  420. }
  421. defer db.Close()
  422. db.Query("slow").Exec()
  423. err = db.Query("void").Exec()
  424. if err != nil {
  425. t.Fatal(err)
  426. }
  427. }
  428. func TestQueryTimeoutClose(t *testing.T) {
  429. srv := NewTestServer(t, defaultProto, context.Background())
  430. defer srv.Stop()
  431. cluster := testCluster(srv.Address, defaultProto)
  432. // Set the timeout arbitrarily low so that the query hits the timeout in a
  433. // timely manner.
  434. cluster.Timeout = 1000 * time.Millisecond
  435. cluster.NumConns = 1
  436. db, err := cluster.CreateSession()
  437. if err != nil {
  438. t.Fatalf("NewCluster: %v", err)
  439. }
  440. ch := make(chan error)
  441. go func() {
  442. err := db.Query("timeout").Exec()
  443. ch <- err
  444. }()
  445. // ensure that the above goroutine gets sheduled
  446. time.Sleep(50 * time.Millisecond)
  447. db.Close()
  448. select {
  449. case err = <-ch:
  450. case <-time.After(1 * time.Second):
  451. t.Fatal("timedout waiting to get a response once cluster is closed")
  452. }
  453. if err != ErrConnectionClosed {
  454. t.Fatalf("expected to get %v got %v", ErrConnectionClosed, err)
  455. }
  456. }
  457. func TestStream0(t *testing.T) {
  458. // TODO: replace this with type check
  459. const expErr = "gocql: received unexpected frame on stream 0"
  460. ctx, cancel := context.WithCancel(context.Background())
  461. defer cancel()
  462. srv := NewTestServer(t, defaultProto, ctx)
  463. defer srv.Stop()
  464. errorHandler := connErrorHandlerFn(func(conn *Conn, err error, closed bool) {
  465. if !srv.isClosed() && !strings.HasPrefix(err.Error(), expErr) {
  466. select {
  467. case <-ctx.Done():
  468. return
  469. default:
  470. t.Errorf("expected to get error prefix %q got %q", expErr, err.Error())
  471. }
  472. }
  473. })
  474. s, err := srv.session()
  475. if err != nil {
  476. t.Fatal(err)
  477. }
  478. defer s.Close()
  479. conn, err := s.connect(srv.host(), errorHandler)
  480. if err != nil {
  481. t.Fatal(err)
  482. }
  483. writer := frameWriterFunc(func(f *framer, streamID int) error {
  484. f.writeQueryFrame(0, "void", &queryParams{})
  485. return f.finishWrite()
  486. })
  487. // need to write out an invalid frame, which we need a connection to do
  488. framer, err := conn.exec(ctx, writer, nil)
  489. if err == nil {
  490. t.Fatal("expected to get an error on stream 0")
  491. } else if !strings.HasPrefix(err.Error(), expErr) {
  492. t.Fatalf("expected to get error prefix %q got %q", expErr, err.Error())
  493. } else if framer != nil {
  494. frame, err := framer.parseFrame()
  495. if err != nil {
  496. t.Fatal(err)
  497. }
  498. t.Fatalf("got frame %v", frame)
  499. }
  500. }
  501. func TestConnClosedBlocked(t *testing.T) {
  502. // issue 664
  503. const proto = 3
  504. srv := NewTestServer(t, proto, context.Background())
  505. defer srv.Stop()
  506. errorHandler := connErrorHandlerFn(func(conn *Conn, err error, closed bool) {
  507. t.Log(err)
  508. })
  509. s, err := srv.session()
  510. if err != nil {
  511. t.Fatal(err)
  512. }
  513. defer s.Close()
  514. conn, err := s.connect(srv.host(), errorHandler)
  515. if err != nil {
  516. t.Fatal(err)
  517. }
  518. if err := conn.conn.Close(); err != nil {
  519. t.Fatal(err)
  520. }
  521. // This will block indefintaly if #664 is not fixed
  522. err = conn.executeQuery(&Query{stmt: "void"}).Close()
  523. if !strings.HasSuffix(err.Error(), "use of closed network connection") {
  524. t.Fatalf("expected to get use of closed networking connection error got: %v\n", err)
  525. }
  526. }
  527. func TestContext_Timeout(t *testing.T) {
  528. srv := NewTestServer(t, defaultProto, context.Background())
  529. defer srv.Stop()
  530. cluster := testCluster(srv.Address, defaultProto)
  531. cluster.Timeout = 5 * time.Second
  532. db, err := cluster.CreateSession()
  533. if err != nil {
  534. t.Fatal(err)
  535. }
  536. defer db.Close()
  537. ctx, cancel := context.WithCancel(context.Background())
  538. cancel()
  539. err = db.Query("timeout").WithContext(ctx).Exec()
  540. if err != context.Canceled {
  541. t.Fatalf("expected to get context cancel error: %v got %v", context.Canceled, err)
  542. }
  543. }
  544. func NewTestServer(t testing.TB, protocol uint8, ctx context.Context) *TestServer {
  545. laddr, err := net.ResolveTCPAddr("tcp", "127.0.0.1:0")
  546. if err != nil {
  547. t.Fatal(err)
  548. }
  549. listen, err := net.ListenTCP("tcp", laddr)
  550. if err != nil {
  551. t.Fatal(err)
  552. }
  553. headerSize := 8
  554. if protocol > protoVersion2 {
  555. headerSize = 9
  556. }
  557. ctx, cancel := context.WithCancel(ctx)
  558. srv := &TestServer{
  559. Address: listen.Addr().String(),
  560. listen: listen,
  561. t: t,
  562. protocol: protocol,
  563. headerSize: headerSize,
  564. ctx: ctx,
  565. cancel: cancel,
  566. }
  567. go srv.closeWatch()
  568. go srv.serve()
  569. return srv
  570. }
  571. func NewSSLTestServer(t testing.TB, protocol uint8, ctx context.Context) *TestServer {
  572. pem, err := ioutil.ReadFile("testdata/pki/ca.crt")
  573. certPool := x509.NewCertPool()
  574. if !certPool.AppendCertsFromPEM(pem) {
  575. t.Fatalf("Failed parsing or appending certs")
  576. }
  577. mycert, err := tls.LoadX509KeyPair("testdata/pki/cassandra.crt", "testdata/pki/cassandra.key")
  578. if err != nil {
  579. t.Fatalf("could not load cert")
  580. }
  581. config := &tls.Config{
  582. Certificates: []tls.Certificate{mycert},
  583. RootCAs: certPool,
  584. }
  585. listen, err := tls.Listen("tcp", "127.0.0.1:0", config)
  586. if err != nil {
  587. t.Fatal(err)
  588. }
  589. headerSize := 8
  590. if protocol > protoVersion2 {
  591. headerSize = 9
  592. }
  593. ctx, cancel := context.WithCancel(ctx)
  594. srv := &TestServer{
  595. Address: listen.Addr().String(),
  596. listen: listen,
  597. t: t,
  598. protocol: protocol,
  599. headerSize: headerSize,
  600. ctx: ctx,
  601. cancel: cancel,
  602. }
  603. go srv.closeWatch()
  604. go srv.serve()
  605. return srv
  606. }
  607. type TestServer struct {
  608. Address string
  609. TimeoutOnStartup int32
  610. t testing.TB
  611. nreq uint64
  612. listen net.Listener
  613. nKillReq int64
  614. compressor Compressor
  615. protocol byte
  616. headerSize int
  617. ctx context.Context
  618. cancel context.CancelFunc
  619. quit chan struct{}
  620. mu sync.Mutex
  621. closed bool
  622. }
  623. func (srv *TestServer) session() (*Session, error) {
  624. return testCluster(srv.Address, protoVersion(srv.protocol)).CreateSession()
  625. }
  626. func (srv *TestServer) host() *HostInfo {
  627. host, err := hostInfo(srv.Address, 9042)
  628. if err != nil {
  629. srv.t.Fatal(err)
  630. }
  631. return host
  632. }
  633. func (srv *TestServer) closeWatch() {
  634. <-srv.ctx.Done()
  635. srv.mu.Lock()
  636. defer srv.mu.Unlock()
  637. srv.closeLocked()
  638. }
  639. func (srv *TestServer) serve() {
  640. defer srv.listen.Close()
  641. for !srv.isClosed() {
  642. conn, err := srv.listen.Accept()
  643. if err != nil {
  644. break
  645. }
  646. go func(conn net.Conn) {
  647. defer conn.Close()
  648. for !srv.isClosed() {
  649. framer, err := srv.readFrame(conn)
  650. if err != nil {
  651. if err == io.EOF {
  652. return
  653. }
  654. srv.errorLocked(err)
  655. return
  656. }
  657. atomic.AddUint64(&srv.nreq, 1)
  658. go srv.process(framer)
  659. }
  660. }(conn)
  661. }
  662. }
  663. func (srv *TestServer) isClosed() bool {
  664. srv.mu.Lock()
  665. defer srv.mu.Unlock()
  666. return srv.closed
  667. }
  668. func (srv *TestServer) closeLocked() {
  669. if srv.closed {
  670. return
  671. }
  672. srv.closed = true
  673. srv.listen.Close()
  674. srv.cancel()
  675. }
  676. func (srv *TestServer) Stop() {
  677. srv.mu.Lock()
  678. defer srv.mu.Unlock()
  679. srv.closeLocked()
  680. }
  681. func (srv *TestServer) errorLocked(err interface{}) {
  682. srv.mu.Lock()
  683. defer srv.mu.Unlock()
  684. if srv.closed {
  685. return
  686. }
  687. srv.t.Error(err)
  688. }
  689. func (srv *TestServer) process(f *framer) {
  690. head := f.header
  691. if head == nil {
  692. srv.errorLocked("process frame with a nil header")
  693. return
  694. }
  695. switch head.op {
  696. case opStartup:
  697. if atomic.LoadInt32(&srv.TimeoutOnStartup) > 0 {
  698. // Do not respond to startup command
  699. // wait until we get a cancel signal
  700. select {
  701. case <-srv.ctx.Done():
  702. return
  703. }
  704. }
  705. f.writeHeader(0, opReady, head.stream)
  706. case opOptions:
  707. f.writeHeader(0, opSupported, head.stream)
  708. f.writeShort(0)
  709. case opQuery:
  710. query := f.readLongString()
  711. first := query
  712. if n := strings.Index(query, " "); n > 0 {
  713. first = first[:n]
  714. }
  715. switch strings.ToLower(first) {
  716. case "kill":
  717. atomic.AddInt64(&srv.nKillReq, 1)
  718. f.writeHeader(0, opError, head.stream)
  719. f.writeInt(0x1001)
  720. f.writeString("query killed")
  721. case "use":
  722. f.writeInt(resultKindKeyspace)
  723. f.writeString(strings.TrimSpace(query[3:]))
  724. case "void":
  725. f.writeHeader(0, opResult, head.stream)
  726. f.writeInt(resultKindVoid)
  727. case "timeout":
  728. <-srv.ctx.Done()
  729. return
  730. case "slow":
  731. go func() {
  732. f.writeHeader(0, opResult, head.stream)
  733. f.writeInt(resultKindVoid)
  734. f.wbuf[0] = srv.protocol | 0x80
  735. select {
  736. case <-srv.ctx.Done():
  737. return
  738. case <-time.After(50 * time.Millisecond):
  739. f.finishWrite()
  740. }
  741. }()
  742. return
  743. default:
  744. f.writeHeader(0, opResult, head.stream)
  745. f.writeInt(resultKindVoid)
  746. }
  747. case opError:
  748. f.writeHeader(0, opError, head.stream)
  749. f.wbuf = append(f.wbuf, f.rbuf...)
  750. default:
  751. f.writeHeader(0, opError, head.stream)
  752. f.writeInt(0)
  753. f.writeString("not supported")
  754. }
  755. f.wbuf[0] = srv.protocol | 0x80
  756. if err := f.finishWrite(); err != nil {
  757. srv.errorLocked(err)
  758. }
  759. }
  760. func (srv *TestServer) readFrame(conn net.Conn) (*framer, error) {
  761. buf := make([]byte, srv.headerSize)
  762. head, err := readHeader(conn, buf)
  763. if err != nil {
  764. return nil, err
  765. }
  766. framer := newFramer(conn, conn, nil, srv.protocol)
  767. err = framer.readFrame(&head)
  768. if err != nil {
  769. return nil, err
  770. }
  771. // should be a request frame
  772. if head.version.response() {
  773. return nil, fmt.Errorf("expected to read a request frame got version: %v", head.version)
  774. } else if head.version.version() != srv.protocol {
  775. return nil, fmt.Errorf("expected to read protocol version 0x%x got 0x%x", srv.protocol, head.version.version())
  776. }
  777. return framer, nil
  778. }