GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

cloudflare/cloudflared

Public

mirrored from https://github.com/cloudflare/cloudflaredAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
173190aa799ece6eee51a0f979e6720f7cbf4bf9

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

cloudflared/datagramsession

datagramsession/manager_test.go

215lines · modecode

RawDownload
Latest commit unavailable.
1package datagramsession
2
3import (
4 "bytes"
5 "context"
6 "fmt"
7 "io"
8 "net"
9 "testing"
10 "time"
11
12 "github.com/google/uuid"
13 "github.com/rs/zerolog"
14 "github.com/stretchr/testify/require"
15 "golang.org/x/sync/errgroup"
16)
17
18func TestManagerServe(t *testing.T) {
19 const (
20 sessions = 20
21 msgs = 50
22 )
23 log := zerolog.Nop()
24 transport := &mockQUICTransport{
25 reqChan: newDatagramChannel(1),
26 respChan: newDatagramChannel(1),
27 }
28 mg := NewManager(transport, &log)
29
30 eyeballTracker := make(map[uuid.UUID]*datagramChannel)
31 for i := 0; i < sessions; i++ {
32 sessionID := uuid.New()
33 eyeballTracker[sessionID] = newDatagramChannel(1)
34 }
35
36 ctx, cancel := context.WithCancel(context.Background())
37 serveDone := make(chan struct{})
38 go func(ctx context.Context) {
39 mg.Serve(ctx)
40 close(serveDone)
41 }(ctx)
42
43 go func(ctx context.Context) {
44 for {
45 sessionID, payload, err := transport.respChan.Receive(ctx)
46 if err != nil {
47 require.Equal(t, context.Canceled, err)
48 return
49 }
50 respChan := eyeballTracker[sessionID]
51 require.NoError(t, respChan.Send(ctx, sessionID, payload))
52 }
53 }(ctx)
54
55 errGroup, ctx := errgroup.WithContext(ctx)
56 for sID, receiver := range eyeballTracker {
57 // Assign loop variables to local variables
58 sessionID := sID
59 eyeballRespReceiver := receiver
60 errGroup.Go(func() error {
61 payload := testPayload(sessionID)
62 expectResp := testResponse(payload)
63
64 cfdConn, originConn := net.Pipe()
65
66 origin := mockOrigin{
67 expectMsgCount: msgs,
68 expectedMsg: payload,
69 expectedResp: expectResp,
70 conn: originConn,
71 }
72 eyeball := mockEyeball{
73 expectMsgCount: msgs,
74 expectedMsg: expectResp,
75 expectSessionID: sessionID,
76 respReceiver: eyeballRespReceiver,
77 }
78
79 reqErrGroup, reqCtx := errgroup.WithContext(ctx)
80 reqErrGroup.Go(func() error {
81 return origin.serve()
82 })
83 reqErrGroup.Go(func() error {
84 return eyeball.serve(reqCtx)
85 })
86
87 session, err := mg.RegisterSession(ctx, sessionID, cfdConn)
88 require.NoError(t, err)
89
90 sessionDone := make(chan struct{})
91 go func() {
92 session.Serve(ctx, time.Minute*2)
93 close(sessionDone)
94 }()
95
96 for i := 0; i < msgs; i++ {
97 require.NoError(t, transport.newRequest(ctx, sessionID, testPayload(sessionID)))
98 }
99
100 // Make sure eyeball and origin have received all messages before unregistering the session
101 require.NoError(t, reqErrGroup.Wait())
102
103 require.NoError(t, mg.UnregisterSession(ctx, sessionID))
104 <-sessionDone
105
106 return nil
107 })
108 }
109
110 require.NoError(t, errGroup.Wait())
111 cancel()
112 transport.close()
113 <-serveDone
114}
115
116type mockOrigin struct {
117 expectMsgCount int
118 expectedMsg []byte
119 expectedResp []byte
120 conn io.ReadWriteCloser
121}
122
123func (mo *mockOrigin) serve() error {
124 expectedMsgLen := len(mo.expectedMsg)
125 readBuffer := make([]byte, expectedMsgLen+1)
126 for i := 0; i < mo.expectMsgCount; i++ {
127 n, err := mo.conn.Read(readBuffer)
128 if err != nil {
129 return err
130 }
131 if n != expectedMsgLen {
132 return fmt.Errorf("Expect to read %d bytes, read %d", expectedMsgLen, n)
133 }
134 if !bytes.Equal(readBuffer[:n], mo.expectedMsg) {
135 return fmt.Errorf("Expect %v, read %v", mo.expectedMsg, readBuffer[:n])
136 }
137
138 _, err = mo.conn.Write(mo.expectedResp)
139 if err != nil {
140 return err
141 }
142 }
143 return nil
144}
145
146func testPayload(sessionID uuid.UUID) []byte {
147 return []byte(fmt.Sprintf("Message from %s", sessionID))
148}
149
150func testResponse(msg []byte) []byte {
151 return []byte(fmt.Sprintf("Response to %v", msg))
152}
153
154type mockEyeball struct {
155 expectMsgCount int
156 expectedMsg []byte
157 expectSessionID uuid.UUID
158 respReceiver *datagramChannel
159}
160
161func (me *mockEyeball) serve(ctx context.Context) error {
162 for i := 0; i < me.expectMsgCount; i++ {
163 sessionID, msg, err := me.respReceiver.Receive(ctx)
164 if err != nil {
165 return err
166 }
167 if sessionID != me.expectSessionID {
168 return fmt.Errorf("Expect session %s, got %s", me.expectSessionID, sessionID)
169 }
170 if !bytes.Equal(msg, me.expectedMsg) {
171 return fmt.Errorf("Expect %v, read %v", me.expectedMsg, msg)
172 }
173 }
174 return nil
175}
176
177// datagramChannel is a channel for Datagram with wrapper to send/receive with context
178type datagramChannel struct {
179 datagramChan chan *newDatagram
180 closedChan chan struct{}
181}
182
183func newDatagramChannel(capacity uint) *datagramChannel {
184 return &datagramChannel{
185 datagramChan: make(chan *newDatagram, capacity),
186 closedChan: make(chan struct{}),
187 }
188}
189
190func (rc *datagramChannel) Send(ctx context.Context, sessionID uuid.UUID, payload []byte) error {
191 select {
192 case <-ctx.Done():
193 return ctx.Err()
194 case <-rc.closedChan:
195 return fmt.Errorf("datagram channel closed")
196 case rc.datagramChan <- &newDatagram{sessionID: sessionID, payload: payload}:
197 return nil
198 }
199}
200
201func (rc *datagramChannel) Receive(ctx context.Context) (uuid.UUID, []byte, error) {
202 select {
203 case <-ctx.Done():
204 return uuid.Nil, nil, ctx.Err()
205 case <-rc.closedChan:
206 return uuid.Nil, nil, fmt.Errorf("datagram channel closed")
207 case msg := <-rc.datagramChan:
208 return msg.sessionID, msg.payload, nil
209 }
210}
211
212func (rc *datagramChannel) Close() {
213 // No need to close msgChan, it will be garbage collect once there is no reference to it
214 close(rc.closedChan)
215}
216