Fix the storage of TCONT in bbsim.

Change-Id: Ia6c702fd0ec6c3c478ffda364115cd076338ac5a
diff --git a/internal/bbsim/api/grpc_api_server.go b/internal/bbsim/api/grpc_api_server.go
index ec25a10..8a9101d 100644
--- a/internal/bbsim/api/grpc_api_server.go
+++ b/internal/bbsim/api/grpc_api_server.go
@@ -89,10 +89,10 @@
 			allocatedGemPorts = append(allocatedGemPorts, resource)
 		}
 
-		for k, v := range pon.AllocatedAllocIds {
+		for _, v := range pon.AllocatedAllocIds {
 			resource := &bbsim.PonAllocatedResources{
-				SerialNumber: common.OnuSnToString(v),
-				Id:           int32(k),
+				SerialNumber: common.OnuSnToString(v.OnuSn),
+				Id:           int32(v.AllocID),
 			}
 			allocatedAllocIds = append(allocatedAllocIds, resource)
 		}
diff --git a/internal/bbsim/devices/onu.go b/internal/bbsim/devices/onu.go
index 1353a50..e6e4e58 100644
--- a/internal/bbsim/devices/onu.go
+++ b/internal/bbsim/devices/onu.go
@@ -382,7 +382,7 @@
 	// clean the ONU state
 	o.Flows = []FlowKey{}
 	o.PonPort.removeOnuId(o.ID)
-	o.PonPort.removeAllocId(o.SerialNumber)
+	o.PonPort.removeAllocIdsForOnuSn(o.SerialNumber)
 	o.PonPort.removeGemPortBySn(o.SerialNumber)
 
 	o.onuAlarmsInfoLock.Lock()
@@ -807,7 +807,7 @@
 			o.MibDataSync = 0
 
 			// if the MIB reset is successful then remove all the stored AllocIds and GemPorts
-			o.PonPort.removeAllocId(o.SerialNumber)
+			o.PonPort.removeAllocIdsForOnuSn(o.SerialNumber)
 			o.PonPort.removeGemPortBySn(o.SerialNumber)
 		}
 	case omci.MibUploadRequestType:
@@ -887,6 +887,7 @@
 			}).Debug("set-onu-admin-lock-state")
 		case me.TContClassID:
 			allocId := msgObj.Attributes[me.TCont_AllocId].(uint16)
+			entityID := msgObj.Attributes["ManagedEntityId"].(uint16)
 
 			// if the AllocId is 255 (0xFF) or 65535 (0xFFFF) it means we are removing it,
 			// otherwise we are adding it
@@ -898,15 +899,15 @@
 					"AllocId":      allocId,
 					"SerialNumber": o.Sn(),
 				}).Trace("freeing-alloc-id-via-omci")
-				o.PonPort.removeAllocId(o.SerialNumber)
+				o.PonPort.removeAllocId(o.PonPortID, o.ID, entityID)
 			} else {
-				if used, sn := o.PonPort.isAllocIdAllocated(allocId); used {
+				if used, allocObj := o.PonPort.isAllocIdAllocated(o.PonPortID, o.ID, entityID); used {
 					onuLogger.WithFields(log.Fields{
 						"IntfId":       o.PonPortID,
 						"OnuId":        o.ID,
 						"AllocId":      allocId,
 						"SerialNumber": o.Sn(),
-					}).Errorf("allocid-already-allocated-to-onu-with-sn-%s", common.OnuSnToString(sn))
+					}).Errorf("allocid-already-allocated-to-onu-with-sn-%s", common.OnuSnToString(allocObj.OnuSn))
 					success = false
 				} else {
 					onuLogger.WithFields(log.Fields{
@@ -916,7 +917,7 @@
 						"AllocId":      allocId,
 						"SerialNumber": o.Sn(),
 					}).Trace("storing-alloc-id-via-omci")
-					o.PonPort.storeAllocId(allocId, o.SerialNumber)
+					o.PonPort.storeAllocId(o.PonPortID, o.ID, entityID, allocId, o.SerialNumber)
 				}
 			}
 		case me.EthernetFrameExtendedPmClassID,
diff --git a/internal/bbsim/devices/onu_omci_test.go b/internal/bbsim/devices/onu_omci_test.go
index 139560c..b6db6b7 100644
--- a/internal/bbsim/devices/onu_omci_test.go
+++ b/internal/bbsim/devices/onu_omci_test.go
@@ -284,7 +284,7 @@
 
 	// create a GemPort and an AllocId for this ONU
 	onu.PonPort.storeGemPort(1024, onu.SerialNumber)
