feat(orc8r): add DeMap endpoints to directoryd (#9739)

magma · Oct 22, 2021 · c025e97 · c025e97
1 parent 5823a10
commit c025e97
Show file tree

Hide file tree

Showing 11 changed files with 832 additions and 140 deletions.
diff --git a/orc8r/cloud/go/services/directoryd/client_api.go b/orc8r/cloud/go/services/directoryd/client_api.go
@@ -84,6 +84,21 @@ func MapHWIDsToHostnames(ctx context.Context, hwidToHostname map[string]string)
 	return nil
 }
 
+// UnmapHWIDsToHostnames removes the {hwid -> hostname} map
+// Derived state, stored in directoryd service.
+func UnmapHWIDsToHostnames(ctx context.Context, hwids []string) error {
+	client, err := getDirectorydClient()
+	if err != nil {
+		return errors.Wrap(err, "failed to get directoryd client")
+	}
+
+	_, err = client.UnmapHWIDsToHostnames(ctx, &protos.UnmapHWIDToHostnameRequest{Hwids: hwids})
+	if err != nil {
+		return fmt.Errorf("failed to ummap hwids to hostnames %v: %s", hwids, err)
+	}
+	return nil
+}
+
 // GetIMSIForSessionID returns the IMSI mapped to by session ID.
 // Derived state, stored in directoryd service.
 // NOTE: this mapping is provided on a best-effort basis, meaning
@@ -125,6 +140,25 @@ func MapSessionIDsToIMSIs(ctx context.Context, networkID string, sessionIDToIMSI
 	return nil
 }
 
