TUN-5422: Define RPC to unregister session

Author: chungthuang

connection/quic.go

@@ -38,7 +38,6 @@ type QUICConnection struct {
 	logger         *zerolog.Logger
 	httpProxy      OriginProxy
 	sessionManager datagramsession.Manager
-	localIP        net.IP
 }
 
 // NewQUICConnection returns a new instance of QUICConnection.
@@ -75,17 +74,11 @@ func NewQUICConnection(
 
 	sessionManager := datagramsession.NewManager(datagramMuxer, observer.log)
 
-	localIP, err := getLocalIP()
-	if err != nil {
-		return nil, err
-	}
-
 	return &QUICConnection{
 		session:        session,
 		httpProxy:      httpProxy,
 		logger:         observer.log,
 		sessionManager: sessionManager,
-		localIP:        localIP,
 	}, nil
 }
 
@@ -197,7 +190,10 @@ func (q *QUICConnection) RegisterUdpSession(ctx context.Context, sessionID uuid.
 	return nil
 }
 
-// TODO: TUN-5422 Implement UnregisterUdpSession RPC
+func (q *QUICConnection) UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID) error {
+	q.sessionManager.UnregisterSession(ctx, sessionID)
+	return nil
+}
 
 // streamReadWriteAcker is a light wrapper over QUIC streams with a callback to send response back to
 // the client.
@@ -292,26 +288,3 @@ func isTransferEncodingChunked(req *http.Request) bool {
 	// separated value as well.
 	return strings.Contains(strings.ToLower(transferEncodingVal), "chunked")
 }
-
-// TODO: TUN-5303: Find the local IP once in ingress package
-// TODO: TUN-5421 allow user to specify which IP to bind to
-func getLocalIP() (net.IP, error) {
-	addrs, err := net.InterfaceAddrs()
-	if err != nil {
-		return nil, err
-	}
-	for _, addr := range addrs {
-		// Find the IP that is not loop back
-		var ip net.IP
-		switch v := addr.(type) {
-		case *net.IPNet:
-			ip = v.IP
-		case *net.IPAddr:
-			ip = v.IP
-		}
-		if !ip.IsLoopback() {
-			return ip, nil
-		}
-	}
-	return nil, fmt.Errorf("cannot determine IP to bind to")
-}

quic/quic_protocol.go

@@ -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()

quic/quic_protocol_test.go

@@ -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

tunnelrpc/pogs/sessionrpc.go

@@ -14,6 +14,7 @@ import (
 
 type SessionManager interface {
 	RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16) error
+	UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID) error
 }
 
 type SessionManager_PogsImpl struct {
@@ -60,6 +61,21 @@ func (i SessionManager_PogsImpl) RegisterUdpSession(p tunnelrpc.SessionManager_r
 	return resp.Marshal(result)
 }
 
+func (i SessionManager_PogsImpl) UnregisterUdpSession(p tunnelrpc.SessionManager_unregisterUdpSession) error {
+	server.Ack(p.Options)
+
+	sessionIDRaw, err := p.Params.SessionId()
+	if err != nil {
+		return err
+	}
+	sessionID, err := uuid.FromBytes(sessionIDRaw)
+	if err != nil {
+		return err
+	}
+
+	return i.impl.UnregisterUdpSession(p.Ctx, sessionID)
+}
+
 type RegisterUdpSessionResponse struct {
 	Err error
 }
@@ -116,3 +132,15 @@ func (c SessionManager_PogsClient) RegisterUdpSession(ctx context.Context, sessi
 	}
 	return response, nil
 }
+
+func (c SessionManager_PogsClient) UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID) error {
+	client := tunnelrpc.SessionManager{Client: c.Client}
+	promise := client.UnregisterUdpSession(ctx, func(p tunnelrpc.SessionManager_unregisterUdpSession_Params) error {
+		if err := p.SetSessionId(sessionID[:]); err != nil {
+			return err
+		}
+		return nil
+	})
+	_, err := promise.Struct()
+	return err
+}

tunnelrpc/tunnelrpc.capnp

@@ -149,4 +149,5 @@ struct RegisterUdpSessionResponse {
 
 interface SessionManager {
     registerUdpSession @0 (sessionId :Data, dstIp :Data, dstPort: UInt16) -> (result :RegisterUdpSessionResponse);
+    unregisterUdpSession @1 (sessionId :Data) -> ();
 }