Merge "Fix the storage of TCONT in bbsim. The TCONTs should be referenced by EntityID and not AllocID. The AllocID is 0xff or 0xffff during reset so we can't find its reference in the map. However the EntityID is always passed correctly to reference the TCONT."
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 e460956..089196f 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["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(entityID)
 			} else {
-				if used, sn := o.PonPort.isAllocIdAllocated(allocId); used {
+				if used, allocObj := o.PonPort.isAllocIdAllocated(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(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 751bf08..b9c2801 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, onu.SerialNumber)
 
 	// send a MibReset
 	err := onu.handleOmciRequest(makeOmciMessage(t, onu, makeOmciMibResetRequest(t)), stream)
diff --git a/internal/bbsim/devices/onu_state_machine_test.go b/internal/bbsim/devices/onu_state_machine_test.go
index 1c8470c..149ae2e 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, 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..86c4b4c 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[uint16]*AllocIDKey),
 			Olt:               &OltDevice{},
 		},
 		OmciResponseRate: 10,
diff --git a/internal/bbsim/devices/pon.go b/internal/bbsim/devices/pon.go
index 4db0e09..9005f2d 100644
--- a/internal/bbsim/devices/pon.go
+++ b/internal/bbsim/devices/pon.go
@@ -30,6 +30,11 @@
 	"module": "PON",
 })
 
+type AllocIDKey struct {
+	OnuSn   *openolt.SerialNumber
+	AllocID uint16
+}
+
 type PonPort struct {
 	// BBSIM Internals
 	ID            uint32
@@ -50,7 +55,7 @@
 	allocatedGemPortsLock sync.RWMutex
 	AllocatedOnuIds       map[uint32]*openolt.SerialNumber
 	allocatedOnuIdsLock   sync.RWMutex
-	AllocatedAllocIds     map[uint16]*openolt.SerialNumber
+	AllocatedAllocIds     map[uint16]*AllocIDKey // key is TCONT entity ID
 	allocatedAllocIdsLock sync.RWMutex
 }
 
@@ -65,7 +70,7 @@
 		Onus:              []*Onu{},
 		AllocatedGemPorts: make(map[uint16]*openolt.SerialNumber),
 		AllocatedOnuIds:   make(map[uint32]*openolt.SerialNumber),
-		AllocatedAllocIds: make(map[uint16]*openolt.SerialNumber),
+		AllocatedAllocIds: make(map[uint16]*AllocIDKey),
 	}
 
 	ponPort.InternalState = fsm.NewFSM(
@@ -291,20 +296,26 @@
 }
 
 // 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(entityID uint16, allocId uint16, onuSn *openolt.SerialNumber) {
 	p.allocatedAllocIdsLock.Lock()
 	defer p.allocatedAllocIdsLock.Unlock()
-	p.AllocatedAllocIds[allocId] = onuSn
+	p.AllocatedAllocIds[entityID] = &AllocIDKey{AllocID: allocId, OnuSn: onuSn}
 }
 
 // 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(entityID uint16) {
 	p.allocatedAllocIdsLock.Lock()
 	defer p.allocatedAllocIdsLock.Unlock()
-	for allocId, sn := range p.AllocatedAllocIds {
-		if sn == onuSn {
-			delete(p.AllocatedAllocIds, allocId)
+	delete(p.AllocatedAllocIds, entityID)
+}
+
+// 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 +323,16 @@
 func (p *PonPort) removeAllAllocIds() {
 	p.allocatedAllocIdsLock.Lock()
 	defer p.allocatedAllocIdsLock.Unlock()
-	p.AllocatedAllocIds = make(map[uint16]*openolt.SerialNumber)
+	p.AllocatedAllocIds = make(map[uint16]*AllocIDKey)
 }
 
 // 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(entityID uint16) (bool, *AllocIDKey) {
 	p.allocatedAllocIdsLock.RLock()
 	defer p.allocatedAllocIdsLock.RUnlock()
 
-	if _, ok := p.AllocatedAllocIds[allocId]; ok {
-		return true, p.AllocatedAllocIds[allocId]
+	if _, ok := p.AllocatedAllocIds[entityID]; ok {
+		return true, p.AllocatedAllocIds[entityID]
 	}
 	return false, nil
 }
diff --git a/internal/bbsim/devices/pon_test.go b/internal/bbsim/devices/pon_test.go
index 5467413..c0cd09b 100644
--- a/internal/bbsim/devices/pon_test.go
+++ b/internal/bbsim/devices/pon_test.go
@@ -129,23 +129,26 @@
 func Test_removeAllocId(t *testing.T) {
 
 	const (
-		allocId1 = 1024
-		allocId2 = 1025
+		entityID1 = 1024
+		entityID2 = 1025
+		allocId1  = 1024
+		allocId2  = 1025
 	)
 
 	pon := &PonPort{
-		AllocatedAllocIds: make(map[uint16]*openolt.SerialNumber),
+		AllocatedAllocIds: make(map[uint16]*AllocIDKey),
 	}
 
-	pon.AllocatedAllocIds[allocId1] = sn1
-	pon.AllocatedAllocIds[allocId2] = sn2
+	pon.AllocatedAllocIds[entityID1] = &AllocIDKey{sn1, allocId1}
+	pon.AllocatedAllocIds[entityID2] = &AllocIDKey{sn2, allocId2}
 
 	assert.Equal(t, len(pon.AllocatedAllocIds), 2)
 
-	pon.removeAllocId(sn1)
+	pon.removeAllocId(entityID1)
 
 	assert.Equal(t, len(pon.AllocatedAllocIds), 1)
-	assert.Nil(t, pon.AllocatedAllocIds[allocId1])
-	assert.Equal(t, pon.AllocatedAllocIds[allocId2], sn2)
+	assert.Nil(t, pon.AllocatedAllocIds[entityID1])
+	assert.NotNil(t, pon.AllocatedAllocIds[entityID2])
+	assert.Equal(t, pon.AllocatedAllocIds[entityID2].OnuSn, sn2)
 
 }