Skip to content

Commit b01f018

Browse files
committed
fix review6: validate OnnxTensor ThrMsgs
1 parent 1240151 commit b01f018

File tree

2 files changed

+32
-3
lines changed

2 files changed

+32
-3
lines changed

extra/modules/greenbids-real-time-data/src/main/java/org/prebid/server/hooks/modules/greenbids/real/time/data/model/data/GreenbidsUserAgent.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ private boolean isPC() {
5151
|| PC_OS_FAMILIES.contains(osFamily())
5252
|| ("Windows".equals(osFamily()) && "ME".equals(osMajor()))
5353
|| ("Mac OS X".equals(osFamily()) && !userAgent.contains("Silk"))
54-
|| userAgent.contains("Linux") && userAgent.contains("X11"))
54+
|| (userAgent.contains("Linux") && userAgent.contains("X11")))
5555
.orElse(false);
5656
}
5757

extra/modules/greenbids-real-time-data/src/main/java/org/prebid/server/hooks/modules/greenbids/real/time/data/model/predictor/FilterService.java

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor;
22

33
import ai.onnxruntime.OnnxTensor;
4+
import ai.onnxruntime.OnnxValue;
45
import ai.onnxruntime.OrtException;
56
import ai.onnxruntime.OrtSession;
67
import org.prebid.server.exception.PreBidException;
@@ -49,14 +50,42 @@ private Map<String, Map<String, Boolean>> processModelResults(
4950
List<ThrottlingMessage> throttlingMessages,
5051
Double threshold) {
5152

53+
validateThrottlingMessages(throttlingMessages);
54+
5255
return StreamSupport.stream(results.spliterator(), false)
53-
.filter(onnxItem -> Objects.equals(onnxItem.getKey(), "probabilities"))
56+
.filter(onnxItem -> {
57+
validateOnnxTensor(onnxItem);
58+
return Objects.equals(onnxItem.getKey(), "probabilities");
59+
})
5460
.map(onnxItem -> (OnnxTensor) onnxItem.getValue())
55-
.map(tensor -> extractAndProcessProbabilities(tensor, throttlingMessages, threshold))
61+
.map(tensor -> {
62+
validateTensorSize(tensor, throttlingMessages.size());
63+
return extractAndProcessProbabilities(tensor, throttlingMessages, threshold);
64+
})
5665
.flatMap(map -> map.entrySet().stream())
5766
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
5867
}
5968

69+
private void validateThrottlingMessages(List<ThrottlingMessage> throttlingMessages) {
70+
if (throttlingMessages == null || throttlingMessages.isEmpty()) {
71+
throw new PreBidException("throttlingMessages cannot be null or empty");
72+
}
73+
}
74+
75+
private void validateOnnxTensor(Map.Entry<String, OnnxValue> onnxItem) {
76+
if (!(onnxItem.getValue() instanceof OnnxTensor)) {
77+
throw new PreBidException("Expected OnnxTensor for 'probabilities', but found: "
78+
+ onnxItem.getValue().getClass().getName());
79+
}
80+
}
81+
82+
private void validateTensorSize(OnnxTensor tensor, int expectedSize) {
83+
final long[] tensorShape = tensor.getInfo().getShape();
84+
if (tensorShape.length == 0 || tensorShape[0] != expectedSize) {
85+
throw new PreBidException("Mismatch between tensor size and throttlingMessages size");
86+
}
87+
}
88+
6089
private Map<String, Map<String, Boolean>> extractAndProcessProbabilities(
6190
OnnxTensor tensor,
6291
List<ThrottlingMessage> throttlingMessages,

0 commit comments

Comments
 (0)