-	onu.PonPort.storeAllocId(1024, onu.SerialNumber)
+	onu.PonPort.storeAllocId(1024, 1024, 0x8001, 1024, onu.SerialNumber)
 
 	// send a MibReset
 	err := onu.handleOmciRequest(makeOmciMessage(t, onu, makeOmciMibResetRequest(t)), stream)
@@ -325,7 +325,7 @@
 
 	// the MDS has to be between 1 and 255, since 0 is valid for a reset
 	assert.Assert(t, onu.MibDataSync > 0)
-	assert.Assert(t, onu.MibDataSync <= 255)
+	// assert.Assert(t, onu.MibDataSync <= 255) // This is always true since 'MibDataSync' is uint8
 }
 
 func Test_GemPortValidation(t *testing.T) {
diff --git a/internal/bbsim/devices/onu_state_machine_test.go b/internal/bbsim/devices/onu_state_machine_test.go
index 1c8470c..a072595 100644
--- a/internal/bbsim/devices/onu_state_machine_test.go
+++ b/internal/bbsim/devices/onu_state_machine_test.go
@@ -50,7 +50,7 @@
 	}
 	onu.onuAlarmsInfo[key] = omcilib.OnuAlarmInfo{SequenceNo: 1, AlarmBitMap: [28]byte{}}
 	onu.PonPort.storeOnuId(onu.ID, onu.SerialNumber)
-	onu.PonPort.storeAllocId(1, onu.SerialNumber)
+	onu.PonPort.storeAllocId(1, 1024, 0x8001, 1024, onu.SerialNumber)
 	onu.PonPort.storeGemPort(1, onu.SerialNumber)
 
 	_ = onu.InternalState.Event(OnuTxDisable)
diff --git a/internal/bbsim/devices/onu_test_helpers.go b/internal/bbsim/devices/onu_test_helpers.go
index e255f5b..ba04b57 100644
--- a/internal/bbsim/devices/onu_test_helpers.go
+++ b/internal/bbsim/devices/onu_test_helpers.go
@@ -153,7 +153,7 @@
 		PonPortID: ponPortId,
 		PonPort: &PonPort{
 			AllocatedGemPorts: make(map[uint16]*openolt.SerialNumber),
-			AllocatedAllocIds: make(map[uint16]*openolt.SerialNumber),
+			AllocatedAllocIds: make(map[AllocIDKey]*AllocIDVal),
 			Olt:               &OltDevice{},
 		},
 		OmciResponseRate: 10,
diff --git a/internal/bbsim/devices/pon.go b/internal/bbsim/devices/pon.go
index 4db0e09..c1b3257 100644
--- a/internal/bbsim/devices/pon.go
+++ b/internal/bbsim/devices/pon.go
@@ -30,6 +30,17 @@
 	"module": "PON",
 })
 
+type AllocIDVal struct {
+	OnuSn   *openolt.SerialNumber
+	AllocID uint16
+}
+
+type AllocIDKey struct {
+	PonID    uint32
+	OnuID    uint32
+	EntityID uint16
+}
+
 type PonPort struct {
 	// BBSIM Internals
 	ID            uint32
@@ -50,7 +61,7 @@
 	allocatedGemPortsLock sync.RWMutex
 	AllocatedOnuIds       map[uint32]*openolt.SerialNumber
 	allocatedOnuIdsLock   sync.RWMutex
-	AllocatedAllocIds     map[uint16]*openolt.SerialNumber
+	AllocatedAllocIds     map[AllocIDKey]*AllocIDVal // key is AllocIDKey
 	allocatedAllocIdsLock sync.RWMutex
 }
 
@@ -65,7 +76,7 @@
 		Onus:              []*Onu{},
 		AllocatedGemPorts: make(map[uint16]*openolt.SerialNumber),
 		AllocatedOnuIds:   make(map[uint32]*openolt.SerialNumber),
-		AllocatedAllocIds: make(map[uint16]*openolt.SerialNumber),
+		AllocatedAllocIds: make(map[AllocIDKey]*AllocIDVal),
 	}
 
 	ponPort.InternalState = fsm.NewFSM(
@@ -291,20 +302,27 @@
 }
 
 // storeAllocId adds the Id to the ONU Ids already allocated to this PON port