+// MapSessionIDsToIMSIs removes {session ID -> IMSI} mapping
+// Derived state, stored in directoryd service.
+func UnmapSessionIDsToIMSIs(ctx context.Context, networkID string, sessionIDs []string) error {
+	client, err := getDirectorydClient()
+	if err != nil {
+		return errors.Wrap(err, "failed to get directoryd client")
+	}
+
+	_, err = client.UnmapSessionIDsToIMSIs(ctx, &protos.UnmapSessionIDToIMSIRequest{
+		NetworkID:  networkID,
+		SessionIDs: sessionIDs,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to unmap session IDs %v under network ID %s: %s", sessionIDs, networkID, err)
+	}
+
+	return nil
+}
+
 // GetHWIDForSgwCTeid returns the HwID mapped to by teid
 // Derived state, stored in directoryd service.
 // NOTE: this mapping is provided on a best-effort basis, meaning
@@ -165,6 +199,24 @@ func MapSgwCTeidToHWID(ctx context.Context, networkID string, teidToHWID map[str
 	return nil
 }
 
+// UnmapSgwCTeidToHWID removes {Teid -> HwId} mapping
+func UnmapSgwCTeidToHWID(ctx context.Context, networkID string, teids []string) error {
+	client, err := getDirectorydClient()
+	if err != nil {
+		return errors.Wrap(err, "failed to get directoryd client")
+	}
+
+	_, err = client.UnmapSgwCTeidToHWID(ctx, &protos.UnmapSgwCTeidToHWIDRequest{
+		NetworkID: networkID,
+		Teids:     teids,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to ummap sgw c teid %v under network ID %s: %s", teids, networkID, err)
+	}
+
+	return nil
+}
+
 //--------------------------
 // State service client APIs
 //--------------------------

diff --git a/orc8r/cloud/go/services/directoryd/client_api_test.go b/orc8r/cloud/go/services/directoryd/client_api_test.go
@@ -155,6 +155,24 @@ func TestDirectorydMethods(t *testing.T) {
 	hwid, err := directoryd.GetHWIDForSgwCTeid(context.Background(), nid0, teid0)
 	assert.NoError(t, err)
 	assert.Equal(t, hwid0, hwid)
+
+	// Remove sid0->imsi0 mapping
+	err = directoryd.UnmapSessionIDsToIMSIs(context.Background(), nid0, []string{sid0})
+	assert.NoError(t, err)
+	_, err = directoryd.GetIMSIForSessionID(context.Background(), nid0, sid0)
+	assert.Error(t, err)
+
+	// Remove hwid->imsi mapping
+	err = directoryd.UnmapHWIDsToHostnames(context.Background(), []string{hwid1})
+	assert.NoError(t, err)
+	_, err = directoryd.GetHostnameForHWID(context.Background(), hwid1)
+	assert.Error(t, err)
+
+	// Remove teid->hwid mapping
+	err = directoryd.UnmapSgwCTeidToHWID(context.Background(), nid0, []string{teid0})
+	assert.NoError(t, err)
+	_, err = directoryd.GetHWIDForSgwCTeid(context.Background(), nid0, teid0)
+	assert.Error(t, err)
 }
 
 func TestDirectorydStateMethods(t *testing.T) {

diff --git a/orc8r/cloud/go/services/directoryd/servicers/servicer.go b/orc8r/cloud/go/services/directoryd/servicers/servicer.go
@@ -36,7 +36,7 @@ func (d *directoryLookupServicer) GetHostnameForHWID(
 ) (*protos.GetHostnameForHWIDResponse, error) {
 	err := req.Validate()
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to validate request")
+		return nil, errors.Wrap(err, "failed to validate GetHostnameForHWIDRequest")
 	}
 
 	hostname, err := d.store.GetHostnameForHWID(req.Hwid)
@@ -48,20 +48,32 @@ func (d *directoryLookupServicer) GetHostnameForHWID(
 func (d *directoryLookupServicer) MapHWIDsToHostnames(ctx context.Context, req *protos.MapHWIDToHostnameRequest) (*protos.Void, error) {
 	err := req.Validate()
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to validate request")
+		return nil, errors.Wrap(err, "failed to validate MapHWIDToHostnameRequest")
 	}
 
 	err = d.store.MapHWIDsToHostnames(req.HwidToHostname)
 
 	return &protos.Void{}, err
 }
 
+func (d *directoryLookupServicer) UnmapHWIDsToHostnames(ctx context.Context, req *protos.UnmapHWIDToHostnameRequest) (*protos.Void, error) {
+	err := req.Validate()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to validate UnmapHWIDToHostnameRequest")
+	}
+
+	err = d.store.UnmapHWIDsToHostnames(req.Hwids)
+
+	return &protos.Void{}, err
+}
+
 func (d *directoryLookupServicer) GetIMSIForSessionID(
 	ctx context.Context, req *protos.GetIMSIForSessionIDRequest,
 ) (*protos.GetIMSIForSessionIDResponse, error) {
 	err := req.Validate()
+
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to validate request")
+		return nil, errors.Wrap(err, "failed to validate GetIMSIForSessionIDRequest")
 	}
 
 	imsi, err := d.store.GetIMSIForSessionID(req.NetworkID, req.SessionID)
@@ -73,31 +85,53 @@ func (d *directoryLookupServicer) GetIMSIForSessionID(
 func (d *directoryLookupServicer) MapSessionIDsToIMSIs(ctx context.Context, req *protos.MapSessionIDToIMSIRequest) (*protos.Void, error) {
 	err := req.Validate()
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to validate request")
+		return nil, errors.Wrap(err, "failed to validate MapSessionIDToIMSIRequest")
 	}
 
 	err = d.store.MapSessionIDsToIMSIs(req.NetworkID, req.SessionIDToIMSI)
 
 	return &protos.Void{}, err
 }
 
+func (d *directoryLookupServicer) UnmapSessionIDsToIMSIs(ctx context.Context, req *protos.UnmapSessionIDToIMSIRequest) (*protos.Void, error) {
+	err := req.Validate()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to validate UnmapSessionIDToIMSIRequest")
+	}
+
+	err = d.store.UnmapSessionIDsToIMSIs(req.NetworkID, req.SessionIDs)
+
+	return &protos.Void{}, err
+}
+
 func (d *directoryLookupServicer) MapSgwCTeidToHWID(ctx context.Context, req *protos.MapSgwCTeidToHWIDRequest) (*protos.Void, error) {
 	err := req.Validate()
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to validate request")
+		return nil, errors.Wrap(err, "failed to validate MapSgwCTeidToHWIDRequest")
 	}
 
 	err = d.store.MapSgwCTeidToHWID(req.NetworkID, req.TeidToHwid)
 
 	return &protos.Void{}, err
 }
 
+func (d *directoryLookupServicer) UnmapSgwCTeidToHWID(ctx context.Context, req *protos.UnmapSgwCTeidToHWIDRequest) (*protos.Void, error) {
+	err := req.Validate()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to validate UnmapSgwCTeidToHWIDRequest")
+	}
+
+	err = d.store.UnmapSgwCTeidToHWID(req.NetworkID, req.Teids)
+
+	return &protos.Void{}, err
+}
+
 func (d *directoryLookupServicer) GetHWIDForSgwCTeid(
 	ctx context.Context, req *protos.GetHWIDForSgwCTeidRequest,
 ) (*protos.GetHWIDForSgwCTeidResponse, error) {
 	err := req.Validate()
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to validate request")
+		return nil, errors.Wrap(err, "failed to validate GetHWIDForSgwCTeidRequest")
 	}
 
 	hwid, err := d.store.GetHWIDForSgwCTeid(req.NetworkID, req.Teid)

