Add other atomic instructions

Fixes #483

Signed-off-by: Dave Thaler <[email protected]>

src/asm_cfg.cpp

@@ -178,7 +178,7 @@ static std::string instype(Instruction ins) {
         return "callx";
     } else if (std::holds_alternative<Mem>(ins)) {
         return std::get<Mem>(ins).is_load ? "load" : "store";
-    } else if (std::holds_alternative<LockAdd>(ins)) {
+    } else if (std::holds_alternative<Atomic>(ins)) {
         return "load_store";
     } else if (std::holds_alternative<Packet>(ins)) {
         return "packet_access";

src/asm_marshal.cpp

@@ -219,7 +219,7 @@ struct MarshalVisitor {
     vector<ebpf_inst> operator()(Mem const& b) {
         Deref access = b.access;
         ebpf_inst res{
-            .opcode = static_cast<uint8_t>((INST_MEM << 5) | width_to_opcode(access.width)),
+            .opcode = static_cast<uint8_t>(INST_MODE_MEM | width_to_opcode(access.width)),
             .dst = 0,
             .src = 0,
             .offset = static_cast<int16_t>(access.offset),
@@ -253,21 +253,21 @@ struct MarshalVisitor {
             .imm = static_cast<int32_t>(b.offset),
         };
         if (b.regoffset) {
-            res.opcode |= (INST_IND << 5);
+            res.opcode |= INST_MODE_IND;
             res.src = b.regoffset->v;
         } else {
-            res.opcode |= (INST_ABS << 5);
+            res.opcode |= INST_MODE_ABS;
         }
         return {res};
     }
 
-    vector<ebpf_inst> operator()(LockAdd const& b) {
+    vector<ebpf_inst> operator()(Atomic const& b) {
         return {ebpf_inst{
-            .opcode = static_cast<uint8_t>(INST_CLS_ST | 0x1 | (INST_XADD << 5) | width_to_opcode(b.access.width)),
+            .opcode = static_cast<uint8_t>(INST_CLS_STX | INST_MODE_ATOMIC | width_to_opcode(b.access.width)),
             .dst = b.access.basereg.v,
             .src = b.valreg.v,
             .offset = static_cast<int16_t>(b.access.offset),
-            .imm = 0}};
+            .imm = (int32_t)b.op}};
     }
 
     vector<ebpf_inst> operator()(IncrementLoopCounter const& ins) {

src/asm_ostream.cpp

@@ -328,10 +328,24 @@ struct InstructionPrinterVisitor {
         }
     }
 
-    void operator()(LockAdd const& b) {
+    void operator()(Atomic const& b) {
         os_ << "lock ";
         print(b.access);
-        os_ << " += " << b.valreg;
+        os_ << " ";
+        bool showfetch = true;
+        switch ((Atomic::Op)((uint32_t)b.op & (uint32_t)Atomic::Op::BASE_MASK)) {
+        case Atomic::Op::ADD: os_ << "+"; break;
+        case Atomic::Op::OR : os_ << "|"; break;
+        case Atomic::Op::AND: os_ << "&"; break;
+        case Atomic::Op::XOR: os_ << "^"; break;
+        case Atomic::Op::XCHG_BASE: os_ << "x"; showfetch = false; break;
+        case Atomic::Op::CMPXCHG_BASE: os_ << "cx"; showfetch = false; break;
+        }
+        os_ << "= " << b.valreg;
+
+        if (showfetch && ((uint32_t)b.op & (uint32_t)Atomic::Op::FETCH)) {
+            os_ << " fetch";
+        }
     }
 
     void operator()(Assume const& b) {

src/asm_ostream.hpp

@@ -59,7 +59,7 @@ inline std::ostream& operator<<(std::ostream& os, Exit const& a) { return os <<
 inline std::ostream& operator<<(std::ostream& os, Jmp const& a) { return os << (Instruction)a; }
 inline std::ostream& operator<<(std::ostream& os, Packet const& a) { return os << (Instruction)a; }
 inline std::ostream& operator<<(std::ostream& os, Mem const& a) { return os << (Instruction)a; }
-inline std::ostream& operator<<(std::ostream& os, LockAdd const& a) { return os << (Instruction)a; }
+inline std::ostream& operator<<(std::ostream& os, Atomic const& a) { return os << (Instruction)a; }
 inline std::ostream& operator<<(std::ostream& os, Assume const& a) { return os << (Instruction)a; }
 inline std::ostream& operator<<(std::ostream& os, Assert const& a) { return os << (Instruction)a; }
 inline std::ostream& operator<<(std::ostream& os, IncrementLoopCounter const& a) { return os << (Instruction)a; }

src/asm_parse.cpp

@@ -34,6 +34,7 @@ using crab::linear_expression_t;
 #define ASSIGN R"_(\s*=\s*)_"
 #define LONGLONG R"_(\s*(ll|)\s*)_"
 #define UNOP R"_((-|be16|be32|be64|le16|le32|le64|swap16|swap32|swap64))_"
+#define ATOMICOP R"_((\+|\||&|\^|x|cx)=)_"
 
 #define PLUSMINUS R"_((\s*[+-])\s*)_"
 #define LPAREN R"_(\s*\(\s*)_"
@@ -81,6 +82,20 @@ static const std::map<std::string, Condition::Op> str_to_cmpop = {
     {"s<", Condition::Op::SLT}, {"s<=", Condition::Op::SLE}, {"s>", Condition::Op::SGT},  {"s>=", Condition::Op::SGE},
 };
 
+static const std::map<std::string, Atomic::Op> str_to_atomicop = {
+    {"+", Atomic::Op::ADD},
+    {"|", Atomic::Op::OR},
+    {"&", Atomic::Op::AND},
+    {"^", Atomic::Op::XOR},
+    {"x", Atomic::Op::XCHG},
+    {"cx", Atomic::Op::CMPXCHG}};
+
+static const std::map<std::string, Atomic::Op> str_to_atomicfetchop = {
+    {"+", Atomic::Op::ADD_FETCH},
+    {"|", Atomic::Op::OR_FETCH},
+    {"&", Atomic::Op::AND_FETCH},
+    {"^", Atomic::Op::XOR_FETCH}};
+
 static const std::map<std::string, int> str_to_width = {
     {"8", 1},
     {"16", 2},
@@ -182,8 +197,17 @@ Instruction parse_instruction(const std::string& line, const std::map<std::strin
             .is_load = false,
         };
     }
-    if (regex_match(text, m, regex("lock " DEREF PAREN(REG PLUSMINUS IMM) " [+]= " REG))) {
-        return LockAdd{.access = deref(m[1], m[2], m[3], m[4]), .valreg = reg(m[5])};
+    if (regex_match(text, m, regex("lock " DEREF PAREN(REG PLUSMINUS IMM) " " ATOMICOP " " REG))) {
+        return Atomic{
+            .op = str_to_atomicop.at(m[5]),
+            .access = deref(m[1], m[2], m[3], m[4]),
+            .valreg = reg(m[6])};
+    }
+    if (regex_match(text, m, regex("lock " DEREF PAREN(REG PLUSMINUS IMM) " " ATOMICOP " " REG " fetch"))) {
+        return Atomic{
+            .op = str_to_atomicfetchop.at(m[5]),
+            .access = deref(m[1], m[2], m[3], m[4]),
+            .valreg = reg(m[6])};
     }
     if (regex_match(text, m, regex("r0 = " DEREF "skb\\[(.*)\\]"))) {
         auto width = str_to_width.at(m[1]);

src/asm_syntax.hpp

@@ -235,11 +235,31 @@ struct Packet {
     constexpr bool operator==(const Packet&) const = default;
 };
 
-/// Special instruction for incrementing values inside shared memory.
-struct LockAdd {
+/// Special instruction for atomically updating values inside shared memory.
+struct Atomic {
+    enum class Op {
+        FETCH = 0x01,
+        BASE_MASK = 0xf0, // Not including the FETCH flag.
+
+        ADD = 0x00,
+        ADD_FETCH = (ADD | FETCH),
+        OR = 0x40,
+        OR_FETCH = (OR | FETCH),
+        AND = 0x50,
+        AND_FETCH = (AND | FETCH),
+        XOR = 0xa0,
+        XOR_FETCH = (XOR | FETCH),
+        XCHG_BASE = 0xe0,    // Not valid by itself
+        CMPXCHG_BASE = 0xf0, // Not valid by itself
+
+        XCHG = (XCHG_BASE | FETCH),
+        CMPXCHG = (CMPXCHG_BASE | FETCH),
+    };
+
+    Op op;
     Deref access;
     Reg valreg;
-    constexpr bool operator==(const LockAdd&) const = default;
+    constexpr bool operator==(const Atomic&) const = default;
 };
 
 /// Not an instruction, just used for failure cases.
@@ -365,7 +385,7 @@ struct IncrementLoopCounter {
     constexpr bool operator==(const IncrementLoopCounter&) const = default;
 };
 
-using Instruction = std::variant<Undefined, Bin, Un, LoadMapFd, Call, Callx, Exit, Jmp, Mem, Packet, LockAdd, Assume, Assert, IncrementLoopCounter>;
+using Instruction = std::variant<Undefined, Bin, Un, LoadMapFd, Call, Callx, Exit, Jmp, Mem, Packet, Atomic, Assume, Assert, IncrementLoopCounter>;
 
 using LabeledInstruction = std::tuple<label_t, Instruction, std::optional<btf_line_info_t>>;
 using InstructionSeq = std::vector<LabeledInstruction>;

src/asm_unmarshal.cpp

@@ -177,6 +177,22 @@ struct Unmarshaller {
         return {};
     }
 
+    auto getAtomicOp(size_t pc, ebpf_inst inst) -> Atomic::Op {
+        switch ((Atomic::Op)inst.imm) {
+        case Atomic::Op::ADD:
+        case Atomic::Op::ADD_FETCH:
+        case Atomic::Op::OR:
+        case Atomic::Op::OR_FETCH:
+        case Atomic::Op::AND:
+        case Atomic::Op::AND_FETCH:
+        case Atomic::Op::XOR:
+        case Atomic::Op::XOR_FETCH:
+        case Atomic::Op::XCHG:
+        case Atomic::Op::CMPXCHG: return (Atomic::Op)inst.imm;
+        }
+        throw InvalidInstruction(pc, "unsupported immediate");
+    }
+
     uint64_t sign_extend(int32_t imm) { return (uint64_t)(int64_t)imm; }
 
     uint64_t zero_extend(int32_t imm) { return (uint64_t)(uint32_t)imm; }
@@ -223,9 +239,10 @@ struct Unmarshaller {
 
         int width = getMemWidth(inst.opcode);
         bool isLD = (inst.opcode & INST_CLS_MASK) == INST_CLS_LD;
-        switch ((inst.opcode & INST_MODE_MASK) >> 5) {
-        case 0: throw InvalidInstruction(pc, inst.opcode);
-        case INST_ABS:
+        switch (inst.opcode & INST_MODE_MASK) {
+        case INST_MODE_IMM:
+            throw InvalidInstruction(pc, inst.opcode);
+        case INST_MODE_ABS:
             if (!info.platform->legacy || !isLD || (width == 8))
                 throw InvalidInstruction(pc, inst.opcode);
             if (inst.dst != 0)
@@ -236,7 +253,7 @@ struct Unmarshaller {
                 throw InvalidInstruction(pc, make_opcode_message("nonzero offset for", inst.opcode));
             return Packet{.width = width, .offset = inst.imm, .regoffset = {}};
 
-        case INST_IND:
+        case INST_MODE_IND:
             if (!info.platform->legacy || !isLD || (width == 8))
                 throw InvalidInstruction(pc, inst.opcode);
             if (inst.dst != 0)
@@ -247,7 +264,7 @@ struct Unmarshaller {
                 throw InvalidInstruction(pc, make_opcode_message("nonzero offset for", inst.opcode));
             return Packet{.width = width, .offset = inst.imm, .regoffset = Reg{inst.src}};
 
-        case INST_MEM: {
+        case INST_MODE_MEM: {
             if (isLD)
                 throw InvalidInstruction(pc, inst.opcode);
             bool isLoad = getMemIsLoad(inst.opcode);
@@ -279,14 +296,13 @@ struct Unmarshaller {
             return res;
         }
 
-        case INST_XADD:
+        case INST_MODE_ATOMIC:
             if (((inst.opcode & INST_CLS_MASK) != INST_CLS_STX) ||
                 ((inst.opcode & INST_SIZE_MASK) != INST_SIZE_W &&
                  (inst.opcode & INST_SIZE_MASK) != INST_SIZE_DW))
                 throw InvalidInstruction(pc, inst.opcode);
-            if (inst.imm != 0)
-                throw InvalidInstruction(pc, "unsupported immediate");
-            return LockAdd{
+            return Atomic{
+                .op = getAtomicOp(pc, inst),
                 .access =
                     Deref{
                         .width = width,

src/assertions.cpp

@@ -170,11 +170,17 @@ class AssertExtractor {
         return res;
     }
 
-    vector<Assert> operator()(LockAdd ins) const {
+    vector<Assert> operator()(Atomic ins) const {
         vector<Assert> res;
-        res.emplace_back(TypeConstraint{ins.access.basereg, TypeGroup::shared});
+        res.emplace_back(TypeConstraint{ins.access.basereg, TypeGroup::pointer});
         res.emplace_back(ValidAccess{ins.access.basereg, ins.access.offset,
                                      Imm{static_cast<uint32_t>(ins.access.width)}, false});
+        if (ins.op == Atomic::Op::CMPXCHG) {
+            // The memory contents pointed to by ins.access will be compared
+            // against the value of the ins.valreg register.  Only numbers are
+            // supported.
+            res.emplace_back(TypeConstraint{ins.valreg, TypeGroup::number});
+        }
         return res;
     }
 

src/crab/ebpf_domain.cpp

@@ -2154,8 +2154,61 @@ void ebpf_domain_t::do_mem_store(const Mem& b, Type val_type, SValue val_svalue,
     });
 }
 
-void ebpf_domain_t::operator()(const LockAdd& a) {
-    // nothing to do here
+void ebpf_domain_t::operator()(const Atomic& a) {
+    if (m_inv.is_bottom())
+        return;
+    if (!m_inv.entail(type_is_pointer(reg_pack(a.access.basereg))) ||
+        !m_inv.entail(type_is_number(reg_pack(a.valreg)))) {
+        return;
+    }
+
+    // Fetch the current value into the R11 pseudo-register.
+    Reg r11{R11_ATOMIC_SCRATCH};
+    Mem mem = {.access = a.access, .value = r11, .is_load = true};
+    (*this)(mem);
+
+    // Compute the new value in R11.
+    Bin bin {.dst = r11, .v = a.valreg, .is64 = (a.access.width == sizeof(uint64_t)), .lddw = false};
+    switch ((Atomic::Op)((uint32_t)a.op & (uint32_t)Atomic::Op::BASE_MASK)) {
+    case Atomic::Op::ADD: bin.op = Bin::Op::ADD; break;
+    case Atomic::Op::OR: bin.op = Bin::Op::OR; break;
+    case Atomic::Op::AND: bin.op = Bin::Op::AND; break;
+    case Atomic::Op::XOR: bin.op = Bin::Op::XOR; break;
+    case Atomic::Op::XCHG_BASE:
+    case Atomic::Op::CMPXCHG_BASE: bin.op = Bin::Op::MOV; break;
+    default: throw std::exception();
+    }
+    (*this)(bin);
+
+    if (a.op == Atomic::Op::CMPXCHG) {
+        // For CMPXCHG, store the original value in r0.
+        Reg r0{R0_RETURN_VALUE};
+        Mem mem = {.access = a.access, .value = r0, .is_load = true};
+        (*this)(mem);
+
+        // For the destination, there are 3 possibilities:
+        // 1) dst.value == r0.value : set R11 to valreg
+        // 2) dst.value != r0.value : don't modify R11
+        // 3) dst.value may or may not == r0.value : set R11 to the union of R11 and valreg
+        // For now we just havoc the value of R11.
+        havoc_register(m_inv, r11);
+    } else if ((uint32_t)a.op & (uint32_t)Atomic::Op::FETCH) {
+        // For other FETCH operations, store the original value in the src register.
+        mem.value = a.valreg;
+        mem.is_load = true;
+        (*this)(mem);
+    }
+
+    // Store the new value back in the original shared memory location.
+    // Note that do_mem_store() currently doesn't track shared memory values,
+    // but stack memory values are tracked and are legal here.
+    mem.value = r11;
+    mem.is_load = false;
+    (*this)(mem);
+
+    // Clear the R11 pseudo-register.
+    havoc_register(m_inv, r11);
+    type_inv.havoc_type(m_inv, r11);
 }
 
 void ebpf_domain_t::operator()(const Call& call) {

src/crab/ebpf_domain.hpp

@@ -67,7 +67,7 @@ class ebpf_domain_t final {
     void operator()(const FuncConstraint&);
     void operator()(const Jmp&);
     void operator()(const LoadMapFd&);
-    void operator()(const LockAdd&);
+    void operator()(const Atomic&);
     void operator()(const Mem&);
     void operator()(const ValidDivisor&);
     void operator()(const Packet&);

src/ebpf_vm_isa.hpp

@@ -45,13 +45,14 @@ enum {
 
     INST_MODE_MASK = 0xe0,
 
-    INST_ABS = 1, // Deprecated
-    INST_IND = 2, // Deprecated
-    INST_MEM = 3,
-    INST_LEN = 4,
-    INST_MSH = 5,
-    INST_XADD = 6,
-    INST_MEM_UNUSED = 7,
+    INST_MODE_IMM = 0x00, // 64-bit immediate instructions
+    INST_MODE_ABS = 0x20, // legacy BPF packet access (absolute)
+    INST_MODE_IND = 0x40, // legacy BPF packet access (indirect)
+    INST_MODE_MEM = 0x60, // regular load and store operations
+    INST_MODE_MEMSX = 0x80, // sign-extension load operations
+    INST_MODE_UNUSED1 = 0xa0,
+    INST_MODE_ATOMIC = 0xc0, // atomic operations
+    INST_MODE_UNUSED2 = 0xe0,
 
     INST_OP_LDDW_IMM = (INST_CLS_LD | INST_SRC_IMM | INST_SIZE_DW), // Special
 
@@ -91,7 +92,8 @@ enum {
     R7 = 7,
     R8 = 8,
     R9 = 9,
-    R10_STACK_POINTER = 10
+    R10_STACK_POINTER = 10,
+    R11_ATOMIC_SCRATCH = 11, // Pseudo-register used internally for atomic instructions.
 };
 
 int opcode_to_width(uint8_t opcode);