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-5422: Define RPC to unregister session

Browse Source
This commit is contained in:
cthuang
2021-11-30 19:58:11 +00:00
committed by Arég Harutyunyan
parent 7e47667b08
commit b73c588254
6 changed files with 417 additions and 227 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
quic/quic_protocol.go
View File

@@ -247,6 +247,10 @@ 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) Close() {
_ = rcs.client.Close()
_ = rcs.transport.Close()

18
quic/quic_protocol_test.go
View File

@@ -131,12 +131,15 @@ func TestRegisterUdpSession(t *testing.T) {
rpcClientStream, err := NewRPCClientStream(context.Background(), clientStream, &logger)
assert.NoError(t, err)
err = rpcClientStream.RegisterUdpSession(context.Background(), rpcServer.sessionID, rpcServer.dstIP, rpcServer.dstPort)
assert.NoError(t, err)
assert.NoError(t, rpcClientStream.RegisterUdpSession(context.Background(), rpcServer.sessionID, rpcServer.dstIP, rpcServer.dstPort))
// Different sessionID, the RPC server should reject the registraion
err = rpcClientStream.RegisterUdpSession(context.Background(), uuid.New(), rpcServer.dstIP, rpcServer.dstPort)
assert.Error(t, err)
assert.Error(t, rpcClientStream.RegisterUdpSession(context.Background(), uuid.New(), rpcServer.dstIP, rpcServer.dstPort))
assert.NoError(t, rpcClientStream.UnregisterUdpSession(context.Background(), rpcServer.sessionID))
// Different sessionID, the RPC server should reject the unregistraion
assert.Error(t, rpcClientStream.UnregisterUdpSession(context.Background(), uuid.New()))
rpcClientStream.Close()
<-sessionRegisteredChan
@@ -161,6 +164,13 @@ func (s mockRPCServer) RegisterUdpSession(ctx context.Context, sessionID uuid.UU
return nil
}
func (s mockRPCServer) UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID) error {
if s.sessionID != sessionID {
return fmt.Errorf("expect session ID %s, got %s", s.sessionID, sessionID)
}
return nil
}
type mockRPCStream struct {
io.ReadCloser
io.WriteCloser