Skip to content

Commit

Permalink
sbox detection (wip, not tested)
Browse files Browse the repository at this point in the history
  • Loading branch information
SJulianS committed May 23, 2024
1 parent 7f603ec commit dd928f8
Show file tree
Hide file tree
Showing 4 changed files with 304 additions and 8 deletions.
2 changes: 1 addition & 1 deletion plugins/hawkeye/include/hawkeye/sbox_lookup.h
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ namespace hal
{
class StateCandidate;

// TODO extract_sbox_components();
Result<std::vector<std::pair<std::set<Gate*>, std::set<Gate*>>>> locate_sboxes(const StateCandidate* candidate);

Result<std::string>
identify_sbox(const StateCandidate* candidate, const std::set<Gate*>& component, const std::set<Gate*>& input_gates, const std::set<Gate*>& output_gates, const SBoxDatabase& db);
Expand Down
9 changes: 9 additions & 0 deletions plugins/hawkeye/include/hawkeye/state_candidate.h
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,9 @@ namespace hal
*/
const std::set<Net*>& get_state_outputs() const;

const std::map<Gate*, std::set<Gate*>>& get_input_ffs_of_gate() const;
const std::map<u32, std::set<Gate*>>& get_longest_distance_to_gate() const;

private:
/**
* The netlist to which the candidate belongs.
Expand Down Expand Up @@ -168,6 +171,12 @@ namespace hal
* The state outputs from the combinational logic computing the next state.
*/
std::set<Net*> m_state_outputs;

std::map<Gate*, std::set<Gate*>> m_gates_reached_by_input_ff;

std::map<Gate*, std::set<Gate*>> m_input_ffs_of_gate;

std::map<u32, std::set<Gate*>> m_longest_distance_to_gate;
};
} // namespace hawkeye
} // namespace hal
195 changes: 195 additions & 0 deletions plugins/hawkeye/src/sbox_lookup.cpp
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
#include "hawkeye/sbox_lookup.h"

