Skip to content

Commit 202f6cf

Browse files
mamaria-kmamaria-k
committed
Add taint to Memory
1 parent 787d94d commit 202f6cf

File tree

3 files changed

+129
-20
lines changed

3 files changed

+129
-20
lines changed

‎lib/Core/Memory.cpp‎

Lines changed: 103 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -110,40 +110,46 @@ ObjectState::ObjectState(const MemoryObject *mo, const Array *array, KType *dt)
110110
: copyOnWriteOwner(0), object(mo), valueOS(ObjectStage(array, nullptr)),
111111
baseOS(ObjectStage(array->size, Expr::createPointer(0), false,
112112
Context::get().getPointerWidth())),
113+
taintOS(ObjectStage(array->size, Expr::createEmptyTaint(), false, Expr::TaintWidth)),
113114
lastUpdate(nullptr), size(array->size), dynamicType(dt), readOnly(false) {
114115
baseOS.initializeToZero();
116+
taintOS.initializeToZero();
115117
}
116118

117119
ObjectState::ObjectState(const MemoryObject *mo, KType *dt)
118120
: copyOnWriteOwner(0), object(mo),
119121
valueOS(ObjectStage(mo->getSizeExpr(), nullptr)),
120122
baseOS(ObjectStage(mo->getSizeExpr(), Expr::createPointer(0), false,
121123
Context::get().getPointerWidth())),
124+
taintOS(ObjectStage(mo->getSizeExpr(), Expr::createEmptyTaint(), false, Expr::TaintWidth)),
122125
lastUpdate(nullptr), size(mo->getSizeExpr()), dynamicType(dt),
123126
readOnly(false) {
124127
baseOS.initializeToZero();
128+
taintOS.initializeToZero();
125129
}
126130

127131
ObjectState::ObjectState(const ObjectState &os)
128132
: copyOnWriteOwner(0), object(os.object), valueOS(os.valueOS),
129-
baseOS(os.baseOS), lastUpdate(os.lastUpdate), size(os.size),
130-
dynamicType(os.dynamicType), readOnly(os.readOnly),
133+
baseOS(os.baseOS), taintOS(os.taintOS), lastUpdate(os.lastUpdate),
134+
size(os.size), dynamicType(os.dynamicType), readOnly(os.readOnly),
131135
wasWritten(os.wasWritten) {}
132136

133137
/***/
134138

135139
void ObjectState::initializeToZero() {
136140
valueOS.initializeToZero();
137141
baseOS.initializeToZero();
142+
taintOS.initializeToZero();
138143
}
139144

