quic_connection_test.go 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819
  1. package connection
  2. import (
  3. "bytes"
  4. "context"
  5. "crypto/rand"
  6. "crypto/rsa"
  7. "crypto/tls"
  8. "crypto/x509"
  9. "encoding/pem"
  10. "fmt"
  11. "io"
  12. "math/big"
  13. "net"
  14. "net/http"
  15. "net/netip"
  16. "net/url"
  17. "strings"
  18. "testing"
  19. "time"
  20. "github.com/gobwas/ws/wsutil"
  21. "github.com/google/uuid"
  22. "github.com/pkg/errors"
  23. "github.com/quic-go/quic-go"
  24. "github.com/rs/zerolog"
  25. "github.com/stretchr/testify/assert"
  26. "github.com/stretchr/testify/require"
  27. "golang.org/x/net/nettest"
  28. "github.com/cloudflare/cloudflared/datagramsession"
  29. "github.com/cloudflare/cloudflared/ingress"
  30. "github.com/cloudflare/cloudflared/packet"
  31. cfdquic "github.com/cloudflare/cloudflared/quic"
  32. "github.com/cloudflare/cloudflared/tracing"
  33. "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
  34. rpcquic "github.com/cloudflare/cloudflared/tunnelrpc/quic"
  35. )
  36. var (
  37. testTLSServerConfig = GenerateTLSConfig()
  38. testQUICConfig = &quic.Config{
  39. KeepAlivePeriod: 5 * time.Second,
  40. EnableDatagrams: true,
  41. }
  42. defaultQUICTimeout = 30 * time.Second
  43. )
  44. var _ ReadWriteAcker = (*streamReadWriteAcker)(nil)
  45. // TestQUICServer tests if a quic server accepts and responds to a quic client with the acceptance protocol.
  46. // It also serves as a demonstration for communication with the QUIC connection started by a cloudflared.
  47. func TestQUICServer(t *testing.T) {
  48. // This is simply a sample websocket frame message.
  49. wsBuf := &bytes.Buffer{}
  50. wsutil.WriteClientBinary(wsBuf, []byte("Hello"))
  51. var tests = []struct {
  52. desc string
  53. dest string
  54. connectionType pogs.ConnectionType
  55. metadata []pogs.Metadata
  56. message []byte
  57. expectedResponse []byte
  58. }{
  59. {
  60. desc: "test http proxy",
  61. dest: "/ok",
  62. connectionType: pogs.ConnectionTypeHTTP,
  63. metadata: []pogs.Metadata{
  64. {
  65. Key: "HttpHeader:Cf-Ray",
  66. Val: "123123123",
  67. },
  68. {
  69. Key: "HttpHost",
  70. Val: "cf.host",
  71. },
  72. {
  73. Key: "HttpMethod",
  74. Val: "GET",
  75. },
  76. },
  77. expectedResponse: []byte("OK"),
  78. },
  79. {
  80. desc: "test http body request streaming",
  81. dest: "/slow_echo_body",
  82. connectionType: pogs.ConnectionTypeHTTP,
  83. metadata: []pogs.Metadata{
  84. {
  85. Key: "HttpHeader:Cf-Ray",
  86. Val: "123123123",
  87. },
  88. {
  89. Key: "HttpHost",
  90. Val: "cf.host",
  91. },
  92. {
  93. Key: "HttpMethod",
  94. Val: "POST",
  95. },
  96. {
  97. Key: "HttpHeader:Content-Length",
  98. Val: "24",
  99. },
  100. },
  101. message: []byte("This is the message body"),
  102. expectedResponse: []byte("This is the message body"),
  103. },
  104. {
  105. desc: "test ws proxy",
  106. dest: "/ws/echo",
  107. connectionType: pogs.ConnectionTypeWebsocket,
  108. metadata: []pogs.Metadata{
  109. {
  110. Key: "HttpHeader:Cf-Cloudflared-Proxy-Connection-Upgrade",
  111. Val: "Websocket",
  112. },
  113. {
  114. Key: "HttpHeader:Another-Header",
  115. Val: "Misc",
  116. },
  117. {
  118. Key: "HttpHost",
  119. Val: "cf.host",
  120. },
  121. {
  122. Key: "HttpMethod",
  123. Val: "get",
  124. },
  125. },
  126. message: wsBuf.Bytes(),
  127. expectedResponse: []byte{0x82, 0x5, 0x48, 0x65, 0x6c, 0x6c, 0x6f},
  128. },
  129. {
  130. desc: "test tcp proxy",
  131. connectionType: pogs.ConnectionTypeTCP,
  132. metadata: []pogs.Metadata{},
  133. message: []byte("Here is some tcp data"),
  134. expectedResponse: []byte("Here is some tcp data"),
  135. },
  136. }
  137. for i, test := range tests {
  138. test := test // capture range variable
  139. t.Run(test.desc, func(t *testing.T) {
  140. ctx, cancel := context.WithCancel(context.Background())
  141. // Start a UDP Listener for QUIC.
  142. udpAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
  143. require.NoError(t, err)
  144. udpListener, err := net.ListenUDP(udpAddr.Network(), udpAddr)
  145. require.NoError(t, err)
  146. defer udpListener.Close()
  147. quicTransport := &quic.Transport{Conn: udpListener, ConnectionIDLength: 16}
  148. quicListener, err := quicTransport.Listen(testTLSServerConfig, testQUICConfig)
  149. require.NoError(t, err)
  150. serverDone := make(chan struct{})
  151. go func() {
  152. quicServer(
  153. ctx, t, quicListener, test.dest, test.connectionType, test.metadata, test.message, test.expectedResponse,
  154. )
  155. close(serverDone)
  156. }()
  157. tunnelConn, _ := testTunnelConnection(t, netip.MustParseAddrPort(udpListener.LocalAddr().String()), uint8(i))
  158. connDone := make(chan struct{})
  159. go func() {
  160. tunnelConn.Serve(ctx)
  161. close(connDone)
  162. }()
  163. <-serverDone
  164. cancel()
  165. <-connDone
  166. })
  167. }
  168. }
  169. type fakeControlStream struct {
  170. ControlStreamHandler
  171. }
  172. func (fakeControlStream) ServeControlStream(ctx context.Context, rw io.ReadWriteCloser, connOptions *pogs.ConnectionOptions, tunnelConfigGetter TunnelConfigJSONGetter) error {
  173. <-ctx.Done()
  174. return nil
  175. }
  176. func (fakeControlStream) IsStopped() bool {
  177. return true
  178. }
  179. func quicServer(
  180. ctx context.Context,
  181. t *testing.T,
  182. listener *quic.Listener,
  183. dest string,
  184. connectionType pogs.ConnectionType,
  185. metadata []pogs.Metadata,
  186. message []byte,
  187. expectedResponse []byte,
  188. ) {
  189. session, err := listener.Accept(ctx)
  190. require.NoError(t, err)
  191. quicStream, err := session.OpenStreamSync(context.Background())
  192. require.NoError(t, err)
  193. stream := cfdquic.NewSafeStreamCloser(quicStream, defaultQUICTimeout, &log)
  194. reqClientStream := rpcquic.RequestClientStream{ReadWriteCloser: stream}
  195. err = reqClientStream.WriteConnectRequestData(dest, connectionType, metadata...)
  196. require.NoError(t, err)
  197. _, err = reqClientStream.ReadConnectResponseData()
  198. require.NoError(t, err)
  199. if message != nil {
  200. // ALPN successful. Write data.
  201. _, err := stream.Write(message)
  202. require.NoError(t, err)
  203. }
  204. response := make([]byte, len(expectedResponse))
  205. _, err = stream.Read(response)
  206. if err != io.EOF {
  207. require.NoError(t, err)
  208. }
  209. // For now it is an echo server. Verify if the same data is returned.
  210. assert.Equal(t, expectedResponse, response)
  211. }
  212. type mockOriginProxyWithRequest struct{}
  213. func (moc *mockOriginProxyWithRequest) ProxyHTTP(w ResponseWriter, tr *tracing.TracedHTTPRequest, isWebsocket bool) error {
  214. // These are a series of crude tests to ensure the headers and http related data is transferred from
  215. // metadata.
  216. r := tr.Request
  217. if r.Method == "" {
  218. return errors.New("method not sent")
  219. }
  220. if r.Host == "" {
  221. return errors.New("host not sent")
  222. }
  223. if len(r.Header) == 0 {
  224. return errors.New("headers not set")
  225. }
  226. if isWebsocket {
  227. return wsEchoEndpoint(w, r)
  228. }
  229. switch r.URL.Path {
  230. case "/ok":
  231. originRespEndpoint(w, http.StatusOK, []byte(http.StatusText(http.StatusOK)))
  232. case "/slow_echo_body":
  233. time.Sleep(5)
  234. fallthrough
  235. case "/echo_body":
  236. resp := &http.Response{
  237. StatusCode: http.StatusOK,
  238. }
  239. _ = w.WriteRespHeaders(resp.StatusCode, resp.Header)
  240. io.Copy(w, r.Body)
  241. case "/error":
  242. return fmt.Errorf("Failed to proxy to origin")
  243. default:
  244. originRespEndpoint(w, http.StatusNotFound, []byte("page not found"))
  245. }
  246. return nil
  247. }
  248. func TestBuildHTTPRequest(t *testing.T) {
  249. var tests = []struct {
  250. name string
  251. connectRequest *pogs.ConnectRequest
  252. body io.ReadCloser
  253. req *http.Request
  254. }{
  255. {
  256. name: "check if http.Request is built correctly with content length",
  257. connectRequest: &pogs.ConnectRequest{
  258. Dest: "http://test.com",
  259. Metadata: []pogs.Metadata{
  260. {
  261. Key: "HttpHeader:Cf-Cloudflared-Proxy-Connection-Upgrade",
  262. Val: "Websocket",
  263. },
  264. {
  265. Key: "HttpHeader:Content-Length",
  266. Val: "514",
  267. },
  268. {
  269. Key: "HttpHeader:Another-Header",
  270. Val: "Misc",
  271. },
  272. {
  273. Key: "HttpHost",
  274. Val: "cf.host",
  275. },
  276. {
  277. Key: "HttpMethod",
  278. Val: "get",
  279. },
  280. },
  281. },
  282. req: &http.Request{
  283. Method: "get",
  284. URL: &url.URL{
  285. Scheme: "http",
  286. Host: "test.com",
  287. },
  288. Proto: "HTTP/1.1",
  289. ProtoMajor: 1,
  290. ProtoMinor: 1,
  291. Header: http.Header{
  292. "Another-Header": []string{"Misc"},
  293. "Content-Length": []string{"514"},
  294. },
  295. ContentLength: 514,
  296. Host: "cf.host",
  297. Body: io.NopCloser(&bytes.Buffer{}),
  298. },
  299. body: io.NopCloser(&bytes.Buffer{}),
  300. },
  301. {
  302. name: "if content length isn't part of request headers, then it's not set",
  303. connectRequest: &pogs.ConnectRequest{
  304. Dest: "http://test.com",
  305. Metadata: []pogs.Metadata{
  306. {
  307. Key: "HttpHeader:Cf-Cloudflared-Proxy-Connection-Upgrade",
  308. Val: "Websocket",
  309. },
  310. {
  311. Key: "HttpHeader:Another-Header",
  312. Val: "Misc",
  313. },
  314. {
  315. Key: "HttpHost",
  316. Val: "cf.host",
  317. },
  318. {
  319. Key: "HttpMethod",
  320. Val: "get",
  321. },
  322. },
  323. },
  324. req: &http.Request{
  325. Method: "get",
  326. URL: &url.URL{
  327. Scheme: "http",
  328. Host: "test.com",
  329. },
  330. Proto: "HTTP/1.1",
  331. ProtoMajor: 1,
  332. ProtoMinor: 1,
  333. Header: http.Header{
  334. "Another-Header": []string{"Misc"},
  335. },
  336. ContentLength: 0,
  337. Host: "cf.host",
  338. Body: http.NoBody,
  339. },
  340. body: io.NopCloser(&bytes.Buffer{}),
  341. },
  342. {
  343. name: "if content length is 0, but transfer-encoding is chunked, body is not nil",
  344. connectRequest: &pogs.ConnectRequest{
  345. Dest: "http://test.com",
  346. Metadata: []pogs.Metadata{
  347. {
  348. Key: "HttpHeader:Another-Header",
  349. Val: "Misc",
  350. },
  351. {
  352. Key: "HttpHeader:Transfer-Encoding",
  353. Val: "chunked",
  354. },
  355. {
  356. Key: "HttpHost",
  357. Val: "cf.host",
  358. },
  359. {
  360. Key: "HttpMethod",
  361. Val: "get",
  362. },
  363. },
  364. },
  365. req: &http.Request{
  366. Method: "get",
  367. URL: &url.URL{
  368. Scheme: "http",
  369. Host: "test.com",
  370. },
  371. Proto: "HTTP/1.1",
  372. ProtoMajor: 1,
  373. ProtoMinor: 1,
  374. Header: http.Header{
  375. "Another-Header": []string{"Misc"},
  376. "Transfer-Encoding": []string{"chunked"},
  377. },
  378. ContentLength: 0,
  379. Host: "cf.host",
  380. Body: io.NopCloser(&bytes.Buffer{}),
  381. },
  382. body: io.NopCloser(&bytes.Buffer{}),
  383. },
  384. {
  385. name: "if content length is 0, but transfer-encoding is gzip,chunked, body is not nil",
  386. connectRequest: &pogs.ConnectRequest{
  387. Dest: "http://test.com",
  388. Metadata: []pogs.Metadata{
  389. {
  390. Key: "HttpHeader:Another-Header",
  391. Val: "Misc",
  392. },
  393. {
  394. Key: "HttpHeader:Transfer-Encoding",
  395. Val: "gzip,chunked",
  396. },
  397. {
  398. Key: "HttpHost",
  399. Val: "cf.host",
  400. },
  401. {
  402. Key: "HttpMethod",
  403. Val: "get",
  404. },
  405. },
  406. },
  407. req: &http.Request{
  408. Method: "get",
  409. URL: &url.URL{
  410. Scheme: "http",
  411. Host: "test.com",
  412. },
  413. Proto: "HTTP/1.1",
  414. ProtoMajor: 1,
  415. ProtoMinor: 1,
  416. Header: http.Header{
  417. "Another-Header": []string{"Misc"},
  418. "Transfer-Encoding": []string{"gzip,chunked"},
  419. },
  420. ContentLength: 0,
  421. Host: "cf.host",
  422. Body: io.NopCloser(&bytes.Buffer{}),
  423. },
  424. body: io.NopCloser(&bytes.Buffer{}),
  425. },
  426. {
  427. name: "if content length is 0, and connect request is a websocket, body is not nil",
  428. connectRequest: &pogs.ConnectRequest{
  429. Type: pogs.ConnectionTypeWebsocket,
  430. Dest: "http://test.com",
  431. Metadata: []pogs.Metadata{
  432. {
  433. Key: "HttpHeader:Another-Header",
  434. Val: "Misc",
  435. },
  436. {
  437. Key: "HttpHost",
  438. Val: "cf.host",
  439. },
  440. {
  441. Key: "HttpMethod",
  442. Val: "get",
  443. },
  444. },
  445. },
  446. req: &http.Request{
  447. Method: "get",
  448. URL: &url.URL{
  449. Scheme: "http",
  450. Host: "test.com",
  451. },
  452. Proto: "HTTP/1.1",
  453. ProtoMajor: 1,
  454. ProtoMinor: 1,
  455. Header: http.Header{
  456. "Another-Header": []string{"Misc"},
  457. },
  458. ContentLength: 0,
  459. Host: "cf.host",
  460. Body: io.NopCloser(&bytes.Buffer{}),
  461. },
  462. body: io.NopCloser(&bytes.Buffer{}),
  463. },
  464. }
  465. log := zerolog.Nop()
  466. for _, test := range tests {
  467. test := test // capture range variable
  468. t.Run(test.name, func(t *testing.T) {
  469. req, err := buildHTTPRequest(context.Background(), test.connectRequest, test.body, 0, &log)
  470. assert.NoError(t, err)
  471. test.req = test.req.WithContext(req.Context())
  472. assert.Equal(t, test.req, req.Request)
  473. })
  474. }
  475. }
  476. func (moc *mockOriginProxyWithRequest) ProxyTCP(ctx context.Context, rwa ReadWriteAcker, tcpRequest *TCPRequest) error {
  477. rwa.AckConnection("")
  478. io.Copy(rwa, rwa)
  479. return nil
  480. }
  481. func TestServeUDPSession(t *testing.T) {
  482. // Start a UDP Listener for QUIC.
  483. udpAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
  484. require.NoError(t, err)
  485. udpListener, err := net.ListenUDP(udpAddr.Network(), udpAddr)
  486. require.NoError(t, err)
  487. defer udpListener.Close()
  488. ctx, cancel := context.WithCancel(context.Background())
  489. // Establish QUIC connection with edge
  490. edgeQUICSessionChan := make(chan quic.Connection)
  491. go func() {
  492. earlyListener, err := quic.Listen(udpListener, testTLSServerConfig, testQUICConfig)
  493. require.NoError(t, err)
  494. edgeQUICSession, err := earlyListener.Accept(ctx)
  495. require.NoError(t, err)
  496. edgeQUICSessionChan <- edgeQUICSession
  497. }()
  498. // Random index to avoid reusing port
  499. tunnelConn, datagramConn := testTunnelConnection(t, netip.MustParseAddrPort(udpListener.LocalAddr().String()), 28)
  500. go tunnelConn.Serve(ctx)
  501. edgeQUICSession := <-edgeQUICSessionChan
  502. serveSession(ctx, datagramConn, edgeQUICSession, closedByOrigin, io.EOF.Error(), t)
  503. serveSession(ctx, datagramConn, edgeQUICSession, closedByTimeout, datagramsession.SessionIdleErr(time.Millisecond*50).Error(), t)
  504. serveSession(ctx, datagramConn, edgeQUICSession, closedByRemote, "eyeball closed connection", t)
  505. cancel()
  506. }
  507. func TestNopCloserReadWriterCloseBeforeEOF(t *testing.T) {
  508. readerWriter := nopCloserReadWriter{ReadWriteCloser: &mockReaderNoopWriter{Reader: strings.NewReader("123456789")}}
  509. buffer := make([]byte, 5)
  510. n, err := readerWriter.Read(buffer)
  511. require.NoError(t, err)
  512. require.Equal(t, n, 5)
  513. // close
  514. require.NoError(t, readerWriter.Close())
  515. // read should get error
  516. n, err = readerWriter.Read(buffer)
  517. require.Equal(t, n, 0)
  518. require.Equal(t, err, fmt.Errorf("closed by handler"))
  519. }
  520. func TestNopCloserReadWriterCloseAfterEOF(t *testing.T) {
  521. readerWriter := nopCloserReadWriter{ReadWriteCloser: &mockReaderNoopWriter{Reader: strings.NewReader("123456789")}}
  522. buffer := make([]byte, 20)
  523. n, err := readerWriter.Read(buffer)
  524. require.NoError(t, err)
  525. require.Equal(t, n, 9)
  526. // force another read to read eof
  527. _, err = readerWriter.Read(buffer)
  528. require.Equal(t, err, io.EOF)
  529. // close
  530. require.NoError(t, readerWriter.Close())
  531. // read should get EOF still
  532. n, err = readerWriter.Read(buffer)
  533. require.Equal(t, n, 0)
  534. require.Equal(t, err, io.EOF)
  535. }
  536. func TestCreateUDPConnReuseSourcePort(t *testing.T) {
  537. edgeIPv4 := netip.MustParseAddrPort("0.0.0.0:0")
  538. edgeIPv6 := netip.MustParseAddrPort("[::]:0")
  539. // We assume the test environment has access to an IPv4 interface
  540. testCreateUDPConnReuseSourcePortForEdgeIP(t, edgeIPv4)
  541. if nettest.SupportsIPv6() {
  542. testCreateUDPConnReuseSourcePortForEdgeIP(t, edgeIPv6)
  543. }
  544. }
  545. func testCreateUDPConnReuseSourcePortForEdgeIP(t *testing.T, edgeIP netip.AddrPort) {
  546. logger := zerolog.Nop()
  547. conn, err := createUDPConnForConnIndex(0, nil, edgeIP, &logger)
  548. require.NoError(t, err)
  549. getPortFunc := func(conn *net.UDPConn) int {
  550. addr := conn.LocalAddr().(*net.UDPAddr)
  551. return addr.Port
  552. }
  553. initialPort := getPortFunc(conn)
  554. // close conn
  555. conn.Close()
  556. // should get the same port as before.
  557. conn, err = createUDPConnForConnIndex(0, nil, edgeIP, &logger)
  558. require.NoError(t, err)
  559. require.Equal(t, initialPort, getPortFunc(conn))
  560. // new index, should get a different port
  561. conn1, err := createUDPConnForConnIndex(1, nil, edgeIP, &logger)
  562. require.NoError(t, err)
  563. require.NotEqual(t, initialPort, getPortFunc(conn1))
  564. // not closing the conn and trying to obtain a new conn for same index should give a different random port
  565. conn, err = createUDPConnForConnIndex(0, nil, edgeIP, &logger)
  566. require.NoError(t, err)
  567. require.NotEqual(t, initialPort, getPortFunc(conn))
  568. }
  569. func serveSession(ctx context.Context, datagramConn *datagramV2Connection, edgeQUICSession quic.Connection, closeType closeReason, expectedReason string, t *testing.T) {
  570. var (
  571. payload = []byte(t.Name())
  572. )
  573. sessionID := uuid.New()
  574. cfdConn, originConn := net.Pipe()
  575. // Registers and run a new session
  576. session, err := datagramConn.sessionManager.RegisterSession(ctx, sessionID, cfdConn)
  577. require.NoError(t, err)
  578. sessionDone := make(chan struct{})
  579. go func() {
  580. datagramConn.serveUDPSession(session, time.Millisecond*50)
  581. close(sessionDone)
  582. }()
  583. // Send a message to the quic session on edge side, it should be deumx to this datagram v2 session
  584. muxedPayload, err := cfdquic.SuffixSessionID(sessionID, payload)
  585. require.NoError(t, err)
  586. muxedPayload, err = cfdquic.SuffixType(muxedPayload, cfdquic.DatagramTypeUDP)
  587. require.NoError(t, err)
  588. err = edgeQUICSession.SendDatagram(muxedPayload)
  589. require.NoError(t, err)
  590. readBuffer := make([]byte, len(payload)+1)
  591. n, err := originConn.Read(readBuffer)
  592. require.NoError(t, err)
  593. require.Equal(t, len(payload), n)
  594. require.True(t, bytes.Equal(payload, readBuffer[:n]))
  595. // Close connection to terminate session
  596. switch closeType {
  597. case closedByOrigin:
  598. originConn.Close()
  599. case closedByRemote:
  600. err = datagramConn.UnregisterUdpSession(ctx, sessionID, expectedReason)
  601. require.NoError(t, err)
  602. case closedByTimeout:
  603. }
  604. if closeType != closedByRemote {
  605. // Session was not closed by remote, so closeUDPSession should be invoked to unregister from remote
  606. unregisterFromEdgeChan := make(chan struct{})
  607. sessionRPCServer := &mockSessionRPCServer{
  608. sessionID: sessionID,
  609. unregisterReason: expectedReason,
  610. calledUnregisterChan: unregisterFromEdgeChan,
  611. }
  612. go runRPCServer(ctx, edgeQUICSession, sessionRPCServer, nil, t)
  613. <-unregisterFromEdgeChan
  614. }
  615. <-sessionDone
  616. }
  617. type closeReason uint8
  618. const (
  619. closedByOrigin closeReason = iota
  620. closedByRemote
  621. closedByTimeout
  622. )
  623. func runRPCServer(ctx context.Context, session quic.Connection, sessionRPCServer pogs.SessionManager, configRPCServer pogs.ConfigurationManager, t *testing.T) {
  624. stream, err := session.AcceptStream(ctx)
  625. require.NoError(t, err)
  626. if stream.StreamID() == 0 {
  627. // Skip the first stream, it's the control stream of the QUIC connection
  628. stream, err = session.AcceptStream(ctx)
  629. require.NoError(t, err)
  630. }
  631. ss := rpcquic.NewCloudflaredServer(
  632. func(_ context.Context, _ *rpcquic.RequestServerStream) error {
  633. return nil
  634. },
  635. sessionRPCServer,
  636. configRPCServer,
  637. 10*time.Second,
  638. )
  639. err = ss.Serve(ctx, stream)
  640. assert.NoError(t, err)
  641. }
  642. type mockSessionRPCServer struct {
  643. sessionID uuid.UUID
  644. unregisterReason string
  645. calledUnregisterChan chan struct{}
  646. }
  647. func (s mockSessionRPCServer) RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16, closeIdleAfter time.Duration, traceContext string) (*pogs.RegisterUdpSessionResponse, error) {
  648. return nil, fmt.Errorf("mockSessionRPCServer doesn't implement RegisterUdpSession")
  649. }
  650. func (s mockSessionRPCServer) UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID, reason string) error {
  651. if s.sessionID != sessionID {
  652. return fmt.Errorf("expect session ID %s, got %s", s.sessionID, sessionID)
  653. }
  654. if s.unregisterReason != reason {
  655. return fmt.Errorf("expect unregister reason %s, got %s", s.unregisterReason, reason)
  656. }
  657. close(s.calledUnregisterChan)
  658. return nil
  659. }
  660. func testTunnelConnection(t *testing.T, serverAddr netip.AddrPort, index uint8) (TunnelConnection, *datagramV2Connection) {
  661. tlsClientConfig := &tls.Config{
  662. InsecureSkipVerify: true,
  663. NextProtos: []string{"argotunnel"},
  664. }
  665. // Start a mock httpProxy
  666. log := zerolog.New(io.Discard)
  667. ctx, cancel := context.WithCancel(context.Background())
  668. defer cancel()
  669. // Dial the QUIC connection to the edge
  670. conn, err := DialQuic(
  671. ctx,
  672. testQUICConfig,
  673. tlsClientConfig,
  674. serverAddr,
  675. nil, // connect on a random port
  676. index,
  677. &log,
  678. )
  679. // Start a session manager for the connection
  680. sessionDemuxChan := make(chan *packet.Session, 4)
  681. datagramMuxer := cfdquic.NewDatagramMuxerV2(conn, &log, sessionDemuxChan)
  682. sessionManager := datagramsession.NewManager(&log, datagramMuxer.SendToSession, sessionDemuxChan)
  683. packetRouter := ingress.NewPacketRouter(nil, datagramMuxer, &log)
  684. datagramConn := &datagramV2Connection{
  685. conn,
  686. sessionManager,
  687. datagramMuxer,
  688. packetRouter,
  689. 15 * time.Second,
  690. 0 * time.Second,
  691. &log,
  692. }
  693. tunnelConn, err := NewTunnelConnection(
  694. ctx,
  695. conn,
  696. index,
  697. &mockOrchestrator{originProxy: &mockOriginProxyWithRequest{}},
  698. datagramConn,
  699. fakeControlStream{},
  700. &pogs.ConnectionOptions{},
  701. 15*time.Second,
  702. 0*time.Second,
  703. 0*time.Second,
  704. &log,
  705. )
  706. require.NoError(t, err)
  707. return tunnelConn, datagramConn
  708. }
  709. type mockReaderNoopWriter struct {
  710. io.Reader
  711. }
  712. func (m *mockReaderNoopWriter) Write(p []byte) (n int, err error) {
  713. return len(p), nil
  714. }
  715. func (m *mockReaderNoopWriter) Close() error {
  716. return nil
  717. }
  718. // GenerateTLSConfig sets up a bare-bones TLS config for a QUIC server
  719. func GenerateTLSConfig() *tls.Config {
  720. key, err := rsa.GenerateKey(rand.Reader, 1024)
  721. if err != nil {
  722. panic(err)
  723. }
  724. template := x509.Certificate{SerialNumber: big.NewInt(1)}
  725. certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key)
  726. if err != nil {
  727. panic(err)
  728. }
  729. keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
  730. certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
  731. tlsCert, err := tls.X509KeyPair(certPEM, keyPEM)
  732. if err != nil {
  733. panic(err)
  734. }
  735. return &tls.Config{
  736. Certificates: []tls.Certificate{tlsCert},
  737. NextProtos: []string{"argotunnel"},
  738. }
  739. }