add other atomic instructions (#558)

Fixes #483
* Fast path for atomic operations on shared memory
---------

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,24 @@ 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) {
+        int32_t imm = (int32_t)b.op;
+        if (b.fetch)
+            imm |= INST_FETCH;
         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 = imm}};
     }
 
     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 (b.op) {
+        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: os_ << "x"; showfetch = false; break;
+        case Atomic::Op::CMPXCHG: os_ << "cx"; showfetch = false; break;
+        }
+        os_ << "= " << b.valreg;
+
+        if (showfetch && b.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,14 @@ 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, int> str_to_width = {
     {"8", 1},
     {"16", 2},
@@ -182,8 +191,13 @@ 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 "( fetch)?"))) {
+        Atomic::Op op = str_to_atomicop.at(m[5]);
+        return Atomic{
+            .op = op,
+            .fetch = m[7].matched || op == Atomic::Op::XCHG || op == Atomic::Op::CMPXCHG,
+            .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,24 @@ 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.
+/// The analysis just treats an atomic operation as a series of consecutive
+/// operations, and the atomicity itself is not significant.
+struct Atomic {
+    enum class Op {
+        ADD = 0x00,
+        OR = 0x40,
+        AND = 0x50,
+        XOR = 0xa0,
+        XCHG = 0xe0,    // Only valid with fetch=true.
+        CMPXCHG = 0xf0, // Only valid with fetch=true.
+    };
+
+    Op op;
+    bool fetch{};
     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 +378,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,21 @@ struct Unmarshaller {
         return {};
     }
 
+    auto getAtomicOp(size_t pc, ebpf_inst inst) -> Atomic::Op {
+        Atomic::Op op = (Atomic::Op)(inst.imm & ~INST_FETCH);
+        switch (op) {
+        case Atomic::Op::XCHG:
+        case Atomic::Op::CMPXCHG:
+            if ((inst.imm & INST_FETCH) == 0)
+                throw InvalidInstruction(pc, "unsupported immediate");
+        case Atomic::Op::ADD:
+        case Atomic::Op::OR:
+        case Atomic::Op::AND:
+        case Atomic::Op::XOR: return op;
+        }
+        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 +238,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 +252,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 +263,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 +295,14 @@ 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),
+                .fetch = (inst.imm & INST_FETCH) == INST_FETCH,
                 .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,69 @@ 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
+// Construct a Bin operation that does the main operation that a given Atomic operation does atomically.
+static Bin atomic_to_bin(const Atomic& a) {
+    Bin bin{
+        .dst = Reg{R11_ATOMIC_SCRATCH}, .v = a.valreg, .is64 = (a.access.width == sizeof(uint64_t)), .lddw = false};
+    switch (a.op) {
+    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:
+    case Atomic::Op::CMPXCHG: bin.op = Bin::Op::MOV; break;
+    default: throw std::exception();
+    }
+    return bin;
+}
+
+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;
+    }
+    if (m_inv.entail(type_is_not_stack(reg_pack(a.access.basereg)))) {
+        // Shared memory regions are volatile so we can just havoc
+        // any register that will be updated.
+        if (a.op == Atomic::Op::CMPXCHG)
+            havoc_register(m_inv, Reg{R0_RETURN_VALUE});
+        else if (a.fetch)
+            havoc_register(m_inv, a.valreg);
+        return;
+    }
+
+    // Fetch the current value into the R11 pseudo-register.
+    const Reg r11{R11_ATOMIC_SCRATCH};
+    (*this)(Mem{.access = a.access, .value = r11, .is_load = true});
+
+    // Compute the new value in R11.
+    (*this)(atomic_to_bin(a));
+
+    if (a.op == Atomic::Op::CMPXCHG) {
+        // For CMPXCHG, store the original value in r0.
+        (*this)(Mem{.access = a.access, .value = Reg{R0_RETURN_VALUE}, .is_load = true});
+
+        // 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 (a.fetch) {
+        // For other FETCH operations, store the original value in the src register.
+        (*this)(Mem{.access = a.access, .value = a.valreg, .is_load = true});
+    }
+
+    // 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.
+    (*this)(Mem{.access = a.access, .value = r11, .is_load = false});
+
+    // 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,16 +45,19 @@ 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
 
+    INST_FETCH = 0x1,
+
     INST_JA = 0x0,
     INST_CALL = 0x8,
     INST_EXIT = 0x9,
@@ -91,7 +94,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);