140145
ref<Expr> ObjectState::read8(unsigned offset) const {
141146
ref<Expr> val = valueOS.readWidth(offset);
142147
ref<Expr> base = baseOS.readWidth(offset);
143-
if (base->isZero()) {
148+
ref<Expr> taint = taintOS.readWidth(offset);
149+
if (base->isZero() && taint->isZero()) {
144150
return val;
145151
} else {
146-
return PointerExpr::create(base, val);
152+
return PointerExpr::create(base, val, taint);
147153
}
148154
}
149155

@@ -155,6 +161,10 @@ ref<Expr> ObjectState::readBase8(unsigned offset) const {
155161
return baseOS.readWidth(offset);
156162
}
157163

164+
ref<Expr> ObjectState::readTaint8(unsigned offset) const {
165+
return taintOS.readWidth(offset);
166+
}
167+
158168
ref<Expr> ObjectState::read8(ref<Expr> offset) const {
159169
assert(!isa<ConstantExpr>(offset) &&
160170
"constant offset passed to symbolic read8");
@@ -177,10 +187,11 @@ ref<Expr> ObjectState::read8(ref<Expr> offset) const {
177187
}
178188
ref<Expr> val = valueOS.readWidth(offset);
179189
ref<Expr> base = baseOS.readWidth(offset);
180-
if (base->isZero()) {
190+
ref<Expr> taint = taintOS.readWidth(offset);
191+
if (base->isZero() && taint->isZero()) {
181192
return val;
182193
} else {
183-
return PointerExpr::create(base, val);
194+
return PointerExpr::create(base, val, taint);
184195
}
185196
}
186197

@@ -230,21 +241,48 @@ ref<Expr> ObjectState::readBase8(ref<Expr> offset) const {
230241
return baseOS.readWidth(offset);
231242
}
232243

244+
ref<Expr> ObjectState::readTaint8(ref<Expr> offset) const {
245+
assert(!isa<ConstantExpr>(offset) &&
246+
"constant offset passed to symbolic read8");
247+
248+
if (object) {
249+
if (ref<ConstantExpr> sizeExpr =
250+
dyn_cast<ConstantExpr>(object->getSizeExpr())) {
251+
auto moSize = sizeExpr->getZExtValue();
252+
if (object && moSize > 4096) {
253+
std::string allocInfo;
254+
object->getAllocInfo(allocInfo);
255+
klee_warning_once(nullptr,
256+
"Symbolic memory access will send the following "
257+
"array of %lu bytes to "
258+
"the constraint solver -- large symbolic arrays may "
259+
"cause significant "
260+
"performance issues: %s",
261+
moSize, allocInfo.c_str());
262+
}
263+
}
264+
}
265+
return taintOS.readWidth(offset);
266+
}
267+
233268
void ObjectState::write8(unsigned offset, uint8_t value) {
234269
valueOS.writeWidth(offset, value);
235270
baseOS.writeWidth(offset,
236271
ConstantExpr::create(0, Context::get().getPointerWidth()));
272+
taintOS.writeWidth(offset, Expr::createEmptyTaint());
237273
}
238274

239275
void ObjectState::write8(unsigned offset, ref<Expr> value) {
240276
wasWritten = true;
241277
if (auto pointer = dyn_cast<PointerExpr>(value)) {
242278
valueOS.writeWidth(offset, pointer->getValue());
243279
baseOS.writeWidth(offset, pointer->getBase());
280+
taintOS.writeWidth(offset, pointer->getTaint());
244281
} else {
245282
valueOS.writeWidth(offset, value);
246283
baseOS.writeWidth(
247284
offset, ConstantExpr::create(0, Context::get().getPointerWidth()));
285+
taintOS.writeWidth(offset, Expr::createEmptyTaint());
248286
}
249287
}
250288

@@ -272,17 +310,22 @@ void ObjectState::write8(ref<Expr> offset, ref<Expr> value) {
272310
if (auto pointer = dyn_cast<PointerExpr>(value)) {
273311
valueOS.writeWidth(offset, pointer->getValue());
274312
baseOS.writeWidth(offset, pointer->getBase());
313+
taintOS.writeWidth(offset, pointer->getTaint());
275314
} else {
276315
valueOS.writeWidth(offset, value);
277316
baseOS.writeWidth(
278317
offset, ConstantExpr::create(0, Context::get().getPointerWidth()));
318+
taintOS.writeWidth(offset, Expr::createEmptyTaint());
279319
}
280320
}
281321

282322
void ObjectState::write(ref<const ObjectState> os) {
283323
wasWritten = true;
324+
284325
valueOS.write(os->valueOS);
285326
baseOS.write(os->baseOS);
327+
taintOS.write(os->taintOS);
328+
286329
lastUpdate = os->lastUpdate;
287330
}
288331

@@ -295,20 +338,22 @@ ref<Expr> ObjectState::read(ref<Expr> offset, Expr::Width width) const {
295338

296339
ref<Expr> val = readValue(offset, width);
297340
ref<Expr> base = readBase(offset, width);
298-
if (base->isZero()) {
341+
ref<Expr> taint = readTaint(offset, width);
342+
if (base->isZero() && taint->isZero()) {
299343
return val;
300344
} else {
301-
return PointerExpr::create(base, val);
345+
return PointerExpr::create(base, val, taint);
302346
}
303347
}
304348

305349
ref<Expr> ObjectState::read(unsigned offset, Expr::Width width) const {
306350
ref<Expr> val = readValue(offset, width);
307351
ref<Expr> base = readBase(offset, width);
308-
if (base->isZero()) {
352+
ref<Expr> taint = readTaint(offset, width);
353+
if (base->isZero() && taint->isZero()) {
309354
return val;
310355
} else {
311-
return PointerExpr::create(base, val);
356+
return PointerExpr::create(base, val, taint);
312357
}
313358
}
314359

@@ -368,10 +413,6 @@ ref<Expr> ObjectState::readBase(ref<Expr> offset, Expr::Width width) const {
368413
if (width == Expr::Bool)
369414
return ExtractExpr::create(readBase8(offset), 0, Expr::Bool);
370415

371-
// Treat bool specially, it is the only non-byte sized write we allow.
372-
if (width == Expr::Bool)
373-
return ExtractExpr::create(readBase8(offset), 0, Expr::Bool);
374-
375416
// Otherwise, follow the slow general case.
376417
unsigned NumBytes = width / 8;
377418
assert(width == NumBytes * 8 && "Invalid read size!");
@@ -416,6 +457,52 @@ ref<Expr> ObjectState::readBase(unsigned offset, Expr::Width width) const {
416457
return Res;
417458
}
418459

460+
ref<Expr> ObjectState::readTaint(ref<Expr> offset, Expr::Width width) const {
461+
// Truncate offset to 32-bits.
462+
offset = ZExtExpr::create(offset, Expr::Int32);
463+
464+
// Check for reads at constant offsets.
465+
if (ConstantExpr *CE = dyn_cast<ConstantExpr>(offset))
466+
return readTaint(CE->getZExtValue(32), width);
467+
468+
// Treat bool specially, it is the only non-byte sized write we allow.
469+
if (width == Expr::Bool)
470+
return ExtractExpr::create(readTaint8(offset), 0, Expr::Bool);
471+
472+
// Otherwise, follow the slow general case.
473+
unsigned NumBytes = width / 8;
474+
assert(width == NumBytes * 8 && "Invalid read size!");
475+
ref<Expr> Res(0);
476+
ref<Expr> null = Expr::createPointer(0);
477+
for (unsigned i = 0; i != NumBytes; ++i) {
478+
unsigned idx = Context::get().isLittleEndian() ? i : (NumBytes - i - 1);
479+
ref<Expr> Byte = readTaint8(
480+
AddExpr::create(offset, ConstantExpr::create(idx, Expr::Int32)));
481+
Res = i ? OrExpr::create(Byte, Res) : Byte;
482+
}
483+
484+
return Res;
485+
}
486+
487+
ref<Expr> ObjectState::readTaint(unsigned offset, Expr::Width width) const {
488+
// Treat bool specially, it is the only non-byte sized write we allow.
489+
if (width == Expr::Bool)
490+
return ExtractExpr::create(readTaint8(offset), 0, Expr::Bool);
491+
492+
// Otherwise, follow the slow general case.
493+
unsigned NumBytes = width / 8;
494+
assert(width == NumBytes * 8 && "Invalid width for read size!");
495+
ref<Expr> Res(0);
496+
ref<Expr> null = Expr::createPointer(0);
497+
for (unsigned i = 0; i != NumBytes; ++i) {
498+
unsigned idx = Context::get().isLittleEndian() ? i : (NumBytes - i - 1);
499+
ref<Expr> Byte = readTaint8(offset + idx);
500+
Res = i ? OrExpr::create(Byte, Res) : Byte;
501+
}
502+
503+
return Res;
504+
}
505+
419506
void ObjectState::write(ref<Expr> offset, ref<Expr> value) {
420507
// Truncate offset to 32-bits.
421508
offset = ZExtExpr::create(offset, Expr::Int32);
@@ -516,6 +603,8 @@ void ObjectState::print() const {
516603
valueOS.print();
517604
llvm::errs() << "\tOffset ObjectStage:\n";
518605
baseOS.print();
606+
llvm::errs() << "\tTaint ObjectStage:\n";
607+
taintOS.print();
519608
}
520609

521610
KType *ObjectState::getDynamicType() const { return dynamicType; }

‎lib/Core/Memory.h‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,11 @@ class ObjectStage {
267267
}
268268
void initializeToZero();
269269

270+
void reset(ref<Expr> newDefault) {
271+
knownSymbolics->reset(std::move(newDefault));
272+
unflushedMask->reset(false);
273+
}
274+
270275
private:
271276
const UpdateList &getUpdates() const;
272277

@@ -291,6 +296,7 @@ class ObjectState {
291296

292297
ObjectStage valueOS;
293298
ObjectStage baseOS;
299+
ObjectStage taintOS;
294300

295301
ref<UpdateNode> lastUpdate;
296302

@@ -318,7 +324,8 @@ class ObjectState {
318324
void initializeToZero();
319325

320326
size_t getSparseStorageEntries() {
321-
return valueOS.getSparseStorageEntries() + baseOS.getSparseStorageEntries();
327+
return valueOS.getSparseStorageEntries() +
328+
baseOS.getSparseStorageEntries() + taintOS.getSparseStorageEntries();
322329
}
323330

324331
void swapObjectHack(MemoryObject *mo) { object = mo; }
@@ -328,10 +335,13 @@ class ObjectState {
328335
ref<Expr> read8(unsigned offset) const;
329336
ref<Expr> readValue(ref<Expr> offset, Expr::Width width) const;
330337
ref<Expr> readBase(ref<Expr> offset, Expr::Width width) const;
338+
ref<Expr> readTaint(ref<Expr> offset, Expr::Width width) const;
331339
ref<Expr> readValue(unsigned offset, Expr::Width width) const;
332340
ref<Expr> readBase(unsigned offset, Expr::Width width) const;
341+
ref<Expr> readTaint(unsigned offset, Expr::Width width) const;
333342
ref<Expr> readValue8(unsigned offset) const;
334343
ref<Expr> readBase8(unsigned offset) const;
344+
ref<Expr> readTaint8(unsigned offset) const;
335345

336346
void write(unsigned offset, ref<Expr> value);
337347
void write(ref<Expr> offset, ref<Expr> value);
@@ -347,10 +357,13 @@ class ObjectState {
347357

348358
KType *getDynamicType() const;
349359

360+
void resetTaint(ref<Expr> newTaint) { taintOS.reset(std::move(newTaint)); }
361+
350362
private:
351363
ref<Expr> read8(ref<Expr> offset) const;
352364
ref<Expr> readValue8(ref<Expr> offset) const;
353365
ref<Expr> readBase8(ref<Expr> offset) const;
366+
ref<Expr> readTaint8(ref<Expr> offset) const;
354367
void write8(unsigned offset, ref<Expr> value);
355368
void write8(ref<Expr> offset, ref<Expr> value);
356369
};

‎lib/Core/SpecialFunctionHandler.cpp‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1253,6 +1253,10 @@ void SpecialFunctionHandler::handleCheckTaintSource(
12531253
"klee_check_taint_source(void*, size_t)");
12541254
return;
12551255
}
1256+
1257+
// uint64_t taintSource = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
1258+
// executor.executeCheckTaintSource(state, target,
1259+
// executor.makePointer(arguments[0]), taintSource);
12561260
}
12571261

12581262
void SpecialFunctionHandler::handleGetTaintRule(
@@ -1265,9 +1269,13 @@ void SpecialFunctionHandler::handleGetTaintRule(
12651269
return;
12661270
}
12671271

1268-
// TODO: now mock
1272+
// // TODO: now mock
12691273
ref<Expr> result = ConstantExpr::create(1, Expr::Int64);
12701274
executor.bindLocal(target, state, result);
1275+
1276+
// uint64_t taintSink = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
1277+
// executor.executeGetTaintRule(state, target,
1278+
// executor.makePointer(arguments[0]), taintSink);
12711279
}
12721280

12731281
void SpecialFunctionHandler::handleTaintHit(klee::ExecutionState &state,
@@ -1279,8 +1287,7 @@ void SpecialFunctionHandler::handleTaintHit(klee::ExecutionState &state,
12791287
return;
12801288
}
12811289

1282-
// printf("klee_taint_hit for rule: %s\n", arguments[0]->toString().c_str());
1283-
1284-
executor.terminateStateOnTargetTaintError(
1285-
state, dyn_cast<ConstantExpr>(arguments[0])->getZExtValue());
1290+
uint64_t ruleId = dyn_cast<ConstantExpr>(arguments[0])->getZExtValue();
1291+
// printf("klee_taint_hit for rule: %zu\n", ruleId);
1292+
executor.terminateStateOnTargetTaintError(state, ruleId);
12861293
}

0 commit comments

Comments
 (0)