void yr_ac_create_failure_links(
YR_ARENA* arena,
YR_AC_AUTOMATON* automaton)
{
YR_AC_STATE_TRANSITION transition;
YR_AC_STATE* current_state;
YR_AC_STATE* failure_state;
YR_AC_STATE* temp_state;
YR_AC_STATE* state;
YR_AC_STATE* transition_state;
YR_AC_STATE* root_state;
YR_AC_MATCH* match;
QUEUE queue;
queue.head = NULL;
queue.tail = NULL;
root_state = automaton->root;
// Set the failure link of root state to itself.
root_state->failure = root_state;
// Push root's children and set their failure link to root.
state = _yr_ac_first_transition(root_state, &transition);
while (state != NULL)
{
_yr_ac_queue_push(&queue, state);
state->failure = root_state;
state = _yr_ac_next_transition(root_state, &transition);
}
// Traverse the trie in BFS order calculating the failure link
// for each state.
while (!_yr_ac_queue_is_empty(&queue))
{
current_state = _yr_ac_queue_pop(&queue);
match = current_state->matches;
if (match != NULL)
{
while (match->next != NULL)
match = match->next;
if (match->backtrack > 0)
match->next = root_state->matches;
}
else
{
current_state->matches = root_state->matches;
}
transition_state = _yr_ac_first_transition(
current_state,
&transition);
while (transition_state != NULL)
{
_yr_ac_queue_push(&queue, transition_state);
failure_state = current_state->failure;
while (1)
{
temp_state = yr_ac_next_state(
failure_state,
transition.input);
if (temp_state != NULL)
{
transition_state->failure = temp_state;
if (transition_state->matches == NULL)
{
transition_state->matches = temp_state->matches;
}
else
{
match = transition_state->matches;
while (match != NULL && match->next != NULL)
match = match->next;
match->next = temp_state->matches;
}
break;
}
else
{
if (failure_state == root_state)
{
transition_state->failure = root_state;
break;
}
else
{
failure_state = failure_state->failure;
}
}
} // while(1)
transition_state = _yr_ac_next_transition(
current_state,
&transition);
}
} // while(!__yr_ac_queue_is_empty(&queue))
}
int yr_ac_add_string(
YR_ARENA* arena,
YR_AC_AUTOMATON* automaton,
YR_STRING* string,
YR_ATOM_LIST_ITEM* atom)
{
int result = ERROR_SUCCESS;
int i;
YR_AC_STATE* state;
YR_AC_STATE* next_state;
YR_AC_MATCH* new_match;
// For each atom create the states in the automaton.
while (atom != NULL)
{
state = automaton->root;
for(i = 0; i < atom->atom_length; i++)
{
next_state = yr_ac_next_state(
state, atom->atom[i]);
if (next_state == NULL)
{
next_state = _yr_ac_create_state(
arena,
state,
atom->atom[i]);
if (next_state == NULL)
return ERROR_INSUFICIENT_MEMORY;
}
state = next_state;
}
result = yr_arena_allocate_struct(
arena,
sizeof(YR_AC_MATCH),
(void**) &new_match,
offsetof(YR_AC_MATCH, string),
offsetof(YR_AC_MATCH, forward_code),
offsetof(YR_AC_MATCH, backward_code),
offsetof(YR_AC_MATCH, next),
EOL);
if (result == ERROR_SUCCESS)
{
new_match->backtrack = state->depth + atom->backtrack;
new_match->string = string;
new_match->forward_code = atom->forward_code;
new_match->backward_code = atom->backward_code;
new_match->next = state->matches;
state->matches = new_match;
}
else
{
break;
}
atom = atom->next;
}
return result;
}
int yr_rules_scan_mem_block(
YR_RULES* rules,
uint8_t* data,
size_t data_size,
int fast_scan_mode,
int timeout,
time_t start_time,
YR_ARENA* matches_arena)
{
YR_AC_STATE* next_state;
YR_AC_MATCH* ac_match;
YR_AC_STATE* current_state;
size_t i;
current_state = rules->automaton->root;
i = 0;
while (i < data_size)
{
ac_match = current_state->matches;
while (ac_match != NULL)
{
if (ac_match->backtrack <= i)
{
FAIL_ON_ERROR(_yr_scan_verify_match(
ac_match,
data,
data_size,
i - ac_match->backtrack,
matches_arena,
fast_scan_mode));
}
ac_match = ac_match->next;
}
next_state = yr_ac_next_state(current_state, data[i]);
while (next_state == NULL && current_state->depth > 0)
{
current_state = current_state->failure;
next_state = yr_ac_next_state(current_state, data[i]);
}
if (next_state != NULL)
current_state = next_state;
i++;
if (timeout > 0 && i % 256 == 0)
{
if (difftime(time(NULL), start_time) > timeout)
return ERROR_SCAN_TIMEOUT;
}
}
ac_match = current_state->matches;
while (ac_match != NULL)
{
if (ac_match->backtrack <= data_size)
{
FAIL_ON_ERROR(_yr_scan_verify_match(
ac_match,
data,
data_size,
data_size - ac_match->backtrack,
matches_arena,
fast_scan_mode));
}
ac_match = ac_match->next;
}
return ERROR_SUCCESS;
}