diff --git a/orc8r/cloud/go/services/directoryd/servicers/servicer_test.go b/orc8r/cloud/go/services/directoryd/servicers/servicer_test.go
@@ -42,6 +42,10 @@ const (
 	sid0  = "some_sessionid_0"
 	sid1  = "some_sessionid_1"
 	sid2  = "some_sessionid_2"
+
+	teid0 = "10"
+	teid1 = "11"
+	teid2 = "12"
 )
 
 func newTestDirectoryLookupServicer(t *testing.T) protos.DirectoryLookupServer {
@@ -91,6 +95,14 @@ func TestDirectoryLookupServicer_HostnameToHWID(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, hn2, res.Hostname)
 
+	// DeMap hwid2
+	put_demap := &protos.UnmapHWIDToHostnameRequest{Hwids: []string{hwid2}}
+	_, err = srv.UnmapHWIDsToHostnames(ctx, put_demap)
+	assert.NoError(t, err)
+	get = &protos.GetHostnameForHWIDRequest{Hwid: hwid2}
+	_, err = srv.GetHostnameForHWID(ctx, get)
+	assert.Error(t, err)
+
 	// hwid0->hostname0 still intact
 	get = &protos.GetHostnameForHWIDRequest{Hwid: hwid0}
 	res, err = srv.GetHostnameForHWID(ctx, get)
@@ -136,6 +148,14 @@ func TestDirectoryLookupServicer_SessionIDToIMSI(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, imsi0, res.Imsi)
 
+	// DeMap sid0 from network nid1
+	put_demap := &protos.UnmapSessionIDToIMSIRequest{NetworkID: nid1, SessionIDs: []string{sid0}}
+	_, err = srv.UnmapSessionIDsToIMSIs(ctx, put_demap)
+	assert.NoError(t, err)
+	get = &protos.GetIMSIForSessionIDRequest{NetworkID: nid1, SessionID: sid0}
+	_, err = srv.GetIMSIForSessionID(ctx, get)
+	assert.Error(t, err)
+
 	// Correctly network-partitioned: {nid0: sid0->imsi0, nid1: sid0->imsi1}
 	put = &protos.MapSessionIDToIMSIRequest{NetworkID: nid0, SessionIDToIMSI: map[string]string{sid0: imsi0}}
 	_, err = srv.MapSessionIDsToIMSIs(ctx, put)
@@ -160,3 +180,74 @@ func TestDirectoryLookupServicer_SessionIDToIMSI(t *testing.T) {
 	_, err = srv.MapSessionIDsToIMSIs(ctx, put)
 	assert.Error(t, err)
 }
+
+func TestDirectoryLookupServicer_TeidToHWID(t *testing.T) {
+	srv := newTestDirectoryLookupServicer(t)
+	stateTestInit.StartTestService(t)
+	ctx := context.Background()
+
+	// Empty initially
+	get := &protos.GetHWIDForSgwCTeidRequest{NetworkID: nid0, Teid: teid0}
+	_, err := srv.GetHWIDForSgwCTeid(ctx, get)
+	assert.Error(t, err)
+
+	// Put and get teid0->hwid0
+	put := &protos.MapSgwCTeidToHWIDRequest{NetworkID: nid0, TeidToHwid: map[string]string{teid0: hwid0}}
+	_, err = srv.MapSgwCTeidToHWID(ctx, put)
+	assert.NoError(t, err)
+	get = &protos.GetHWIDForSgwCTeidRequest{NetworkID: nid0, Teid: teid0}
+	res, err := srv.GetHWIDForSgwCTeid(ctx, get)
+	assert.NoError(t, err)
+	assert.Equal(t, hwid0, res.Hwid)
+
+	// Put and get sid1->imsi1, sid2->imsi2
+	put = &protos.MapSgwCTeidToHWIDRequest{NetworkID: nid0, TeidToHwid: map[string]string{teid1: hwid1, teid2: hwid2}}
+	_, err = srv.MapSgwCTeidToHWID(ctx, put)
+	assert.NoError(t, err)
+	get = &protos.GetHWIDForSgwCTeidRequest{NetworkID: nid0, Teid: teid1}
+	res, err = srv.GetHWIDForSgwCTeid(ctx, get)
+	assert.NoError(t, err)
+	assert.Equal(t, hwid1, res.Hwid)
+	get = &protos.GetHWIDForSgwCTeidRequest{NetworkID: nid0, Teid: teid2}
+	res, err = srv.GetHWIDForSgwCTeid(ctx, get)
+	assert.NoError(t, err)
+	assert.Equal(t, hwid2, res.Hwid)
+
+	// sid0->imsi0 still intact
+	get = &protos.GetHWIDForSgwCTeidRequest{NetworkID: nid0, Teid: teid0}
+	res, err = srv.GetHWIDForSgwCTeid(ctx, get)
+	assert.NoError(t, err)
+	assert.Equal(t, hwid0, res.Hwid)
+
+	// Correctly network-partitioned: {nid0: sid0->imsi0, nid1: sid0->imsi1}
+	put = &protos.MapSgwCTeidToHWIDRequest{NetworkID: nid0, TeidToHwid: map[string]string{teid0: hwid0}}
+	_, err = srv.MapSgwCTeidToHWID(ctx, put)
+	assert.NoError(t, err)
+	put = &protos.MapSgwCTeidToHWIDRequest{NetworkID: nid1, TeidToHwid: map[string]string{teid0: hwid2}}
+	_, err = srv.MapSgwCTeidToHWID(ctx, put)
+	assert.NoError(t, err)
+	get = &protos.GetHWIDForSgwCTeidRequest{NetworkID: nid0, Teid: teid0}
+	res, err = srv.GetHWIDForSgwCTeid(ctx, get)
+	assert.NoError(t, err)
+	assert.Equal(t, hwid0, res.Hwid)
+	get = &protos.GetHWIDForSgwCTeidRequest{NetworkID: nid1, Teid: teid0}
+	res, err = srv.GetHWIDForSgwCTeid(ctx, get)
+	assert.NoError(t, err)
+	assert.Equal(t, hwid2, res.Hwid)
+
+	// DeMap teid0 from network nid1
+	put_demap := &protos.UnmapSgwCTeidToHWIDRequest{NetworkID: nid1, Teids: []string{teid0}}
+	_, err = srv.UnmapSgwCTeidToHWID(ctx, put_demap)
+	assert.NoError(t, err)
+	get = &protos.GetHWIDForSgwCTeidRequest{NetworkID: nid1, Teid: teid0}
+	_, err = srv.GetHWIDForSgwCTeid(ctx, get)
+	assert.Error(t, err)
+
+	// Fail with empty network ID
+	get = &protos.GetHWIDForSgwCTeidRequest{Teid: teid0}
+	_, err = srv.GetHWIDForSgwCTeid(ctx, get)
+	assert.Error(t, err)
+	put = &protos.MapSgwCTeidToHWIDRequest{TeidToHwid: map[string]string{teid0: hwid0}}
+	_, err = srv.MapSgwCTeidToHWID(ctx, put)
+	assert.Error(t, err)
+}
diff --git a/orc8r/cloud/go/services/directoryd/storage/storage.go b/orc8r/cloud/go/services/directoryd/storage/storage.go
@@ -22,15 +22,24 @@ type DirectorydStorage interface {
 	// MapHWIDsToHostnames maps {hwid -> hostname}.
 	MapHWIDsToHostnames(hwidToHostname map[string]string) error
 
+	// UnmapSgwCTeidToHWID removes {hwid -> hostname} for a specific hwid
+	UnmapHWIDsToHostnames(hwids []string) error
+
 	// GetIMSIForSessionID returns the IMSI mapped to by session ID.
 	GetIMSIForSessionID(networkID, sessionID string) (string, error)
 
 	// MapSessionIDsToIMSIs maps {session ID -> IMSI}.
 	MapSessionIDsToIMSIs(networkID string, sessionIDToIMSI map[string]string) error
 
+	// UnmapSessionIDsToIMSIs removes {session ID -> IMSI} for a specific sessionID
+	UnmapSessionIDsToIMSIs(networkID string, sessionIDs []string) error
+
 	// GetHWIDForSgwCTeid returns the HwId mapped by teid
 	GetHWIDForSgwCTeid(networkID, teid string) (string, error)
 
 	// MapSgwCTeidToHWID maps {teid -> hwid}
 	MapSgwCTeidToHWID(networkID string, s8TeidToHwid map[string]string) error
+
+	// UnmapSgwCTeidToHWID removes {teid -> hwid} for a specific teid
+	UnmapSgwCTeidToHWID(networkID string, teids []string) error
 }