@@ -10,6 +10,8 @@ import scorex.util.Extensions._
1010import scorex .util .serialization .{Reader , Writer }
1111import scorex .util .{ModifierId , ScorexLogging , bytesToId , idToBytes }
1212
13+ import scala .collection .immutable
14+
1315case class ModifiersData (typeId : ModifierTypeId , modifiers : Map [ModifierId , Array [Byte ]])
1416
1517case class InvData (typeId : ModifierTypeId , ids : Seq [ModifierId ])
@@ -131,13 +133,15 @@ class ModifiersSpec(maxMessageSize: Int) extends MessageSpecV1[ModifiersData] wi
131133 override val messageCode : MessageCode = MessageCode
132134 override val messageName : String = MessageName
133135
136+ private val HeaderLength = 5 // msg type Id + modifiersCount
137+
134138 override def serialize (data : ModifiersData , w : Writer ): Unit = {
135139
136140 val typeId = data.typeId
137141 val modifiers = data.modifiers
138142 require(modifiers.nonEmpty, " empty modifiers list" )
139143
140- val (msgCount, msgSize) = modifiers.foldLeft((0 , 5 )) { case ((c, s), (id, modifier)) =>
144+ val (msgCount, msgSize) = modifiers.foldLeft((0 , HeaderLength )) { case ((c, s), (id, modifier)) =>
141145 val size = s + NodeViewModifier .ModifierIdSize + 4 + modifier.length
142146 val count = if (size <= maxMessageSize) c + 1 else c
143147 count -> size
@@ -154,21 +158,27 @@ class ModifiersSpec(maxMessageSize: Int) extends MessageSpecV1[ModifiersData] wi
154158 }
155159
156160 if (msgSize > maxMessageSize) {
157- log.warn(s " Message with modifiers ${modifiers.keySet} have size $msgSize exceeding limit $maxMessageSize. " +
161+ log.warn(s " Message with modifiers ${modifiers.keySet} has size $msgSize exceeding limit $maxMessageSize. " +
158162 s " Sending ${w.length() - start} bytes instead " )
159163 }
160164 }
161165
162166 override def parse (r : Reader ): ModifiersData = {
163- val typeId = ModifierTypeId @@ r.getByte()
164- val count = r.getUInt().toIntExact
165- val seq = (0 until count).map { _ =>
167+ val typeId = ModifierTypeId @@ r.getByte() // 1 byte
168+ val count = r.getUInt().toIntExact // 8 bytes
169+ val resMap = immutable.Map .newBuilder[ModifierId , Array [Byte ]]
170+ (0 until count).foldLeft(HeaderLength ) { case (msgSize, _) =>
166171 val id = bytesToId(r.getBytes(NodeViewModifier .ModifierIdSize ))
167172 val objBytesCnt = r.getUInt().toIntExact
173+ val newMsgSize = msgSize + NodeViewModifier .ModifierIdSize + objBytesCnt
174+ if (newMsgSize > maxMessageSize) {
175+ throw new Exception (" Too big message with modifiers, size: " + maxMessageSize)
176+ }
168177 val obj = r.getBytes(objBytesCnt)
169- id -> obj
178+ resMap += (id -> obj)
179+ newMsgSize
170180 }
171- ModifiersData (typeId, seq.toMap )
181+ ModifiersData (typeId, resMap.result() )
172182 }
173183}
174184
0 commit comments