hoelee/cloudflared
1
0
Fork 0
mirror of https://github.com/cloudflare/cloudflared.git synced 2025-07-27 20:50:00 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

TUN-5494: Send a RPC with terminate reason to edge if the session is closed locally

Browse Source
This commit is contained in:
cthuang
2021-12-14 22:52:47 +00:00
parent 70e675f42c
commit ebae7a7024
12 changed files with 563 additions and 297 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
quic/quic_protocol.go
View File

@@ -251,8 +251,8 @@ func (rcs *RPCClientStream) RegisterUdpSession(ctx context.Context, sessionID uu
return resp.Err
}
func (rcs *RPCClientStream) UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID) error {
return rcs.client.UnregisterUdpSession(ctx, sessionID)
func (rcs *RPCClientStream) UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID, message string) error {
return rcs.client.UnregisterUdpSession(ctx, sessionID, message)
}
func (rcs *RPCClientStream) Close() {

28
quic/quic_protocol_test.go
View File

@@ -114,11 +114,13 @@ func TestRegisterUdpSession(t *testing.T) {
clientStream := mockRPCStream{clientReader, clientWriter}
serverStream := mockRPCStream{serverReader, serverWriter}
unregisterMessage := "closed by eyeball"
rpcServer := mockRPCServer{
sessionID: uuid.New(),
dstIP: net.IP{172, 16, 0, 1},
dstPort: 8000,
closeIdleAfter: testCloseIdleAfterHint,
sessionID: uuid.New(),
dstIP: net.IP{172, 16, 0, 1},
dstPort: 8000,
closeIdleAfter: testCloseIdleAfterHint,
unregisterMessage: unregisterMessage,
}
logger := zerolog.Nop()
sessionRegisteredChan := make(chan struct{})
@@ -142,20 +144,21 @@ func TestRegisterUdpSession(t *testing.T) {
// Different sessionID, the RPC server should reject the registraion
assert.Error(t, rpcClientStream.RegisterUdpSession(context.Background(), uuid.New(), rpcServer.dstIP, rpcServer.dstPort, testCloseIdleAfterHint))
assert.NoError(t, rpcClientStream.UnregisterUdpSession(context.Background(), rpcServer.sessionID))
assert.NoError(t, rpcClientStream.UnregisterUdpSession(context.Background(), rpcServer.sessionID, unregisterMessage))
// Different sessionID, the RPC server should reject the unregistraion
assert.Error(t, rpcClientStream.UnregisterUdpSession(context.Background(), uuid.New()))
assert.Error(t, rpcClientStream.UnregisterUdpSession(context.Background(), uuid.New(), unregisterMessage))
rpcClientStream.Close()
<-sessionRegisteredChan
}
type mockRPCServer struct {
sessionID uuid.UUID
dstIP net.IP
dstPort uint16
closeIdleAfter time.Duration
sessionID uuid.UUID
dstIP net.IP
dstPort uint16
closeIdleAfter time.Duration
unregisterMessage string
}
func (s mockRPCServer) RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16, closeIdleAfter time.Duration) error {
@@ -174,10 +177,13 @@ func (s mockRPCServer) RegisterUdpSession(ctx context.Context, sessionID uuid.UU
return nil
}
func (s mockRPCServer) UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID) error {
func (s mockRPCServer) UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID, message string) error {
if s.sessionID != sessionID {
return fmt.Errorf("expect session ID %s, got %s", s.sessionID, sessionID)
}
if s.unregisterMessage != message {
return fmt.Errorf("expect unregister message %s, got %s", s.unregisterMessage, message)
}
return nil
}