conn_test.go 23 KB

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