-func (p *PonPort) storeAllocId(allocId uint16, onuSn *openolt.SerialNumber) {
+func (p *PonPort) storeAllocId(ponID uint32, onuID uint32, entityID uint16, allocId uint16, onuSn *openolt.SerialNumber) {
 	p.allocatedAllocIdsLock.Lock()
 	defer p.allocatedAllocIdsLock.Unlock()
-	p.AllocatedAllocIds[allocId] = onuSn
+	p.AllocatedAllocIds[AllocIDKey{ponID, onuID, entityID}] = &AllocIDVal{onuSn, allocId}
 }
 
 // removeAllocId removes the AllocId from the allocated resources
-// this is done via SN as the AllocId is not remove but set to a default value
-func (p *PonPort) removeAllocId(onuSn *openolt.SerialNumber) {
+func (p *PonPort) removeAllocId(ponID uint32, onuID uint32, entityID uint16) {
 	p.allocatedAllocIdsLock.Lock()
 	defer p.allocatedAllocIdsLock.Unlock()
-	for allocId, sn := range p.AllocatedAllocIds {
-		if sn == onuSn {
-			delete(p.AllocatedAllocIds, allocId)
+	allocKey := AllocIDKey{ponID, onuID, entityID}
+	delete(p.AllocatedAllocIds, allocKey)
+}
+
+// removeAllocIdsForOnuSn removes the all AllocIds for the given onu serial number
+func (p *PonPort) removeAllocIdsForOnuSn(onuSn *openolt.SerialNumber) {
+	p.allocatedAllocIdsLock.Lock()
+	defer p.allocatedAllocIdsLock.Unlock()
+	for id, allocObj := range p.AllocatedAllocIds {
+		if onuSn == allocObj.OnuSn {
+			delete(p.AllocatedAllocIds, id)
 		}
 	}
 }
@@ -312,16 +330,16 @@
 func (p *PonPort) removeAllAllocIds() {
 	p.allocatedAllocIdsLock.Lock()
 	defer p.allocatedAllocIdsLock.Unlock()
-	p.AllocatedAllocIds = make(map[uint16]*openolt.SerialNumber)
+	p.AllocatedAllocIds = make(map[AllocIDKey]*AllocIDVal)
 }
 
 // isAllocIdAllocated returns whether this AllocId is already in use on this PON
-func (p *PonPort) isAllocIdAllocated(allocId uint16) (bool, *openolt.SerialNumber) {
+func (p *PonPort) isAllocIdAllocated(ponID uint32, onuID uint32, entityID uint16) (bool, *AllocIDVal) {
 	p.allocatedAllocIdsLock.RLock()
 	defer p.allocatedAllocIdsLock.RUnlock()
-
-	if _, ok := p.AllocatedAllocIds[allocId]; ok {
-		return true, p.AllocatedAllocIds[allocId]
+	allocKey := AllocIDKey{ponID, onuID, entityID}
+	if _, ok := p.AllocatedAllocIds[allocKey]; ok {
+		return true, p.AllocatedAllocIds[allocKey]
 	}
 	return false, nil
 }
diff --git a/internal/bbsim/devices/pon_test.go b/internal/bbsim/devices/pon_test.go
index 5467413..f08c37c 100644
--- a/internal/bbsim/devices/pon_test.go
+++ b/internal/bbsim/devices/pon_test.go
@@ -128,24 +128,24 @@
 // the allocId is never removed, is always set to either 255 or 65535
 func Test_removeAllocId(t *testing.T) {
 
-	const (
-		allocId1 = 1024
-		allocId2 = 1025
-	)
+	const entityID1 uint16 = 1024
+	const entityID2 uint16 = 1025
+	const allocId1 uint16 = 1024
+	const allocId2 uint16 = 1025
 
 	pon := &PonPort{
-		AllocatedAllocIds: make(map[uint16]*openolt.SerialNumber),
+		AllocatedAllocIds: make(map[AllocIDKey]*AllocIDVal),
 	}
 
-	pon.AllocatedAllocIds[allocId1] = sn1
-	pon.AllocatedAllocIds[allocId2] = sn2
+	pon.AllocatedAllocIds[AllocIDKey{0, 1, entityID1}] = &AllocIDVal{sn1, allocId1}
+	pon.AllocatedAllocIds[AllocIDKey{1, 1, entityID2}] = &AllocIDVal{sn2, allocId2}
 
 	assert.Equal(t, len(pon.AllocatedAllocIds), 2)
 
-	pon.removeAllocId(sn1)
+	pon.removeAllocId(0, 1, entityID1)
 
 	assert.Equal(t, len(pon.AllocatedAllocIds), 1)
-	assert.Nil(t, pon.AllocatedAllocIds[allocId1])
-	assert.Equal(t, pon.AllocatedAllocIds[allocId2], sn2)
-
+	assert.NotContains(t, pon.AllocatedAllocIds, AllocIDKey{0, 1, entityID1})
+	assert.Contains(t, pon.AllocatedAllocIds, AllocIDKey{1, 1, entityID2})
+	assert.Equal(t, pon.AllocatedAllocIds[AllocIDKey{1, 1, entityID2}].OnuSn, sn2)
 }