#include "graph_algorithm/algorithms/components.h"
#include "graph_algorithm/algorithms/subgraph.h"
#include "hal_core/netlist/boolean_function.h"
#include "hal_core/netlist/decorators/boolean_function_net_decorator.h"
#include "hal_core/netlist/decorators/subgraph_netlist_decorator.h"
Expand All @@ -12,6 +14,199 @@ namespace hal
{
namespace hawkeye
{
Result<std::vector<std::pair<std::set<Gate*>, std::set<Gate*>>>> locate_sboxes(const StateCandidate* candidate)
{
const auto* nl = candidate->get_netlist();
const auto* graph = candidate->get_graph();

// get initial set of components
auto comp_res = graph_algorithm::get_connected_components(graph, false).map<std::vector<std::vector<Gate*>>>([graph](const auto& comps) -> Result<std::vector<std::vector<Gate*>>> {
std::vector<std::vector<Gate*>> res;
for (const auto& c : comps)
{
if (const auto gates_res = graph->get_gates_from_vertices(c); gates_res.is_ok())
{
res.push_back(gates_res.get());
}
else
{
return ERR(gates_res.get_error());
}
}
return OK(res);
});
if (comp_res.is_error())
{
return ERR(comp_res.get_error());
}
auto components = comp_res.get();

const auto& state_input_reg = candidate->get_input_reg();
const auto& state_output_reg = candidate->get_output_reg();

std::vector<std::pair<std::set<Gate*>, std::set<Gate*>>> res;

for (const auto& component : components)
{
// gather FFs of the component that are also part of the state input reg
std::set<Gate*> component_input_ffs;
std::set_intersection(component.begin(), component.end(), state_input_reg.begin(), state_input_reg.end(), std::inserter(component_input_ffs, component_input_ffs.begin()));

u32 number_input_ffs = component_input_ffs.size();
if (number_input_ffs < 3)
{
// too small for S-box
continue;
}
else if (number_input_ffs <= 8)
{
// assume to have found a single S-box
std::set<Gate*> sbox_output_gates;

for (auto* cand_gate : component)
{
// skip FFs
if (cand_gate->get_type()->has_property(GateTypeProperty::ff))
{
continue;
}

// output gates are all combinational gates that have no other successors but the state output reg
const auto suc_gates = cand_gate->get_unique_successors();
if (std::none_of(suc_gates.begin(), suc_gates.end(), [&state_output_reg](Gate* g) { return state_output_reg.find(g) == state_output_reg.end(); }))
{
sbox_output_gates.insert(cand_gate);
}
}

// create S-box candidate if input size equals output size
if (sbox_output_gates.size() == number_input_ffs)
{
res.push_back(std::make_pair(std::move(component_input_ffs), std::move(sbox_output_gates)));
}
continue;
}

// try to split component in smaller sub-components (assuming component is combination of S-box and linear layer)
std::set<Gate*> current_subset = component_input_ffs;
std::vector<std::vector<std::set<Gate*>>> input_groupings;

// abuse that map keys are sorted, hence rbegin() will return max distance in map
const auto& longest_dist_to_gates = candidate->get_longest_distance_to_gate();
for (u32 step = 1; step < longest_dist_to_gates.rbegin()->first; step++)
{
if (const auto dist_it = longest_dist_to_gates.find(step); dist_it != longest_dist_to_gates.end())
{
const auto& new_gates = std::get<1>(*dist_it);
current_subset.insert(new_gates.begin(), new_gates.end());
}
else
{
break;
}

// generate subgraph of new sub-component
auto subgraph_res = graph_algorithm::get_subgraph(graph, current_subset);
if (subgraph_res.is_error())
{
return ERR(subgraph_res.get_error());
}
auto subgraph = std::move(subgraph_res.get());

auto comp_res = graph_algorithm::get_connected_components(subgraph.get(), false);
if (comp_res.is_error())
{
return ERR(comp_res.get_error());
}

// determine input groups feeding into distinct sub-circuits
std::set<u32> lens;
std::vector<std::vector<Gate*>> subcomponents;
std::vector<std::set<Gate*>> input_groups;
for (const auto& comp : comp_res.get())
{
auto gates_res = graph->get_gates_from_vertices(comp);
if (gates_res.is_error())
{
return ERR(gates_res.get_error());
}
auto comp_gates = gates_res.get();

// only consider sub-components connected to at least one input FF
if (std::any_of(comp_gates.begin(), comp_gates.end(), [&component_input_ffs](Gate* g) { return component_input_ffs.find(g) != component_input_ffs.end(); }))
{
std::set<Gate*> input_group;
std::set_intersection(comp_gates.begin(), comp_gates.end(), state_input_reg.begin(), state_input_reg.end(), std::inserter(input_group, input_group.begin()));
lens.insert(input_group.size());
input_groups.push_back(std::move(input_group));
subcomponents.push_back(std::move(comp_gates));
}
}

// all input groups should have same size and comprise more than one input
if (lens.size() == 1 && input_groups.at(0).size() > 1 && input_groups.size() > 1)
{
input_groupings.push_back(input_groups);
}
}

const auto& input_ffs_of_gate = candidate->get_input_ffs_of_gate();

std::vector<std::pair<std::set<Gate*>, std::set<Gate*>>> sboxes_input_output_gates;
for (const auto& input_groups : input_groupings)
{
for (const auto& input_group : input_groups)
{
std::set<Gate*> output_group;
for (auto auto* comp_gate : component)
{
// disregard output FFs
if (state_output_reg.find(comp_gate) != state_output_reg.end())
{
continue;
}

// disregard gates that only depend on a single input FF
if (input_ffs_of_gate.at(comp_gate).size() <= 1)
{
continue;
}

// disregard gates that do not only depend on the input FFs of the sub-component
if (!std::includes(input_group.begin(), input_group.end(), input_ffs_of_gate.at(comp_gate).begin(), input_ffs_of_gate.at(comp_gate).end()))
{
continue;
}

// disregard gates for which no successor is dependent on an additional (external) input
auto sucs = comp_gate->get_unique_successors();
if (std::all_of(sucs.begin(), sucs.end(), [&input_ffs_of_gate, &input_group](auto* sg) {
return std::includes(input_group.begin(), input_group.end(), input_ffs_of_gate.at(sg).begin(), input_ffs_of_gate.at(sg).end())
}))
{
continue;
}

// TODO kick out inverters

// TODO kick out everything that depends on other outputs only

output_group.insert(comp_gate);
}

if (output_group.size() == input_group.size() && output_group.size() <= 8)
{
sboxes_input_output_gates.push_back(std::make_pair(input_group, output_group));
}

// TODO other cases
}
}

// TODO call identify_sbox on all pairs sboxes_input_output_gates
}
}

Result<std::string>
identify_sbox(const StateCandidate* candidate, const std::set<Gate*>& component, const std::set<Gate*>& input_gates, const std::set<Gate*>& output_gates, const SBoxDatabase& db)
{
Expand Down
106 changes: 99 additions & 7 deletions plugins/hawkeye/src/state_candidate.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,11 @@ namespace hal
log_info("hawkeye", "start isolating state logic...");
auto start = std::chrono::system_clock::now();

const auto& cand_input_reg = candidate->get_input_reg();
for (const auto* out_ff : candidate->get_output_reg())
const auto& state_input_reg = candidate->get_input_reg();
const auto& state_output_reg = candidate->get_output_reg();

// DFS from output reg backwards
for (const auto* out_ff : state_output_reg)
{
auto ff_data_predecessors = out_ff->get_predecessors([](const GatePin* p, const Endpoint* _) { return p->get_type() == PinType::data; });

Expand Down Expand Up @@ -101,7 +104,7 @@ namespace hal
if (predecessor_gate->get_type()->has_property(GateTypeProperty::ff))
{
// if predecessor is part of input state reg, fill set of next state logic
if (cand_input_reg.find(predecessor_gate) != cand_input_reg.end())
if (state_input_reg.find(predecessor_gate) != state_input_reg.end())
{
state_inputs.insert(next_predecessor->get_net());
state_logic.insert(current_gate);
Expand Down Expand Up @@ -189,7 +192,7 @@ namespace hal

state_cand->m_size = candidate->get_size();

for (const auto* g : candidate->get_input_reg())
for (const auto* g : state_input_reg)
{
auto* new_g = copied_nl->create_gate(g->get_id(), g->get_type(), g->get_name());
state_cand->m_in_reg.insert(new_g);
Expand All @@ -207,9 +210,9 @@ namespace hal
copy_out_eps_of_gate(copied_nl, g, new_g);
}

if (candidate->get_input_reg() != candidate->get_output_reg())
if (state_input_reg != state_output_reg)
{
for (const auto* g : candidate->get_output_reg())
for (const auto* g : state_output_reg)
{
auto* new_g = copied_nl->create_gate(g->get_id(), g->get_type(), g->get_name());
state_cand->m_out_reg.insert(new_g);
Expand All @@ -220,7 +223,7 @@ namespace hal
else
{
// create separate FF instances for state output register so that input and output register are distinct
for (const auto* g : candidate->get_output_reg())
for (const auto* g : state_output_reg)
{
// only differences: do not enforce ID (already taken by in_reg FF) and append suffix to name
auto* new_g = copied_nl->create_gate(g->get_type(), g->get_name() + "_OUT");
Expand All @@ -247,6 +250,85 @@ namespace hal
}
state_cand->m_graph = std::move(nl_graph_res.get());

// DFS from input reg forwards
std::map<Gate*, u32> gate_to_longest_distance;
for (auto* in_ff : state_cand->m_in_reg)
{
std::vector<Gate*> stack = {in_ff};
std::vector<Gate*> previous;
while (!stack.empty())
{
auto* current_gate = stack.back();

// pop stack if last gate on stack has been dealt with completely
if (!previous.empty() && previous.back() == current_gate)
{
stack.pop_back();
previous.pop_back();
continue;
}

state_cand->m_input_ffs_of_gate[current_gate].insert(in_ff);
state_cand->m_gates_reached_by_input_ff[in_ff].insert(current_gate);

// expand towards successors
bool added = false;
for (auto* next_successor : current_gate->get_successors())
{
auto* successor_gate = next_successor->get_gate();
if (successor_gate->get_type()->has_property(GateTypeProperty::ff))
{
// if successor is part of output state reg, fill set of gates reached by input FF
if (state_cand->m_out_reg.find(successor_gate) != state_cand->m_out_reg.end())
{
state_cand->m_input_ffs_of_gate[successor_gate].insert(in_ff);
state_cand->m_gates_reached_by_input_ff[in_ff].insert(successor_gate);
}
}
else if (successor_gate->get_type()->has_property(GateTypeProperty::combinational))
{
// if successor is part of next state logic, add gate to stack
if (state_cand->m_state_logic.find(successor_gate) != state_cand->m_state_logic.end())
{
stack.push_back(successor_gate);
added = true;

u32 current_distance = previous.size() + 1;
if (const auto dist_it = gate_to_longest_distance.find(successor_gate); dist_it != gate_to_longest_distance.end())
{
u32 stored_distance = dist_it->second;
if (stored_distance < current_distance)
{
gate_to_longest_distance[successor_gate] = current_distance;
}
}
else
{
gate_to_longest_distance[successor_gate] = current_distance;
}
}
}
}

if (added)
{
// push current gate to previous if progress was made
previous.push_back(current_gate);
}
else
{
// otherwise pop last element from stack as it has been dealt with already
stack.pop_back();
}
}
}

// invert gate_to_longest_distance map to fill m_longest_distance_to_gate
for (const auto& [gate, distance] : gate_to_longest_distance)
{
state_cand->m_longest_distance_to_gate[distance].insert(gate);
}

return OK(std::move(state_cand));
}

Expand Down Expand Up @@ -299,5 +381,15 @@ namespace hal
{
return m_state_outputs;
}

const std::map<Gate*, std::set<Gate*>>& StateCandidate::get_input_ffs_of_gate() const
{
return m_input_ffs_of_gate;
}

const std::map<u32, std::set<Gate*>>& StateCandidate::get_longest_distance_to_gate() const
{
return m_longest_distance_to_gate;
}
} // namespace hawkeye
} // namespace hal

0 comments on commit dd928f8

Please sign in to comment.