Source code
Revision control
Copy as Markdown
Other Tools
/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
#include "nsBayesianFilter.h"
#include "nsIInputStream.h"
#include "nsIStreamListener.h"
#include "nsNetUtil.h"
#include "nsIMsgMessageService.h"
#include "nsMsgUtils.h" // for GetMessageServiceFromURI
#include "prnetdb.h"
#include "nsIMsgWindow.h"
#include "mozilla/Logging.h"
#include "nsAppDirectoryServiceDefs.h"
#include "nsUnicharUtils.h"
#include "nsDirectoryServiceUtils.h"
#include "nsIMIMEHeaderParam.h"
#include "nsNetCID.h"
#include "nsIMsgMailNewsUrl.h"
#include "nsIPrefService.h"
#include "nsIPrefBranch.h"
#include "nsIStringEnumerator.h"
#include "nsIObserverService.h"
#include "nsIChannel.h"
#include "nsIMailChannel.h"
#include "nsDependentSubstring.h"
#include "nsMemory.h"
#include "nsUnicodeProperties.h"
#include "mozilla/ArenaAllocatorExtensions.h" // for ArenaStrdup
using namespace mozilla;
using mozilla::intl::Script;
using mozilla::intl::UnicodeProperties;
// needed to mark attachment flag on the db hdr
#include "nsIMsgHdr.h"
// needed to strip html out of the body
#include "nsIParserUtils.h"
#include "nsIDocumentEncoder.h"
#include "nsIncompleteGamma.h"
#include <math.h>
#include <prmem.h>
#include "nsIMsgTraitService.h"
#include "mozilla/Services.h"
#include "mozilla/Attributes.h"
#include <cstdlib> // for std::abs(int/long)
#include <cmath> // for std::abs(float/double)
static mozilla::LazyLogModule BayesianFilterLogModule("BayesianFilter");
#define kDefaultJunkThreshold .99 // we override this value via a pref
static const char* kBayesianFilterTokenDelimiters = " \t\n\r\f.";
static unsigned int kMinLengthForToken =
3; // lower bound on the number of characters in a word before we treat it
// as a token
static unsigned int kMaxLengthForToken =
12; // upper bound on the number of characters in a word to be declared as
// a token
#define FORGED_RECEIVED_HEADER_HINT "may be forged"_ns
#ifndef M_LN2
# define M_LN2 0.69314718055994530942
#endif
#ifndef M_E
# define M_E 2.7182818284590452354
#endif
// provide base implementation of hash lookup of a string
struct BaseToken : public PLDHashEntryHdr {
const char* mWord;
};
// token for a particular message
// mCount, mAnalysisLink are initialized to zero by the hash code
struct Token : public BaseToken {
uint32_t mCount;
uint32_t mAnalysisLink; // index in mAnalysisStore of the AnalysisPerToken
// object for the first trait for this token
// Helper to support Tokenizer::copyTokens()
void clone(const Token& other) {
mWord = other.mWord;
mCount = other.mCount;
mAnalysisLink = other.mAnalysisLink;
}
};
// token stored in a training file for a group of messages
// mTraitLink is initialized to 0 by the hash code
struct CorpusToken : public BaseToken {
uint32_t mTraitLink; // index in mTraitStore of the TraitPerToken
// object for the first trait for this token
};
// set the value of a TraitPerToken object
TraitPerToken::TraitPerToken(uint32_t aTraitId, uint32_t aCount)
: mId(aTraitId), mCount(aCount), mNextLink(0) {}
// shorthand representations of trait ids for junk and good
static const uint32_t kJunkTrait = nsIJunkMailPlugin::JUNK_TRAIT;
static const uint32_t kGoodTrait = nsIJunkMailPlugin::GOOD_TRAIT;
// set the value of an AnalysisPerToken object
AnalysisPerToken::AnalysisPerToken(uint32_t aTraitIndex, double aDistance,
double aProbability)
: mTraitIndex(aTraitIndex),
mDistance(aDistance),
mProbability(aProbability),
mNextLink(0) {}
// the initial size of the AnalysisPerToken linked list storage
const uint32_t kAnalysisStoreCapacity = 2048;
// the initial size of the TraitPerToken linked list storage
const uint32_t kTraitStoreCapacity = 16384;
// Size of Auto arrays representing per trait information
const uint32_t kTraitAutoCapacity = 10;
TokenEnumeration::TokenEnumeration(PLDHashTable* table)
: mIterator(table->Iter()) {}
inline bool TokenEnumeration::hasMoreTokens() { return !mIterator.Done(); }
inline BaseToken* TokenEnumeration::nextToken() {
auto token = static_cast<BaseToken*>(mIterator.Get());
mIterator.Next();
return token;
}
// member variables
static const PLDHashTableOps gTokenTableOps = {
PLDHashTable::HashStringKey, PLDHashTable::MatchStringKey,
PLDHashTable::MoveEntryStub, PLDHashTable::ClearEntryStub, nullptr};
TokenHash::TokenHash(uint32_t aEntrySize)
: mTokenTable(&gTokenTableOps, aEntrySize, 128) {
mEntrySize = aEntrySize;
}
TokenHash::~TokenHash() {}
nsresult TokenHash::clearTokens() {
// we re-use the tokenizer when classifying multiple messages,
// so this gets called after every message classification.
mTokenTable.ClearAndPrepareForLength(128);
mWordPool.Clear();
return NS_OK;
}
char* TokenHash::copyWord(const char* word, uint32_t len) {
return ArenaStrdup(Substring(word, len), mWordPool);
}
inline BaseToken* TokenHash::get(const char* word) {
PLDHashEntryHdr* entry = mTokenTable.Search(word);
if (entry) return static_cast<BaseToken*>(entry);
return NULL;
}
BaseToken* TokenHash::add(const char* word) {
if (!word || !*word) {
NS_ERROR("Trying to add a null word");
return nullptr;
}
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug, ("add word: %s", word));
PLDHashEntryHdr* entry = mTokenTable.Add(word, mozilla::fallible);
BaseToken* token = static_cast<BaseToken*>(entry);
if (token) {
if (token->mWord == NULL) {
uint32_t len = strlen(word);
NS_ASSERTION(len != 0, "adding zero length word to tokenizer");
if (!len)
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug,
("adding zero length word to tokenizer"));
token->mWord = copyWord(word, len);
NS_ASSERTION(token->mWord, "copyWord failed");
if (!token->mWord) {
MOZ_LOG(BayesianFilterLogModule, LogLevel::Error,
("copyWord failed: %s (%d)", word, len));
mTokenTable.RawRemove(entry);
return NULL;
}
}
}
return token;
}
inline uint32_t TokenHash::countTokens() { return mTokenTable.EntryCount(); }
inline TokenEnumeration TokenHash::getTokens() {
return TokenEnumeration(&mTokenTable);
}
Tokenizer::Tokenizer()
: TokenHash(sizeof(Token)),
mBodyDelimiters(kBayesianFilterTokenDelimiters),
mHeaderDelimiters(kBayesianFilterTokenDelimiters),
mCustomHeaderTokenization(false),
mMaxLengthForToken(kMaxLengthForToken),
mIframeToDiv(false) {
nsresult rv;
nsCOMPtr<nsIPrefService> prefs =
do_GetService(NS_PREFSERVICE_CONTRACTID, &rv);
NS_ENSURE_SUCCESS_VOID(rv);
nsCOMPtr<nsIPrefBranch> prefBranch;
rv = prefs->GetBranch("mailnews.bayesian_spam_filter.",
getter_AddRefs(prefBranch));
NS_ENSURE_SUCCESS_VOID(rv); // no branch defined, just use defaults
/*
* RSS feeds store their summary as alternate content of an iframe. But due
* the tokenizer to replace the iframe with div for tokenization.
*/
rv = prefBranch->GetBoolPref("iframe_to_div", &mIframeToDiv);
if (NS_FAILED(rv)) mIframeToDiv = false;
/*
* the list of delimiters used to tokenize the message and body
* defaults to the value in kBayesianFilterTokenDelimiters, but may be
* set with the following preferences for the body and header
* separately.
*
* \t, \n, \v, \f, \r, and \\ will be escaped to their normal
* C-library values, all other two-letter combinations beginning with \
* will be ignored.
*/
prefBranch->GetCharPref("body_delimiters", mBodyDelimiters);
if (!mBodyDelimiters.IsEmpty())
UnescapeCString(mBodyDelimiters);
else // prefBranch empties the result when it fails :(
mBodyDelimiters.Assign(kBayesianFilterTokenDelimiters);
prefBranch->GetCharPref("header_delimiters", mHeaderDelimiters);
if (!mHeaderDelimiters.IsEmpty())
UnescapeCString(mHeaderDelimiters);
else
mHeaderDelimiters.Assign(kBayesianFilterTokenDelimiters);
/*
* Extensions may wish to enable or disable tokenization of certain headers.
* Define any headers to enable/disable in a string preference like this:
* "mailnews.bayesian_spam_filter.tokenizeheader.headername"
*
* where "headername" is the header to tokenize. For example, to tokenize the
* header "x-spam-status" use the preference:
*
* "mailnews.bayesian_spam_filter.tokenizeheader.x-spam-status"
*
* The value of the string preference will be interpreted in one of
* four ways, depending on the value:
*
* If "false" then do not tokenize that header
* If "full" then add the entire header value as a token,
* without breaking up into subtokens using delimiters
* If "standard" then tokenize the header using as delimiters the current
* value of the generic header delimiters
* Any other string is interpreted as a list of delimiters to use to parse
* the header. \t, \n, \v, \f, \r, and \\ will be escaped to their normal
* C-library values, all other two-letter combinations beginning with \
* will be ignored.
*
* Header names in the preference should be all lower case
*
* Extensions may also set the maximum length of a token (default is
* kMaxLengthForToken) by setting the int preference:
* "mailnews.bayesian_spam_filter.maxlengthfortoken"
*/
nsTArray<nsCString> headers;
// get customized maximum token length
int32_t maxLengthForToken;
rv = prefBranch->GetIntPref("maxlengthfortoken", &maxLengthForToken);
mMaxLengthForToken =
NS_SUCCEEDED(rv) ? uint32_t(maxLengthForToken) : kMaxLengthForToken;
rv = prefs->GetBranch("mailnews.bayesian_spam_filter.tokenizeheader.",
getter_AddRefs(prefBranch));
if (NS_SUCCEEDED(rv)) rv = prefBranch->GetChildList("", headers);
if (NS_SUCCEEDED(rv)) {
mCustomHeaderTokenization = true;
for (auto& header : headers) {
nsCString value;
prefBranch->GetCharPref(header.get(), value);
if (value.EqualsLiteral("false")) {
mDisabledHeaders.AppendElement(header);
continue;
}
mEnabledHeaders.AppendElement(header);
if (value.EqualsLiteral("standard"))
value.SetIsVoid(true); // Void means use default delimiter
else if (value.EqualsLiteral("full"))
value.Truncate(); // Empty means add full header
else
UnescapeCString(value);
mEnabledHeadersDelimiters.AppendElement(value);
}
}
}
Tokenizer::~Tokenizer() {}
inline Token* Tokenizer::get(const char* word) {
return static_cast<Token*>(TokenHash::get(word));
}
Token* Tokenizer::add(const char* word, uint32_t count) {
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug,
("add word: %s (count=%d)", word, count));
Token* token = static_cast<Token*>(TokenHash::add(word));
if (token) {
token->mCount += count; // hash code initializes this to zero
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug,
("adding word to tokenizer: %s (count=%d) (mCount=%d)", word, count,
token->mCount));
}
return token;
}
static bool isDecimalNumber(const char* word) {
const char* p = word;
if (*p == '-') ++p;
char c;
while ((c = *p++)) {
if (!isdigit((unsigned char)c)) return false;
}
return true;
}
static bool isASCII(const char* word) {
const unsigned char* p = (const unsigned char*)word;
unsigned char c;
while ((c = *p++)) {
if (c > 127) return false;
}
return true;
}
inline bool isUpperCase(char c) { return ('A' <= c) && (c <= 'Z'); }
static char* toLowerCase(char* str) {
char c, *p = str;
while ((c = *p++)) {
if (isUpperCase(c)) p[-1] = c + ('a' - 'A');
}
return str;
}
void Tokenizer::addTokenForHeader(const char* aTokenPrefix, nsACString& aValue,
bool aTokenizeValue,
const char* aDelimiters) {
if (aValue.Length()) {
ToLowerCase(aValue);
if (!aTokenizeValue) {
nsCString tmpStr;
tmpStr.Assign(aTokenPrefix);
tmpStr.Append(':');
tmpStr.Append(aValue);
add(tmpStr.get());
} else {
char* word;
nsCString str(aValue);
char* next = str.BeginWriting();
const char* delimiters =
!aDelimiters ? mHeaderDelimiters.get() : aDelimiters;
while ((word = NS_strtok(delimiters, &next)) != NULL) {
if (strlen(word) < kMinLengthForToken) continue;
if (isDecimalNumber(word)) continue;
if (isASCII(word)) {
nsCString tmpStr;
tmpStr.Assign(aTokenPrefix);
tmpStr.Append(':');
tmpStr.Append(word);
add(tmpStr.get());
}
}
}
}
}
void Tokenizer::tokenizeAttachments(
nsTArray<RefPtr<nsIPropertyBag2>>& attachments) {
for (auto attachment : attachments) {
nsCString contentType;
ToLowerCase(contentType);
attachment->GetPropertyAsAUTF8String(u"contentType"_ns, contentType);
addTokenForHeader("attachment/content-type", contentType);
nsCString displayName;
attachment->GetPropertyAsAUTF8String(u"displayName"_ns, displayName);
ToLowerCase(displayName);
addTokenForHeader("attachment/filename", displayName);
}
}
void Tokenizer::tokenizeHeaders(nsTArray<nsCString>& aHeaderNames,
nsTArray<nsCString>& aHeaderValues) {
nsCString headerValue;
nsAutoCString
headerName; // we'll be normalizing all header names to lower case
for (uint32_t i = 0; i < aHeaderNames.Length(); i++) {
headerName = aHeaderNames[i];
ToLowerCase(headerName);
headerValue = aHeaderValues[i];
bool headerProcessed = false;
if (mCustomHeaderTokenization) {
// Process any exceptions set from preferences
for (uint32_t i = 0; i < mEnabledHeaders.Length(); i++)
if (headerName.Equals(mEnabledHeaders[i])) {
if (mEnabledHeadersDelimiters[i].IsVoid())
// tokenize with standard delimiters for all headers
addTokenForHeader(headerName.get(), headerValue, true);
else if (mEnabledHeadersDelimiters[i].IsEmpty())
// do not break the header into tokens
addTokenForHeader(headerName.get(), headerValue);
else
// use the delimiter in mEnabledHeadersDelimiters
addTokenForHeader(headerName.get(), headerValue, true,
mEnabledHeadersDelimiters[i].get());
headerProcessed = true;
break; // we found the header, no need to look for more custom values
}
for (uint32_t i = 0; i < mDisabledHeaders.Length(); i++) {
if (headerName.Equals(mDisabledHeaders[i])) {
headerProcessed = true;
break;
}
}
if (headerProcessed) continue;
}
switch (headerName.First()) {
case 'c':
if (headerName.EqualsLiteral("content-type")) {
nsresult rv;
nsCOMPtr<nsIMIMEHeaderParam> mimehdrpar =
do_GetService(NS_MIMEHEADERPARAM_CONTRACTID, &rv);
if (NS_FAILED(rv)) break;
// extract the charset parameter
nsCString parameterValue;
mimehdrpar->GetParameterInternal(headerValue, "charset", nullptr,
nullptr,
getter_Copies(parameterValue));
addTokenForHeader("charset", parameterValue);
// create a token containing just the content type
mimehdrpar->GetParameterInternal(headerValue, "type", nullptr,
nullptr,
getter_Copies(parameterValue));
if (!parameterValue.Length())
mimehdrpar->GetParameterInternal(
headerValue, nullptr /* use first unnamed param */, nullptr,
nullptr, getter_Copies(parameterValue));
addTokenForHeader("content-type/type", parameterValue);
// XXX: should we add a token for the entire content-type header as
// well or just these parts we have extracted?
}
break;
case 'r':
if (headerName.EqualsLiteral("received")) {
// look for the string "may be forged" in the received headers.
// sendmail sometimes adds this hint This does not compile on linux
// yet. Need to figure out why. Commenting out for now if
// (FindInReadable(FORGED_RECEIVED_HEADER_HINT, headerValue))
// addTokenForHeader(headerName.get(), FORGED_RECEIVED_HEADER_HINT);
}
// leave out reply-to
break;
case 's':
if (headerName.EqualsLiteral("subject")) {
// we want to tokenize the subject
addTokenForHeader(headerName.get(), headerValue, true);
}
// important: leave out sender field. Too strong of an indicator
break;
case 'x': // (2) X-Mailer / user-agent works best if it is untokenized,
// just fold the case and any leading/trailing white space
// all headers beginning with x-mozilla are being changed by us, so
// ignore
if (StringBeginsWith(headerName, "x-mozilla"_ns)) break;
// fall through
[[fallthrough]];
case 'u':
addTokenForHeader(headerName.get(), headerValue);
break;
default:
addTokenForHeader(headerName.get(), headerValue);
break;
} // end switch
}
}
void Tokenizer::tokenize_ascii_word(char* aWord) {
// always deal with normalized lower case strings
toLowerCase(aWord);
uint32_t wordLength = strlen(aWord);
// if the wordLength is within our accepted token limit, then add it
if (wordLength >= kMinLengthForToken && wordLength <= mMaxLengthForToken)
add(aWord);
else if (wordLength > mMaxLengthForToken) {
// don't skip over the word if it looks like an email address,
// there is value in adding tokens for addresses
nsDependentCString word(aWord,
wordLength); // CHEAP, no allocation occurs here...
// XXX: i think the 40 byte check is just for perf reasons...if the email
// address is longer than that then forget about it.
const char* atSign = strchr(aWord, '@');
if (wordLength < 40 && strchr(aWord, '.') && atSign &&
!strchr(atSign + 1, '@')) {
uint32_t numBytesToSep = atSign - aWord;
if (numBytesToSep <
wordLength - 1) // if the @ sign is the last character, it must not
// be an email address
{
// split the john@foo.com into john and foo.com, treat them as separate
// tokens
nsCString emailNameToken;
emailNameToken.AssignLiteral("email name:");
emailNameToken.Append(Substring(word, 0, numBytesToSep++));
add(emailNameToken.get());
nsCString emailAddrToken;
emailAddrToken.AssignLiteral("email addr:");
emailAddrToken.Append(
Substring(word, numBytesToSep, wordLength - numBytesToSep));
add(emailAddrToken.get());
return;
}
}
// there is value in generating a token indicating the number
// of characters we are skipping. We'll round to the nearest 10
nsCString skipToken;
skipToken.AssignLiteral("skip:");
skipToken.Append(word[0]);
skipToken.Append(' ');
skipToken.AppendInt((wordLength / 10) * 10);
add(skipToken.get());
}
}
// Copied from mozilla/intl/lwbrk/WordBreaker.cpp
#define ASCII_IS_ALPHA(c) \
((('a' <= (c)) && ((c) <= 'z')) || (('A' <= (c)) && ((c) <= 'Z')))
#define ASCII_IS_DIGIT(c) (('0' <= (c)) && ((c) <= '9'))
#define ASCII_IS_SPACE(c) \
((' ' == (c)) || ('\t' == (c)) || ('\r' == (c)) || ('\n' == (c)))
#define IS_ALPHABETICAL_SCRIPT(c) ((c) < 0x2E80)
// we change the beginning of IS_HAN from 0x4e00 to 0x3400 to relfect
// Unicode 3.0
#define IS_HAN(c) \
((0x3400 <= (c)) && ((c) <= 0x9fff)) || ((0xf900 <= (c)) && ((c) <= 0xfaff))
#define IS_KATAKANA(c) ((0x30A0 <= (c)) && ((c) <= 0x30FF))
#define IS_HIRAGANA(c) ((0x3040 <= (c)) && ((c) <= 0x309F))
#define IS_HALFWIDTHKATAKANA(c) ((0xFF60 <= (c)) && ((c) <= 0xFF9F))
// Return true if aChar belongs to a SEAsian script that is written without
// word spaces, so we need to use the "complex breaker" to find possible word
// (How well this works depends on the level of platform support for finding
// possible line breaks - or possible word boundaries - in the particular
// script. Thai, at least, works pretty well on the major desktop OSes. If
// the script is not supported by the platform, we just won't find any useful
// boundaries.)
static bool IsScriptioContinua(char16_t aChar) {
Script sc = UnicodeProperties::GetScriptCode(aChar);
return sc == Script::THAI || sc == Script::MYANMAR || sc == Script::KHMER ||
sc == Script::JAVANESE || sc == Script::BALINESE ||
sc == Script::SUNDANESE || sc == Script::LAO;
}
// one subtract and one conditional jump should be faster than two conditional
// jump on most recent system.
#define IN_RANGE(x, low, high) ((uint16_t)((x) - (low)) <= (high) - (low))
#define IS_JA_HIRAGANA(x) IN_RANGE(x, 0x3040, 0x309F)
// swapping the range using xor operation to reduce conditional jump.
#define IS_JA_KATAKANA(x) \
(IN_RANGE(x ^ 0x0004, 0x30A0, 0x30FE) || (IN_RANGE(x, 0xFF66, 0xFF9F)))
#define IS_JA_KANJI(x) \
(IN_RANGE(x, 0x2E80, 0x2FDF) || IN_RANGE(x, 0x4E00, 0x9FAF))
#define IS_JA_KUTEN(x) (((x) == 0x3001) || ((x) == 0xFF64) || ((x) == 0xFF0E))
#define IS_JA_TOUTEN(x) (((x) == 0x3002) || ((x) == 0xFF61) || ((x) == 0xFF0C))
#define IS_JA_SPACE(x) ((x) == 0x3000)
#define IS_JA_FWLATAIN(x) IN_RANGE(x, 0xFF01, 0xFF5E)
#define IS_JA_FWNUMERAL(x) IN_RANGE(x, 0xFF10, 0xFF19)
#define IS_JAPANESE_SPECIFIC(x) \
(IN_RANGE(x, 0x3040, 0x30FF) || IN_RANGE(x, 0xFF01, 0xFF9F))
enum char_class {
others = 0,
space,
hiragana,
katakana,
kanji,
kuten,
touten,
kigou,
fwlatain,
ascii
};
static char_class getCharClass(char16_t c) {
char_class charClass = others;
if (IS_JA_HIRAGANA(c))
charClass = hiragana;
else if (IS_JA_KATAKANA(c))
charClass = katakana;
else if (IS_JA_KANJI(c))
charClass = kanji;
else if (IS_JA_KUTEN(c))
charClass = kuten;
else if (IS_JA_TOUTEN(c))
charClass = touten;
else if (IS_JA_FWLATAIN(c))
charClass = fwlatain;
return charClass;
}
static bool isJapanese(const char* word) {
nsString text = NS_ConvertUTF8toUTF16(word);
const char16_t* p = (const char16_t*)text.get();
char16_t c;
// it is japanese chunk if it contains any hiragana or katakana.
while ((c = *p++))
if (IS_JAPANESE_SPECIFIC(c)) return true;
return false;
}
static bool isFWNumeral(const char16_t* p1, const char16_t* p2) {
for (; p1 < p2; p1++)
if (!IS_JA_FWNUMERAL(*p1)) return false;
return true;
}
// The japanese tokenizer was added as part of Bug #277354
void Tokenizer::tokenize_japanese_word(char* chunk) {
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug,
("entering tokenize_japanese_word(%s)", chunk));
nsString srcStr = NS_ConvertUTF8toUTF16(chunk);
const char16_t* p1 = srcStr.get();
const char16_t* p2 = p1;
if (!*p2) return;
char_class cc = getCharClass(*p2);
while (*(++p2)) {
if (cc == getCharClass(*p2)) continue;
nsCString token = NS_ConvertUTF16toUTF8(p1, p2 - p1);
if ((!isDecimalNumber(token.get())) && (!isFWNumeral(p1, p2))) {
nsCString tmpStr;
tmpStr.AppendLiteral("JA:");
tmpStr.Append(token);
add(tmpStr.get());
}
cc = getCharClass(*p2);
p1 = p2;
}
}
nsresult Tokenizer::stripHTML(const nsAString& inString, nsAString& outString) {
uint32_t flags = nsIDocumentEncoder::OutputLFLineBreak |
nsIDocumentEncoder::OutputNoScriptContent |
nsIDocumentEncoder::OutputNoFramesContent |
nsIDocumentEncoder::OutputBodyOnly;
nsCOMPtr<nsIParserUtils> utils = do_GetService(NS_PARSERUTILS_CONTRACTID);
return utils->ConvertToPlainText(inString, flags, 80, outString);
}
enum WordBreakClass : uint8_t {
kWbClassSpace = 0,
kWbClassAlphaLetter,
kWbClassPunct,
kWbClassHanLetter,
kWbClassKatakanaLetter,
kWbClassHiraganaLetter,
kWbClassHWKatakanaLetter,
kWbClassScriptioContinua
};
WordBreakClass GetWordBreakClass(char16_t c) {
// begin of the hack
if (IS_ALPHABETICAL_SCRIPT(c)) {
if (IS_ASCII(c)) {
if (ASCII_IS_SPACE(c)) {
return WordBreakClass::kWbClassSpace;
}
if (ASCII_IS_ALPHA(c) || ASCII_IS_DIGIT(c) || (c == '_')) {
return WordBreakClass::kWbClassAlphaLetter;
}
return WordBreakClass::kWbClassPunct;
}
if (c == 0x00A0 /*NBSP*/) {
return WordBreakClass::kWbClassSpace;
}
if (mozilla::unicode::GetGenCategory(c) == nsUGenCategory::kPunctuation) {
return WordBreakClass::kWbClassPunct;
}
if (IsScriptioContinua(c)) {
return WordBreakClass::kWbClassScriptioContinua;
}
return WordBreakClass::kWbClassAlphaLetter;
}
if (IS_HAN(c)) {
return WordBreakClass::kWbClassHanLetter;
}
if (IS_KATAKANA(c)) {
return kWbClassKatakanaLetter;
}
if (IS_HIRAGANA(c)) {
return WordBreakClass::kWbClassHiraganaLetter;
}
if (IS_HALFWIDTHKATAKANA(c)) {
return WordBreakClass::kWbClassHWKatakanaLetter;
}
if (mozilla::unicode::GetGenCategory(c) == nsUGenCategory::kPunctuation) {
return WordBreakClass::kWbClassPunct;
}
if (IsScriptioContinua(c)) {
return WordBreakClass::kWbClassScriptioContinua;
}
return WordBreakClass::kWbClassAlphaLetter;
}
nsresult Tokenizer::ScannerNext(const char16_t* text, int32_t length,
int32_t pos, bool isLastBuffer, int32_t* begin,
int32_t* end, bool* _retval) {
// if we reach the end, just return
if (pos >= length) {
*begin = pos;
*end = pos;
*_retval = false;
return NS_OK;
}
WordBreakClass char_class = GetWordBreakClass(text[pos]);
// If we are in Chinese mode, return one Han letter at a time.
// We should not do this if we are in Japanese or Korean mode.
if (WordBreakClass::kWbClassHanLetter == char_class) {
*begin = pos;
*end = pos + 1;
*_retval = true;
return NS_OK;
}
int32_t next;
// Find the next "word".
next =
mozilla::intl::WordBreaker::Next(text, (uint32_t)length, (uint32_t)pos);
// If we don't have enough text to make decision, return.
if (next == NS_WORDBREAKER_NEED_MORE_TEXT) {
*begin = pos;
*end = isLastBuffer ? length : pos;
*_retval = isLastBuffer;
return NS_OK;
}
// If what we got is space or punct, look at the next break.
if (char_class == WordBreakClass::kWbClassSpace ||
char_class == WordBreakClass::kWbClassPunct) {
// If the next "word" is not letters,
// call itself recursively with the new pos.
return ScannerNext(text, length, next, isLastBuffer, begin, end, _retval);
}
// For the rest, return.
*begin = pos;
*end = next;
*_retval = true;
return NS_OK;
}
void Tokenizer::tokenize(const char* aText) {
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug, ("tokenize: %s", aText));
// strip out HTML tags before we begin processing
// uggh but first we have to blow up our string into UCS2
// since that's what the document encoder wants. UTF8/UCS2, I wish we all
// spoke the same language here..
nsString text = NS_ConvertUTF8toUTF16(aText);
nsString strippedUCS2;
// RSS feeds store their summary information as an iframe. But due to
// workaround, allow an option to replace iframe with div in the message
// text. We disable by default, since most people won't be applying bayes
// to RSS
if (mIframeToDiv) {
text.ReplaceSubstring(u"<iframe"_ns, u"<div"_ns);
text.ReplaceSubstring(u"/iframe>"_ns, u"/div>"_ns);
}
stripHTML(text, strippedUCS2);
// convert 0x3000(full width space) into 0x0020
char16_t* substr_start = strippedUCS2.BeginWriting();
char16_t* substr_end = strippedUCS2.EndWriting();
while (substr_start != substr_end) {
if (*substr_start == 0x3000) *substr_start = 0x0020;
++substr_start;
}
nsCString strippedStr = NS_ConvertUTF16toUTF8(strippedUCS2);
char* strippedText = strippedStr.BeginWriting();
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug,
("tokenize stripped html: %s", strippedText));
char* word;
char* next = strippedText;
while ((word = NS_strtok(mBodyDelimiters.get(), &next)) != NULL) {
if (!*word) continue;
if (isDecimalNumber(word)) continue;
if (isASCII(word))
tokenize_ascii_word(word);
else if (isJapanese(word))
tokenize_japanese_word(word);
else {
nsresult rv;
// Convert this word from UTF-8 into UCS2.
NS_ConvertUTF8toUTF16 uword(word);
ToLowerCase(uword);
const char16_t* utext = uword.get();
int32_t len = uword.Length(), pos = 0, begin, end;
bool gotUnit;
while (pos < len) {
rv = ScannerNext(utext, len, pos, true, &begin, &end, &gotUnit);
if (NS_SUCCEEDED(rv) && gotUnit) {
NS_ConvertUTF16toUTF8 utfUnit(utext + begin, end - begin);
add(utfUnit.get());
// Advance to end of current unit.
pos = end;
} else {
break;
}
}
}
}
}
// helper function to un-escape \n, \t, etc from a CString
void Tokenizer::UnescapeCString(nsCString& aCString) {
nsAutoCString result;
const char* readEnd = aCString.EndReading();
result.SetLength(aCString.Length());
char* writeStart = result.BeginWriting();
char* writeIter = writeStart;
bool inEscape = false;
for (const char* readIter = aCString.BeginReading(); readIter != readEnd;
readIter++) {
if (!inEscape) {
if (*readIter == '\\')
inEscape = true;
else
*(writeIter++) = *readIter;
} else {
inEscape = false;
switch (*readIter) {
case '\\':
*(writeIter++) = '\\';
break;
case 't':
*(writeIter++) = '\t';
break;
case 'n':
*(writeIter++) = '\n';
break;
case 'v':
*(writeIter++) = '\v';
break;
case 'f':
*(writeIter++) = '\f';
break;
case 'r':
*(writeIter++) = '\r';
break;
default:
// all other escapes are ignored
break;
}
}
}
result.Truncate(writeIter - writeStart);
aCString.Assign(result);
}
Token* Tokenizer::copyTokens() {
uint32_t count = countTokens();
if (count > 0) {
Token* tokens = new Token[count];
if (tokens) {
Token* tp = tokens;
TokenEnumeration e(&mTokenTable);
while (e.hasMoreTokens()) {
Token* src = static_cast<Token*>(e.nextToken());
tp->clone(*src);
++tp;
}
}
return tokens;
}
return NULL;
}
class TokenAnalyzer {
public:
virtual ~TokenAnalyzer() {}
virtual void analyzeTokens(Tokenizer& tokenizer) = 0;
void setTokenListener(nsIStreamListener* aTokenListener) {
mTokenListener = aTokenListener;
}
void setSource(const nsACString& sourceURI) { mTokenSource = sourceURI; }
nsCOMPtr<nsIStreamListener> mTokenListener;
nsCString mTokenSource;
};
/**
* This class downloads the raw content of an email message, buffering until
* complete segments are seen, that is until a linefeed is seen, although
* any of the valid token separators would do. This could be a further
* refinement.
*/
class TokenStreamListener : public nsIStreamListener {
public:
NS_DECL_ISUPPORTS
NS_DECL_NSIREQUESTOBSERVER
NS_DECL_NSISTREAMLISTENER
explicit TokenStreamListener(TokenAnalyzer* analyzer);
protected:
virtual ~TokenStreamListener();
TokenAnalyzer* mAnalyzer;
char* mBuffer;
uint32_t mBufferSize;
uint32_t mLeftOverCount;
Tokenizer mTokenizer;
bool mSetAttachmentFlag;
};
const uint32_t kBufferSize = 16384;
TokenStreamListener::TokenStreamListener(TokenAnalyzer* analyzer)
: mAnalyzer(analyzer),
mBuffer(NULL),
mBufferSize(kBufferSize),
mLeftOverCount(0),
mSetAttachmentFlag(false) {}
TokenStreamListener::~TokenStreamListener() {
delete[] mBuffer;
delete mAnalyzer;
}
NS_IMPL_ISUPPORTS(TokenStreamListener, nsIRequestObserver, nsIStreamListener)
/* void onStartRequest (in nsIRequest aRequest); */
NS_IMETHODIMP TokenStreamListener::OnStartRequest(nsIRequest* aRequest) {
mLeftOverCount = 0;
if (!mBuffer) {
mBuffer = new char[mBufferSize];
NS_ENSURE_TRUE(mBuffer, NS_ERROR_OUT_OF_MEMORY);
}
return NS_OK;
}
/* void onDataAvailable (in nsIRequest aRequest, in nsIInputStream aInputStream,
* in unsigned long long aOffset, in unsigned long aCount); */
NS_IMETHODIMP TokenStreamListener::OnDataAvailable(nsIRequest* aRequest,
nsIInputStream* aInputStream,
uint64_t aOffset,
uint32_t aCount) {
nsresult rv = NS_OK;
while (aCount > 0) {
uint32_t readCount, totalCount = (aCount + mLeftOverCount);
if (totalCount >= mBufferSize) {
readCount = mBufferSize - mLeftOverCount - 1;
} else {
readCount = aCount;
}
// mBuffer is supposed to be allocated in onStartRequest. But something
// is causing that to not happen, so as a last-ditch attempt we'll
// do it here.
if (!mBuffer) {
mBuffer = new char[mBufferSize];
NS_ENSURE_TRUE(mBuffer, NS_ERROR_OUT_OF_MEMORY);
}
char* buffer = mBuffer;
rv = aInputStream->Read(buffer + mLeftOverCount, readCount, &readCount);
if (NS_FAILED(rv)) break;
if (readCount == 0) {
rv = NS_ERROR_UNEXPECTED;
NS_WARNING("failed to tokenize");
break;
}
aCount -= readCount;
/* consume the tokens up to the last legal token delimiter in the buffer. */
totalCount = (readCount + mLeftOverCount);
buffer[totalCount] = '\0';
char* lastDelimiter = NULL;
char* scan = buffer + totalCount;
while (scan > buffer) {
if (strchr(mTokenizer.mBodyDelimiters.get(), *--scan)) {
lastDelimiter = scan;
break;
}
}
if (lastDelimiter) {
*lastDelimiter = '\0';
mTokenizer.tokenize(buffer);
uint32_t consumedCount = 1 + (lastDelimiter - buffer);
mLeftOverCount = totalCount - consumedCount;
if (mLeftOverCount)
memmove(buffer, buffer + consumedCount, mLeftOverCount);
} else {
/* didn't find a delimiter, keep the whole buffer around. */
mLeftOverCount = totalCount;
if (totalCount >= (mBufferSize / 2)) {
uint32_t newBufferSize = mBufferSize * 2;
char* newBuffer = new char[newBufferSize];
NS_ENSURE_TRUE(newBuffer, NS_ERROR_OUT_OF_MEMORY);
memcpy(newBuffer, mBuffer, mLeftOverCount);
delete[] mBuffer;
mBuffer = newBuffer;
mBufferSize = newBufferSize;
}
}
}
return rv;
}
/* void onStopRequest (in nsIRequest aRequest, in nsresult aStatusCode); */
NS_IMETHODIMP TokenStreamListener::OnStopRequest(nsIRequest* aRequest,
nsresult aStatusCode) {
nsCOMPtr<nsIMailChannel> mailChannel = do_QueryInterface(aRequest);
if (mailChannel) {
nsTArray<nsCString> headerNames;
nsTArray<nsCString> headerValues;
mailChannel->GetHeaderNames(headerNames);
mailChannel->GetHeaderValues(headerValues);
mTokenizer.tokenizeHeaders(headerNames, headerValues);
nsTArray<RefPtr<nsIPropertyBag2>> attachments;
mailChannel->GetAttachments(attachments);
mTokenizer.tokenizeAttachments(attachments);
}
if (mLeftOverCount) {
/* assume final buffer is complete. */
mBuffer[mLeftOverCount] = '\0';
mTokenizer.tokenize(mBuffer);
}
/* finally, analyze the tokenized message. */
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug,
("analyze the tokenized message"));
if (mAnalyzer) mAnalyzer->analyzeTokens(mTokenizer);
return NS_OK;
}
/* Implementation file */
NS_IMPL_ISUPPORTS(nsBayesianFilter, nsIMsgFilterPlugin, nsIJunkMailPlugin,
nsIMsgCorpus, nsISupportsWeakReference, nsIObserver)
nsBayesianFilter::nsBayesianFilter() : mTrainingDataDirty(false) {
int32_t junkThreshold = 0;
nsresult rv;
nsCOMPtr<nsIPrefBranch> pPrefBranch(
do_GetService(NS_PREFSERVICE_CONTRACTID, &rv));
if (pPrefBranch)
pPrefBranch->GetIntPref("mail.adaptivefilters.junk_threshold",
&junkThreshold);
mJunkProbabilityThreshold = (static_cast<double>(junkThreshold)) / 100.0;
if (mJunkProbabilityThreshold == 0 || mJunkProbabilityThreshold >= 1)
mJunkProbabilityThreshold = kDefaultJunkThreshold;
MOZ_LOG(BayesianFilterLogModule, LogLevel::Warning,
("junk probability threshold: %f", mJunkProbabilityThreshold));
mCorpus.readTrainingData();
// get parameters for training data flushing, from the prefs
nsCOMPtr<nsIPrefBranch> prefBranch;
nsCOMPtr<nsIPrefService> prefs =
do_GetService(NS_PREFSERVICE_CONTRACTID, &rv);
NS_ASSERTION(NS_SUCCEEDED(rv), "failed accessing preferences service");
rv = prefs->GetBranch(nullptr, getter_AddRefs(prefBranch));
NS_ASSERTION(NS_SUCCEEDED(rv), "failed getting preferences branch");
rv = prefBranch->GetIntPref(
"mailnews.bayesian_spam_filter.flush.minimum_interval",
&mMinFlushInterval);
// it is not a good idea to allow a minimum interval of under 1 second
if (NS_FAILED(rv) || (mMinFlushInterval <= 1000))
mMinFlushInterval = DEFAULT_MIN_INTERVAL_BETWEEN_WRITES;
rv = prefBranch->GetIntPref("mailnews.bayesian_spam_filter.junk_maxtokens",
&mMaximumTokenCount);
if (NS_FAILED(rv))
mMaximumTokenCount = 0; // which means do not limit token counts
MOZ_LOG(BayesianFilterLogModule, LogLevel::Warning,
("maximum junk tokens: %d", mMaximumTokenCount));
// give a default capacity to the memory structure used to store
// per-message/per-trait token data
mAnalysisStore.SetCapacity(kAnalysisStoreCapacity);
// dummy 0th element. Index 0 means "end of list" so we need to
// start from 1
AnalysisPerToken analysisPT(0, 0.0, 0.0);
mAnalysisStore.AppendElement(analysisPT);
mNextAnalysisIndex = 1;
}
nsresult nsBayesianFilter::Init() {
nsCOMPtr<nsIObserverService> observerService =
mozilla::services::GetObserverService();
if (observerService)
observerService->AddObserver(this, "profile-before-change", true);
return NS_OK;
}
void nsBayesianFilter::TimerCallback(nsITimer* aTimer, void* aClosure) {
// we will flush the training data to disk after enough time has passed
// since the first time a message has been classified after the last flush
nsBayesianFilter* filter = static_cast<nsBayesianFilter*>(aClosure);
filter->mCorpus.writeTrainingData(filter->mMaximumTokenCount);
filter->mTrainingDataDirty = false;
}
nsBayesianFilter::~nsBayesianFilter() {
if (mTimer) {
mTimer->Cancel();
mTimer = nullptr;
}
// call shutdown when we are going away in case we need
// to flush the training set to disk
Shutdown();
}
// this object is used for one call to classifyMessage or classifyMessages().
// So if we're classifying multiple messages, this object will be used for each
// message. It's going to hold a reference to itself, basically, to stay in
// memory.
class MessageClassifier : public TokenAnalyzer {
public:
// full classifier with arbitrary traits
MessageClassifier(nsBayesianFilter* aFilter,
nsIJunkMailClassificationListener* aJunkListener,
nsIMsgTraitClassificationListener* aTraitListener,
nsIMsgTraitDetailListener* aDetailListener,
const nsTArray<uint32_t>& aProTraits,
const nsTArray<uint32_t>& aAntiTraits,
nsIMsgWindow* aMsgWindow,
const nsTArray<nsCString>& aMessageURIs)
: mFilter(aFilter),
mJunkMailPlugin(aFilter),
mJunkListener(aJunkListener),
mTraitListener(aTraitListener),
mDetailListener(aDetailListener),
mProTraits(aProTraits.Clone()),
mAntiTraits(aAntiTraits.Clone()),
mMsgWindow(aMsgWindow),
mMessageURIs(aMessageURIs.Clone()),
mCurMessageToClassify(0) {
MOZ_ASSERT(aProTraits.Length() == aAntiTraits.Length());
}
// junk-only classifier
MessageClassifier(nsBayesianFilter* aFilter,
nsIJunkMailClassificationListener* aJunkListener,
nsIMsgWindow* aMsgWindow,
const nsTArray<nsCString>& aMessageURIs)
: mFilter(aFilter),
mJunkMailPlugin(aFilter),
mJunkListener(aJunkListener),
mTraitListener(nullptr),
mDetailListener(nullptr),
mMsgWindow(aMsgWindow),
mMessageURIs(aMessageURIs.Clone()),
mCurMessageToClassify(0) {
mProTraits.AppendElement(kJunkTrait);
mAntiTraits.AppendElement(kGoodTrait);
}
virtual ~MessageClassifier() {}
virtual void analyzeTokens(Tokenizer& tokenizer) {
mFilter->classifyMessage(tokenizer, mTokenSource, mProTraits, mAntiTraits,
mJunkListener, mTraitListener, mDetailListener);
tokenizer.clearTokens();
classifyNextMessage();
}
virtual void classifyNextMessage() {
if (++mCurMessageToClassify < mMessageURIs.Length()) {
MOZ_LOG(BayesianFilterLogModule, LogLevel::Warning,
("classifyNextMessage(%s)",
mMessageURIs[mCurMessageToClassify].get()));
mFilter->tokenizeMessage(mMessageURIs[mCurMessageToClassify], mMsgWindow,
this);
} else {
// call all listeners with null parameters to signify end of batch
if (mJunkListener)
mJunkListener->OnMessageClassified(EmptyCString(),
nsIJunkMailPlugin::UNCLASSIFIED, 0);
if (mTraitListener) {
nsTArray<uint32_t> nullTraits;
nsTArray<uint32_t> nullPercents;
mTraitListener->OnMessageTraitsClassified(EmptyCString(), nullTraits,
nullPercents);
}
mTokenListener =
nullptr; // this breaks the circular ref that keeps this object alive
// so we will be destroyed as a result.
}
}
private:
nsBayesianFilter* mFilter;
nsCOMPtr<nsIJunkMailPlugin> mJunkMailPlugin;
nsCOMPtr<nsIJunkMailClassificationListener> mJunkListener;
nsCOMPtr<nsIMsgTraitClassificationListener> mTraitListener;
nsCOMPtr<nsIMsgTraitDetailListener> mDetailListener;
nsTArray<uint32_t> mProTraits;
nsTArray<uint32_t> mAntiTraits;
nsCOMPtr<nsIMsgWindow> mMsgWindow;
nsTArray<nsCString> mMessageURIs;
uint32_t mCurMessageToClassify; // 0-based index
};
nsresult nsBayesianFilter::tokenizeMessage(const nsACString& aMessageURI,
nsIMsgWindow* aMsgWindow,
TokenAnalyzer* aAnalyzer) {
nsCOMPtr<nsIMsgMessageService> msgService;
nsresult rv =
GetMessageServiceFromURI(aMessageURI, getter_AddRefs(msgService));
NS_ENSURE_SUCCESS(rv, rv);
aAnalyzer->setSource(aMessageURI);
nsCOMPtr<nsIURI> dummyNull;
return msgService->StreamMessage(
aMessageURI, aAnalyzer->mTokenListener, aMsgWindow, nullptr,
true /* convert data */, "filter"_ns, false, getter_AddRefs(dummyNull));
}
// a TraitAnalysis is the per-token representation of the statistical
// calculations, basically created to group information that is then
// sorted by mDistance
struct TraitAnalysis {
uint32_t mTokenIndex;
double mDistance;
double mProbability;
};
// comparator required to sort an nsTArray
class compareTraitAnalysis {
public:
bool Equals(const TraitAnalysis& a, const TraitAnalysis& b) const {
return a.mDistance == b.mDistance;
}
bool LessThan(const TraitAnalysis& a, const TraitAnalysis& b) const {
return a.mDistance < b.mDistance;
}
};
inline double dmax(double x, double y) { return (x > y ? x : y); }
inline double dmin(double x, double y) { return (x < y ? x : y); }
// Chi square functions are implemented by an incomplete gamma function.
// Note that chi2P's callers multiply the arguments by 2 but chi2P
// divides them by 2 again. Inlining chi2P gives the compiler a
// chance to notice this.
// Both chi2P and nsIncompleteGammaP set *error negative on domain
// errors and nsIncompleteGammaP sets it posivive on internal errors.
// This may be useful but the chi2P callers treat any error as fatal.
// Note that converting unsigned ints to floating point can be slow on
// some platforms (like Intel) so use signed quantities for the numeric
// routines.
static inline double chi2P(double chi2, double nu, int32_t* error) {
// domain checks; set error and return a dummy value
if (chi2 < 0.0 || nu <= 0.0) {
*error = -1;
return 0.0;
}
// reversing the arguments is intentional
return nsIncompleteGammaP(nu / 2.0, chi2 / 2.0, error);
}
void nsBayesianFilter::classifyMessage(
Tokenizer& tokenizer, const nsACString& messageURI,
nsTArray<uint32_t>& aProTraits, nsTArray<uint32_t>& aAntiTraits,
nsIJunkMailClassificationListener* listener,
nsIMsgTraitClassificationListener* aTraitListener,
nsIMsgTraitDetailListener* aDetailListener) {
if (aProTraits.Length() != aAntiTraits.Length()) {
NS_ERROR("Each Pro trait needs a matching Anti trait");
return;
}
Token* tokens = tokenizer.copyTokens();
uint32_t tokenCount;
if (!tokens) {
// This can happen with problems with UTF conversion
NS_ERROR("Trying to classify a null or invalid message");
tokenCount = 0;
// don't return so that we still call the listeners
} else {
tokenCount = tokenizer.countTokens();
}
/* this part is similar to the Graham algorithm with some adjustments. */
uint32_t traitCount = aProTraits.Length();
// pro message counts per trait index
AutoTArray<uint32_t, kTraitAutoCapacity> numProMessages;
// anti message counts per trait index
AutoTArray<uint32_t, kTraitAutoCapacity> numAntiMessages;
// array of pro aliases per trait index
AutoTArray<nsTArray<uint32_t>, kTraitAutoCapacity> proAliasArrays;
// array of anti aliases per trait index
AutoTArray<nsTArray<uint32_t>, kTraitAutoCapacity> antiAliasArrays;
// construct the outgoing listener arrays
AutoTArray<uint32_t, kTraitAutoCapacity> traits;
AutoTArray<uint32_t, kTraitAutoCapacity> percents;
if (traitCount > kTraitAutoCapacity) {
traits.SetCapacity(traitCount);
percents.SetCapacity(traitCount);
numProMessages.SetCapacity(traitCount);
numAntiMessages.SetCapacity(traitCount);
proAliasArrays.SetCapacity(traitCount);
antiAliasArrays.SetCapacity(traitCount);
}
nsresult rv;
nsCOMPtr<nsIMsgTraitService> traitService(
do_GetService("@mozilla.org/msg-trait-service;1", &rv));
if (NS_FAILED(rv)) {
NS_ERROR("Failed to get trait service");
MOZ_LOG(BayesianFilterLogModule, LogLevel::Error,
("Failed to get trait service"));
}
// get aliases and message counts for the pro and anti traits
for (uint32_t traitIndex = 0; traitIndex < traitCount; traitIndex++) {
nsresult rv;
// pro trait
nsTArray<uint32_t> proAliases;
uint32_t proTrait = aProTraits[traitIndex];
if (traitService) {
rv = traitService->GetAliases(proTrait, proAliases);
if (NS_FAILED(rv)) {
NS_ERROR("trait service failed to get aliases");
MOZ_LOG(BayesianFilterLogModule, LogLevel::Error,
("trait service failed to get aliases"));
}
}
proAliasArrays.AppendElement(proAliases.Clone());
uint32_t proMessageCount = mCorpus.getMessageCount(proTrait);
for (uint32_t aliasIndex = 0; aliasIndex < proAliases.Length();
aliasIndex++)
proMessageCount += mCorpus.getMessageCount(proAliases[aliasIndex]);
numProMessages.AppendElement(proMessageCount);
// anti trait
nsTArray<uint32_t> antiAliases;
uint32_t antiTrait = aAntiTraits[traitIndex];
if (traitService) {
rv = traitService->GetAliases(antiTrait, antiAliases);
if (NS_FAILED(rv)) {
NS_ERROR("trait service failed to get aliases");
MOZ_LOG(BayesianFilterLogModule, LogLevel::Error,
("trait service failed to get aliases"));
}
}
antiAliasArrays.AppendElement(antiAliases.Clone());
uint32_t antiMessageCount = mCorpus.getMessageCount(antiTrait);
for (uint32_t aliasIndex = 0; aliasIndex < antiAliases.Length();
aliasIndex++)
antiMessageCount += mCorpus.getMessageCount(antiAliases[aliasIndex]);
numAntiMessages.AppendElement(antiMessageCount);
}
for (uint32_t i = 0; i < tokenCount; ++i) {
Token& token = tokens[i];
CorpusToken* t = mCorpus.get(token.mWord);
if (!t) continue;
for (uint32_t traitIndex = 0; traitIndex < traitCount; traitIndex++) {
uint32_t iProCount = mCorpus.getTraitCount(t, aProTraits[traitIndex]);
// add in any counts for aliases to proTrait
for (uint32_t aliasIndex = 0;
aliasIndex < proAliasArrays[traitIndex].Length(); aliasIndex++)
iProCount +=
mCorpus.getTraitCount(t, proAliasArrays[traitIndex][aliasIndex]);
double proCount = static_cast<double>(iProCount);
uint32_t iAntiCount = mCorpus.getTraitCount(t, aAntiTraits[traitIndex]);
// add in any counts for aliases to antiTrait
for (uint32_t aliasIndex = 0;
aliasIndex < antiAliasArrays[traitIndex].Length(); aliasIndex++)
iAntiCount +=
mCorpus.getTraitCount(t, antiAliasArrays[traitIndex][aliasIndex]);
double antiCount = static_cast<double>(iAntiCount);
double prob, denom;
// Prevent a divide by zero error by setting defaults for prob
// If there are no matching tokens at all, ignore.
if (antiCount == 0.0 && proCount == 0.0) continue;
// if only anti match, set probability to 0%
if (proCount == 0.0) prob = 0.0;
// if only pro match, set probability to 100%
else if (antiCount == 0.0)
prob = 1.0;
// not really needed, but just to be sure check the denom as well
else if ((denom = proCount * numAntiMessages[traitIndex] +
antiCount * numProMessages[traitIndex]) == 0.0)
continue;
else
prob = (proCount * numAntiMessages[traitIndex]) / denom;
double n = proCount + antiCount;
prob = (0.225 + n * prob) / (.45 + n);
double distance = std::abs(prob - 0.5);
if (distance >= .1) {
mozilla::DebugOnly<nsresult> rv =
setAnalysis(token, traitIndex, distance, prob);
NS_ASSERTION(NS_SUCCEEDED(rv), "Problem in setAnalysis");
}
}
}
for (uint32_t traitIndex = 0; traitIndex < traitCount; traitIndex++) {
AutoTArray<TraitAnalysis, 1024> traitAnalyses;
// copy valid tokens into an array to sort
for (uint32_t tokenIndex = 0; tokenIndex < tokenCount; tokenIndex++) {
uint32_t storeIndex = getAnalysisIndex(tokens[tokenIndex], traitIndex);
if (storeIndex) {
TraitAnalysis ta = {tokenIndex, mAnalysisStore[storeIndex].mDistance,
mAnalysisStore[storeIndex].mProbability};
traitAnalyses.AppendElement(ta);
}
}
// sort the array by the distances
traitAnalyses.Sort(compareTraitAnalysis());
uint32_t count = traitAnalyses.Length();
uint32_t first, last = count;
const uint32_t kMaxTokens = 150;
first = (count > kMaxTokens) ? count - kMaxTokens : 0;
// Setup the arrays to save details if needed
nsTArray<double> sArray;
nsTArray<double> hArray;
uint32_t usedTokenCount = (count > kMaxTokens) ? kMaxTokens : count;
if (aDetailListener) {
sArray.SetCapacity(usedTokenCount);
hArray.SetCapacity(usedTokenCount);
}
double H = 1.0, S = 1.0;
int32_t Hexp = 0, Sexp = 0;
uint32_t goodclues = 0;
int e;
// index from end to analyze most significant first
for (uint32_t ip1 = last; ip1 != first; --ip1) {
TraitAnalysis& ta = traitAnalyses[ip1 - 1];
if (ta.mDistance > 0.0) {
goodclues++;
double value = ta.mProbability;
S *= (1.0 - value);
H *= value;
if (S < 1e-200) {
S = frexp(S, &e);
Sexp += e;
}
if (H < 1e-200) {
H = frexp(H, &e);
Hexp += e;
}
MOZ_LOG(BayesianFilterLogModule, LogLevel::Warning,
("token probability (%s) is %f", tokens[ta.mTokenIndex].mWord,
ta.mProbability));
}
if (aDetailListener) {
sArray.AppendElement(log(S) + Sexp * M_LN2);
hArray.AppendElement(log(H) + Hexp * M_LN2);
}
}
S = log(S) + Sexp * M_LN2;
H = log(H) + Hexp * M_LN2;
double prob;
if (goodclues > 0) {
int32_t chi_error;
S = chi2P(-2.0 * S, 2.0 * goodclues, &chi_error);
if (!chi_error) H = chi2P(-2.0 * H, 2.0 * goodclues, &chi_error);
// if any error toss the entire calculation
if (!chi_error)
prob = (S - H + 1.0) / 2.0;
else
prob = 0.5;
} else
prob = 0.5;
if (aDetailListener) {
// Prepare output arrays
nsTArray<uint32_t> tokenPercents(usedTokenCount);
nsTArray<uint32_t> runningPercents(usedTokenCount);
nsTArray<nsString> tokenStrings(usedTokenCount);
double clueCount = 1.0;
for (uint32_t tokenIndex = 0; tokenIndex < usedTokenCount; tokenIndex++) {
TraitAnalysis& ta = traitAnalyses[last - 1 - tokenIndex];
int32_t chi_error;
S = chi2P(-2.0 * sArray[tokenIndex], 2.0 * clueCount, &chi_error);
if (!chi_error)
H = chi2P(-2.0 * hArray[tokenIndex], 2.0 * clueCount, &chi_error);
clueCount += 1.0;
double runningProb;
if (!chi_error)
runningProb = (S - H + 1.0) / 2.0;
else
runningProb = 0.5;
runningPercents.AppendElement(
static_cast<uint32_t>(runningProb * 100. + .5));
tokenPercents.AppendElement(
static_cast<uint32_t>(ta.mProbability * 100. + .5));
tokenStrings.AppendElement(
NS_ConvertUTF8toUTF16(tokens[ta.mTokenIndex].mWord));
}
aDetailListener->OnMessageTraitDetails(messageURI, aProTraits[traitIndex],
tokenStrings, tokenPercents,
runningPercents);
}
uint32_t proPercent = static_cast<uint32_t>(prob * 100. + .5);
// directly classify junk to maintain backwards compatibility
if (aProTraits[traitIndex] == kJunkTrait) {
bool isJunk = (prob >= mJunkProbabilityThreshold);
MOZ_LOG(BayesianFilterLogModule, LogLevel::Info,
("%s is junk probability = (%f) HAM SCORE:%f SPAM SCORE:%f",
PromiseFlatCString(messageURI).get(), prob, H, S));
// the algorithm in "A Plan For Spam" assumes that you have a large good
// corpus and a large junk corpus.
// that won't be the case with users who first use the junk mail trait
// so, we do certain things to encourage them to train.
//
// if there are no good tokens, assume the message is junk
// this will "encourage" the user to train
// and if there are no bad tokens, assume the message is not junk
// this will also "encourage" the user to train
// see bug #194238
if (listener && !mCorpus.getMessageCount(kGoodTrait))
isJunk = true;
else if (listener && !mCorpus.getMessageCount(kJunkTrait))
isJunk = false;
if (listener)
listener->OnMessageClassified(
messageURI,
isJunk ? nsMsgJunkStatus(nsIJunkMailPlugin::JUNK)
: nsMsgJunkStatus(nsIJunkMailPlugin::GOOD),
proPercent);
}
if (aTraitListener) {
traits.AppendElement(aProTraits[traitIndex]);
percents.AppendElement(proPercent);
}
}
if (aTraitListener)
aTraitListener->OnMessageTraitsClassified(messageURI, traits, percents);
delete[] tokens;
// reuse mAnalysisStore without clearing memory
mNextAnalysisIndex = 1;
// but shrink it back to the default size
if (mAnalysisStore.Length() > kAnalysisStoreCapacity)
mAnalysisStore.RemoveElementsAt(
kAnalysisStoreCapacity,
mAnalysisStore.Length() - kAnalysisStoreCapacity);
mAnalysisStore.Compact();
}
void nsBayesianFilter::classifyMessage(
Tokenizer& tokens, const nsACString& messageURI,
nsIJunkMailClassificationListener* aJunkListener) {
AutoTArray<uint32_t, 1> proTraits;
AutoTArray<uint32_t, 1> antiTraits;
proTraits.AppendElement(kJunkTrait);
antiTraits.AppendElement(kGoodTrait);
classifyMessage(tokens, messageURI, proTraits, antiTraits, aJunkListener,
nullptr, nullptr);
}
NS_IMETHODIMP
nsBayesianFilter::Observe(nsISupports* aSubject, const char* aTopic,
const char16_t* someData) {
if (!strcmp(aTopic, "profile-before-change")) Shutdown();
return NS_OK;
}
/* void shutdown (); */
NS_IMETHODIMP nsBayesianFilter::Shutdown() {
if (mTrainingDataDirty) mCorpus.writeTrainingData(mMaximumTokenCount);
mTrainingDataDirty = false;
return NS_OK;
}
/* readonly attribute boolean shouldDownloadAllHeaders; */
NS_IMETHODIMP nsBayesianFilter::GetShouldDownloadAllHeaders(
bool* aShouldDownloadAllHeaders) {
// bayesian filters work on the whole msg body currently.
*aShouldDownloadAllHeaders = false;
return NS_OK;
}
/* void classifyMessage (in string aMsgURL, in nsIJunkMailClassificationListener
* aListener); */
NS_IMETHODIMP nsBayesianFilter::ClassifyMessage(
const nsACString& aMessageURL, nsIMsgWindow* aMsgWindow,
nsIJunkMailClassificationListener* aListener) {
AutoTArray<nsCString, 1> urls = {PromiseFlatCString(aMessageURL)};
MessageClassifier* analyzer =
new MessageClassifier(this, aListener, aMsgWindow, urls);
NS_ENSURE_TRUE(analyzer, NS_ERROR_OUT_OF_MEMORY);
TokenStreamListener* tokenListener = new TokenStreamListener(analyzer);
NS_ENSURE_TRUE(tokenListener, NS_ERROR_OUT_OF_MEMORY);
analyzer->setTokenListener(tokenListener);
return tokenizeMessage(aMessageURL, aMsgWindow, analyzer);
}
/* void classifyMessages(in Array<ACString> aMsgURIs,
* in nsIMsgWindow aMsgWindow,
* in nsIJunkMailClassificationListener aListener); */
NS_IMETHODIMP nsBayesianFilter::ClassifyMessages(
const nsTArray<nsCString>& aMsgURLs, nsIMsgWindow* aMsgWindow,
nsIJunkMailClassificationListener* aListener) {
TokenAnalyzer* analyzer =
new MessageClassifier(this, aListener, aMsgWindow, aMsgURLs);
NS_ENSURE_TRUE(analyzer, NS_ERROR_OUT_OF_MEMORY);
TokenStreamListener* tokenListener = new TokenStreamListener(analyzer);
NS_ENSURE_TRUE(tokenListener, NS_ERROR_OUT_OF_MEMORY);
analyzer->setTokenListener(tokenListener);
return tokenizeMessage(aMsgURLs[0], aMsgWindow, analyzer);
}
nsresult nsBayesianFilter::setAnalysis(Token& token, uint32_t aTraitIndex,
double aDistance, double aProbability) {
uint32_t nextLink = token.mAnalysisLink;
uint32_t lastLink = 0;
uint32_t linkCount = 0, maxLinks = 100;
// try to find an existing element. Limit the search to maxLinks
// as a precaution
for (linkCount = 0; nextLink && linkCount < maxLinks; linkCount++) {
AnalysisPerToken& rAnalysis = mAnalysisStore[nextLink];
if (rAnalysis.mTraitIndex == aTraitIndex) {
rAnalysis.mDistance = aDistance;
rAnalysis.mProbability = aProbability;
return NS_OK;
}
lastLink = nextLink;
nextLink = rAnalysis.mNextLink;
}
if (linkCount >= maxLinks) return NS_ERROR_FAILURE;
// trait does not exist, so add it
AnalysisPerToken analysis(aTraitIndex, aDistance, aProbability);
if (mAnalysisStore.Length() == mNextAnalysisIndex)
mAnalysisStore.InsertElementAt(mNextAnalysisIndex, analysis);
else if (mAnalysisStore.Length() > mNextAnalysisIndex)
mAnalysisStore.ReplaceElementsAt(mNextAnalysisIndex, 1, analysis);
else // we can only insert at the end of the array
return NS_ERROR_FAILURE;
if (lastLink)
// the token had at least one link, so update the last link to point to
// the new item
mAnalysisStore[lastLink].mNextLink = mNextAnalysisIndex;
else
// need to update the token's first link
token.mAnalysisLink = mNextAnalysisIndex;
mNextAnalysisIndex++;
return NS_OK;
}
uint32_t nsBayesianFilter::getAnalysisIndex(Token& token,
uint32_t aTraitIndex) {
uint32_t nextLink;
uint32_t linkCount = 0, maxLinks = 100;
for (nextLink = token.mAnalysisLink; nextLink && linkCount < maxLinks;
linkCount++) {
AnalysisPerToken& rAnalysis = mAnalysisStore[nextLink];
if (rAnalysis.mTraitIndex == aTraitIndex) return nextLink;
nextLink = rAnalysis.mNextLink;
}
NS_ASSERTION(linkCount < maxLinks, "corrupt analysis store");
// Trait not found, indicate by zero
return 0;
}
NS_IMETHODIMP nsBayesianFilter::ClassifyTraitsInMessage(
const nsACString& aMsgURI, const nsTArray<uint32_t>& aProTraits,
const nsTArray<uint32_t>& aAntiTraits,
nsIMsgTraitClassificationListener* aTraitListener, nsIMsgWindow* aMsgWindow,
nsIJunkMailClassificationListener* aJunkListener) {
AutoTArray<nsCString, 1> uris = {PromiseFlatCString(aMsgURI)};
return ClassifyTraitsInMessages(uris, aProTraits, aAntiTraits, aTraitListener,
aMsgWindow, aJunkListener);
}
NS_IMETHODIMP nsBayesianFilter::ClassifyTraitsInMessages(
const nsTArray<nsCString>& aMsgURIs, const nsTArray<uint32_t>& aProTraits,
const nsTArray<uint32_t>& aAntiTraits,
nsIMsgTraitClassificationListener* aTraitListener, nsIMsgWindow* aMsgWindow,
nsIJunkMailClassificationListener* aJunkListener) {
MOZ_ASSERT(aProTraits.Length() == aAntiTraits.Length());
MessageClassifier* analyzer =
new MessageClassifier(this, aJunkListener, aTraitListener, nullptr,
aProTraits, aAntiTraits, aMsgWindow, aMsgURIs);
TokenStreamListener* tokenListener = new TokenStreamListener(analyzer);
analyzer->setTokenListener(tokenListener);
return tokenizeMessage(aMsgURIs[0], aMsgWindow, analyzer);
}
class MessageObserver : public TokenAnalyzer {
public:
MessageObserver(nsBayesianFilter* filter,
const nsTArray<uint32_t>& aOldClassifications,
const nsTArray<uint32_t>& aNewClassifications,
nsIJunkMailClassificationListener* aJunkListener,
nsIMsgTraitClassificationListener* aTraitListener)
: mFilter(filter),
mJunkMailPlugin(filter),
mJunkListener(aJunkListener),
mTraitListener(aTraitListener),
mOldClassifications(aOldClassifications.Clone()),
mNewClassifications(aNewClassifications.Clone()) {}
virtual void analyzeTokens(Tokenizer& tokenizer) {
mFilter->observeMessage(tokenizer, mTokenSource, mOldClassifications,
mNewClassifications, mJunkListener, mTraitListener);
// release reference to listener, which will allow us to go away as well.
mTokenListener = nullptr;
}
private:
nsBayesianFilter* mFilter;
nsCOMPtr<nsIJunkMailPlugin> mJunkMailPlugin;
nsCOMPtr<nsIJunkMailClassificationListener> mJunkListener;
nsCOMPtr<nsIMsgTraitClassificationListener> mTraitListener;
nsTArray<uint32_t> mOldClassifications;
nsTArray<uint32_t> mNewClassifications;
};
NS_IMETHODIMP nsBayesianFilter::SetMsgTraitClassification(
const nsACString& aMsgURI, const nsTArray<uint32_t>& aOldTraits,
const nsTArray<uint32_t>& aNewTraits,
nsIMsgTraitClassificationListener* aTraitListener, nsIMsgWindow* aMsgWindow,
nsIJunkMailClassificationListener* aJunkListener) {
MessageObserver* analyzer = new MessageObserver(
this, aOldTraits, aNewTraits, aJunkListener, aTraitListener);
NS_ENSURE_TRUE(analyzer, NS_ERROR_OUT_OF_MEMORY);
TokenStreamListener* tokenListener = new TokenStreamListener(analyzer);
NS_ENSURE_TRUE(tokenListener, NS_ERROR_OUT_OF_MEMORY);
analyzer->setTokenListener(tokenListener);
return tokenizeMessage(aMsgURI, aMsgWindow, analyzer);
}
// set new message classifications for a message
void nsBayesianFilter::observeMessage(
Tokenizer& tokenizer, const nsACString& messageURL,
nsTArray<uint32_t>& oldClassifications,
nsTArray<uint32_t>& newClassifications,
nsIJunkMailClassificationListener* aJunkListener,
nsIMsgTraitClassificationListener* aTraitListener) {
bool trainingDataWasDirty = mTrainingDataDirty;
// Uhoh...if the user is re-training then the message may already be
// classified and we are classifying it again with the same classification.
// the old code would have removed the tokens for this message then added them
// back. But this really hurts the message occurrence count for tokens if you
// just removed training.dat and are re-training. See Bug #237095 for more
// details. What can we do here? Well we can skip the token removal step if
// the classifications are the same and assume the user is just re-training.
// But this then allows users to re-classify the same message on the same
// training set over and over again leading to data skew. But that's all I can
// think to do right now to address this.....
uint32_t oldLength = oldClassifications.Length();
for (uint32_t index = 0; index < oldLength; index++) {
uint32_t trait = oldClassifications.ElementAt(index);
// skip removing if trait is also in the new set
if (newClassifications.Contains(trait)) continue;
// remove the tokens from the token set it is currently in
uint32_t messageCount;
messageCount = mCorpus.getMessageCount(trait);
if (messageCount > 0) {
mCorpus.setMessageCount(trait, messageCount - 1);
mCorpus.forgetTokens(tokenizer, trait, 1);
mTrainingDataDirty = true;
}
}
nsMsgJunkStatus newClassification = nsIJunkMailPlugin::UNCLASSIFIED;
uint32_t junkPercent =
0; // 0 here is no possibility of meeting the classification
uint32_t newLength = newClassifications.Length();
for (uint32_t index = 0; index < newLength; index++) {
uint32_t trait = newClassifications.ElementAt(index);
mCorpus.setMessageCount(trait, mCorpus.getMessageCount(trait) + 1);
mCorpus.rememberTokens(tokenizer, trait, 1);
mTrainingDataDirty = true;
if (aJunkListener) {
if (trait == kJunkTrait) {
junkPercent = nsIJunkMailPlugin::IS_SPAM_SCORE;
newClassification = nsIJunkMailPlugin::JUNK;
} else if (trait == kGoodTrait) {
junkPercent = nsIJunkMailPlugin::IS_HAM_SCORE;
newClassification = nsIJunkMailPlugin::GOOD;
}
}
}
if (aJunkListener)
aJunkListener->OnMessageClassified(messageURL, newClassification,
junkPercent);
if (aTraitListener) {
// construct the outgoing listener arrays
AutoTArray<uint32_t, kTraitAutoCapacity> traits;
AutoTArray<uint32_t, kTraitAutoCapacity> percents;
uint32_t newLength = newClassifications.Length();
if (newLength > kTraitAutoCapacity) {
traits.SetCapacity(newLength);
percents.SetCapacity(newLength);
}
traits.AppendElements(newClassifications);
for (uint32_t index = 0; index < newLength; index++)
percents.AppendElement(100); // This is 100 percent, or certainty
aTraitListener->OnMessageTraitsClassified(messageURL, traits, percents);
}
if (mTrainingDataDirty && !trainingDataWasDirty) {
// if training data became dirty just now, schedule flush
// mMinFlushInterval msec from now
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug,
("starting training data flush timer %i msec", mMinFlushInterval));
nsresult rv = NS_NewTimerWithFuncCallback(
getter_AddRefs(mTimer), nsBayesianFilter::TimerCallback, (void*)this,
mMinFlushInterval, nsITimer::TYPE_ONE_SHOT,
"nsBayesianFilter::TimerCallback", nullptr);
if (NS_FAILED(rv)) {
NS_WARNING("Could not start nsBayesianFilter timer");
}
}
}
NS_IMETHODIMP nsBayesianFilter::GetUserHasClassified(bool* aResult) {
*aResult = ((mCorpus.getMessageCount(kGoodTrait) +
mCorpus.getMessageCount(kJunkTrait)) &&
mCorpus.countTokens());
return NS_OK;
}
// Set message classification (only allows junk and good)
NS_IMETHODIMP nsBayesianFilter::SetMessageClassification(
const nsACString& aMsgURL, nsMsgJunkStatus aOldClassification,
nsMsgJunkStatus aNewClassification, nsIMsgWindow* aMsgWindow,
nsIJunkMailClassificationListener* aListener) {
AutoTArray<uint32_t, 1> oldClassifications;
AutoTArray<uint32_t, 1> newClassifications;
// convert between classifications and trait
if (aOldClassification == nsIJunkMailPlugin::JUNK)
oldClassifications.AppendElement(kJunkTrait);
else if (aOldClassification == nsIJunkMailPlugin::GOOD)
oldClassifications.AppendElement(kGoodTrait);
if (aNewClassification == nsIJunkMailPlugin::JUNK)
newClassifications.AppendElement(kJunkTrait);
else if (aNewClassification == nsIJunkMailPlugin::GOOD)
newClassifications.AppendElement(kGoodTrait);
MessageObserver* analyzer = new MessageObserver(
this, oldClassifications, newClassifications, aListener, nullptr);
NS_ENSURE_TRUE(analyzer, NS_ERROR_OUT_OF_MEMORY);
TokenStreamListener* tokenListener = new TokenStreamListener(analyzer);
NS_ENSURE_TRUE(tokenListener, NS_ERROR_OUT_OF_MEMORY);
analyzer->setTokenListener(tokenListener);
return tokenizeMessage(aMsgURL, aMsgWindow, analyzer);
}
NS_IMETHODIMP nsBayesianFilter::ResetTrainingData() {
return mCorpus.resetTrainingData();
}
NS_IMETHODIMP nsBayesianFilter::DetailMessage(
const nsACString& aMsgURI, uint32_t aProTrait, uint32_t aAntiTrait,
nsIMsgTraitDetailListener* aDetailListener, nsIMsgWindow* aMsgWindow) {
AutoTArray<uint32_t, 1> proTraits = {aProTrait};
AutoTArray<uint32_t, 1> antiTraits = {aAntiTrait};
AutoTArray<nsCString, 1> uris = {PromiseFlatCString(aMsgURI)};
MessageClassifier* analyzer =
new MessageClassifier(this, nullptr, nullptr, aDetailListener, proTraits,
antiTraits, aMsgWindow, uris);
NS_ENSURE_TRUE(analyzer, NS_ERROR_OUT_OF_MEMORY);
TokenStreamListener* tokenListener = new TokenStreamListener(analyzer);
NS_ENSURE_TRUE(tokenListener, NS_ERROR_OUT_OF_MEMORY);
analyzer->setTokenListener(tokenListener);
return tokenizeMessage(aMsgURI, aMsgWindow, analyzer);
}
// nsIMsgCorpus implementation
NS_IMETHODIMP nsBayesianFilter::CorpusCounts(uint32_t aTrait,
uint32_t* aMessageCount,
uint32_t* aTokenCount) {
NS_ENSURE_ARG_POINTER(aTokenCount);
*aTokenCount = mCorpus.countTokens();
if (aTrait && aMessageCount) *aMessageCount = mCorpus.getMessageCount(aTrait);
return NS_OK;
}
NS_IMETHODIMP nsBayesianFilter::ClearTrait(uint32_t aTrait) {
return mCorpus.ClearTrait(aTrait);
}
NS_IMETHODIMP
nsBayesianFilter::UpdateData(nsIFile* aFile, bool aIsAdd,
const nsTArray<uint32_t>& aFromTraits,
const nsTArray<uint32_t>& aToTraits) {
MOZ_ASSERT(aFromTraits.Length() == aToTraits.Length());
return mCorpus.UpdateData(aFile, aIsAdd, aFromTraits, aToTraits);
}
NS_IMETHODIMP
nsBayesianFilter::GetTokenCount(const nsACString& aWord, uint32_t aTrait,
uint32_t* aCount) {
NS_ENSURE_ARG_POINTER(aCount);
CorpusToken* t = mCorpus.get(PromiseFlatCString(aWord).get());
uint32_t count = mCorpus.getTraitCount(t, aTrait);
*aCount = count;
return NS_OK;
}
/* Corpus Store */
/*
Format of the training file for version 1:
[0xFEEDFACE]
[number good messages][number bad messages]
[number good tokens]
[count][length of word]word
...
[number bad tokens]
[count][length of word]word
...
Format of the trait file for version 1:
[0xFCA93601] (the 01 is the version)
for each trait to write
[id of trait to write] (0 means end of list)
[number of messages per trait]
for each token with non-zero count
[count]
[length of word]word
*/
CorpusStore::CorpusStore()
: TokenHash(sizeof(CorpusToken)),
mNextTraitIndex(1) // skip 0 since index=0 will mean end of linked list
{
getTrainingFile(getter_AddRefs(mTrainingFile));
mTraitStore.SetCapacity(kTraitStoreCapacity);
TraitPerToken traitPT(0, 0);
mTraitStore.AppendElement(traitPT); // dummy 0th element
}
CorpusStore::~CorpusStore() {}
inline int writeUInt32(FILE* stream, uint32_t value) {
value = PR_htonl(value);
return fwrite(&value, sizeof(uint32_t), 1, stream);
}
inline int readUInt32(FILE* stream, uint32_t* value) {
int n = fread(value, sizeof(uint32_t), 1, stream);
if (n == 1) {
*value = PR_ntohl(*value);
}
return n;
}
void CorpusStore::forgetTokens(Tokenizer& aTokenizer, uint32_t aTraitId,
uint32_t aCount) {
// if we are forgetting the tokens for a message, should only
// subtract 1 from the occurrence count for that token in the training set
// because we assume we only bumped the training set count once per messages
// containing the token.
TokenEnumeration tokens = aTokenizer.getTokens();
while (tokens.hasMoreTokens()) {
CorpusToken* token = static_cast<CorpusToken*>(tokens.nextToken());
remove(token->mWord, aTraitId, aCount);
}
}
void CorpusStore::rememberTokens(Tokenizer& aTokenizer, uint32_t aTraitId,
uint32_t aCount) {
TokenEnumeration tokens = aTokenizer.getTokens();
while (tokens.hasMoreTokens()) {
CorpusToken* token = static_cast<CorpusToken*>(tokens.nextToken());
if (!token) {
NS_ERROR("null token");
continue;
}
add(token->mWord, aTraitId, aCount);
}
}
bool CorpusStore::writeTokens(FILE* stream, bool shrink, uint32_t aTraitId) {
uint32_t tokenCount = countTokens();
uint32_t newTokenCount = 0;
// calculate the tokens for this trait to write
TokenEnumeration tokens = getTokens();
for (uint32_t i = 0; i < tokenCount; ++i) {
CorpusToken* token = static_cast<CorpusToken*>(tokens.nextToken());
uint32_t count = getTraitCount(token, aTraitId);
// Shrinking the token database is accomplished by dividing all token counts
// by 2. If shrinking, we'll ignore counts < 2, otherwise only ignore counts
// of < 1
if ((shrink && count > 1) || (!shrink && count)) newTokenCount++;
}
if (writeUInt32(stream, newTokenCount) != 1) return false;
if (newTokenCount > 0) {
TokenEnumeration tokens = getTokens();
for (uint32_t i = 0; i < tokenCount; ++i) {
CorpusToken* token = static_cast<CorpusToken*>(tokens.nextToken());
uint32_t wordCount = getTraitCount(token, aTraitId);
if (shrink) wordCount /= 2;
if (!wordCount) continue; // Don't output zero count words
if (writeUInt32(stream, wordCount) != 1) return false;
uint32_t tokenLength = strlen(token->mWord);
if (writeUInt32(stream, tokenLength) != 1) return false;
if (fwrite(token->mWord, tokenLength, 1, stream) != 1) return false;
}
}
return true;
}
bool CorpusStore::readTokens(FILE* stream, int64_t fileSize, uint32_t aTraitId,
bool aIsAdd) {
uint32_t tokenCount;
if (readUInt32(stream, &tokenCount) != 1) return false;
int64_t fpos = ftell(stream);
if (fpos < 0) return false;
uint32_t bufferSize = 4096;
char* buffer = new char[bufferSize];
if (!buffer) return false;
for (uint32_t i = 0; i < tokenCount; ++i) {
uint32_t count;
if (readUInt32(stream, &count) != 1) break;
uint32_t size;
if (readUInt32(stream, &size) != 1) break;
fpos += 8;
if (fpos + size > fileSize) {
delete[] buffer;
return false;
}
if (size >= bufferSize) {
delete[] buffer;
while (size >= bufferSize) {
bufferSize *= 2;
if (bufferSize == 0) return false;
}
buffer = new char[bufferSize];
if (!buffer) return false;
}
if (fread(buffer, size, 1, stream) != 1) break;
fpos += size;
buffer[size] = '\0';
if (aIsAdd)
add(buffer, aTraitId, count);
else
remove(buffer, aTraitId, count);
}
delete[] buffer;
return true;
}
nsresult CorpusStore::getTrainingFile(nsIFile** aTrainingFile) {
// should we cache the profile manager's directory?
nsCOMPtr<nsIFile> profileDir;
nsresult rv = NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
getter_AddRefs(profileDir));
NS_ENSURE_SUCCESS(rv, rv);
rv = profileDir->Append(u"training.dat"_ns);
NS_ENSURE_SUCCESS(rv, rv);
return profileDir->QueryInterface(NS_GET_IID(nsIFile), (void**)aTrainingFile);
}
nsresult CorpusStore::getTraitFile(nsIFile** aTraitFile) {
// should we cache the profile manager's directory?
nsCOMPtr<nsIFile> profileDir;
nsresult rv = NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
getter_AddRefs(profileDir));
NS_ENSURE_SUCCESS(rv, rv);
rv = profileDir->Append(u"traits.dat"_ns);
NS_ENSURE_SUCCESS(rv, rv);
return profileDir->QueryInterface(NS_GET_IID(nsIFile), (void**)aTraitFile);
}
static const char kMagicCookie[] = {'\xFE', '\xED', '\xFA', '\xCE'};
// random string used to identify trait file and version (last byte is version)
static const char kTraitCookie[] = {'\xFC', '\xA9', '\x36', '\x01'};
void CorpusStore::writeTrainingData(uint32_t aMaximumTokenCount) {
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug,
("writeTrainingData() entered"));
if (!mTrainingFile) return;
/*
* For backwards compatibility, write the good and junk tokens to
* training.dat; additional traits are added to a different file
*/
// open the file, and write out training data
FILE* stream;
nsresult rv = mTrainingFile->OpenANSIFileDesc("wb", &stream);
if (NS_FAILED(rv)) return;
// If the number of tokens exceeds our limit, set the shrink flag
bool shrink = false;
if ((aMaximumTokenCount > 0) && // if 0, do not limit tokens
(countTokens() > aMaximumTokenCount)) {
shrink = true;
MOZ_LOG(BayesianFilterLogModule, LogLevel::Warning,
("shrinking token data file"));
}
// We implement shrink by dividing counts by two
uint32_t shrinkFactor = shrink ? 2 : 1;
if (!((fwrite(kMagicCookie, sizeof(kMagicCookie), 1, stream) == 1) &&
(writeUInt32(stream, getMessageCount(kGoodTrait) / shrinkFactor)) &&
(writeUInt32(stream, getMessageCount(kJunkTrait) / shrinkFactor)) &&
writeTokens(stream, shrink, kGoodTrait) &&
writeTokens(stream, shrink, kJunkTrait))) {
NS_WARNING("failed to write training data.");
fclose(stream);
// delete the training data file, since it is potentially corrupt.
mTrainingFile->Remove(false);
} else {
fclose(stream);
}
/*
* Write the remaining data to a second file traits.dat
*/
if (!mTraitFile) {
getTraitFile(getter_AddRefs(mTraitFile));
if (!mTraitFile) return;
}
// open the file, and write out training data
rv = mTraitFile->OpenANSIFileDesc("wb", &stream);
if (NS_FAILED(rv)) return;
uint32_t numberOfTraits = mMessageCounts.Length();
bool error;
while (1) // break on error or done
{
if ((error = (fwrite(kTraitCookie, sizeof(kTraitCookie), 1, stream) != 1)))
break;
for (uint32_t index = 0; index < numberOfTraits; index++) {
uint32_t trait = mMessageCountsId[index];
if (trait == 1 || trait == 2)
continue; // junk traits are stored in training.dat
if ((error = (writeUInt32(stream, trait) != 1))) break;
if ((error = (writeUInt32(stream, mMessageCounts[index] / shrinkFactor) !=
1)))
break;
if ((error = !writeTokens(stream, shrink, trait))) break;
}
break;
}
// we add a 0 at the end to represent end of trait list
error = writeUInt32(stream, 0) != 1;
fclose(stream);
if (error) {
NS_WARNING("failed to write trait data.");
// delete the trait data file, since it is probably corrupt.
mTraitFile->Remove(false);
}
if (shrink) {
// We'll clear the tokens, and read them back in from the file.
// Yes this is slower than in place, but this is a rare event.
if (countTokens()) {
clearTokens();
for (uint32_t index = 0; index < numberOfTraits; index++)
mMessageCounts[index] = 0;
}
readTrainingData();
}
}
void CorpusStore::readTrainingData() {
/*
* To maintain backwards compatibility, good and junk traits
* are stored in a file "training.dat"
*/
if (!mTrainingFile) return;
bool exists;
nsresult rv = mTrainingFile->Exists(&exists);
if (NS_FAILED(rv) || !exists) return;
FILE* stream;
rv = mTrainingFile->OpenANSIFileDesc("rb", &stream);
if (NS_FAILED(rv)) return;
int64_t fileSize;
rv = mTrainingFile->GetFileSize(&fileSize);
if (NS_FAILED(rv)) return;
// FIXME: should make sure that the tokenizers are empty.
char cookie[4];
uint32_t goodMessageCount = 0, junkMessageCount = 0;
if (!((fread(cookie, sizeof(cookie), 1, stream) == 1) &&
(memcmp(cookie, kMagicCookie, sizeof(cookie)) == 0) &&
(readUInt32(stream, &goodMessageCount) == 1) &&
(readUInt32(stream, &junkMessageCount) == 1) &&
readTokens(stream, fileSize, kGoodTrait, true) &&
readTokens(stream, fileSize, kJunkTrait, true))) {
NS_WARNING("failed to read training data.");
MOZ_LOG(BayesianFilterLogModule, LogLevel::Error,
("failed to read training data."));
}
setMessageCount(kGoodTrait, goodMessageCount);
setMessageCount(kJunkTrait, junkMessageCount);
fclose(stream);
/*
* Additional traits are stored in traits.dat
*/
if (!mTraitFile) {
getTraitFile(getter_AddRefs(mTraitFile));
if (!mTraitFile) return;
}
rv = mTraitFile->Exists(&exists);
if (NS_FAILED(rv) || !exists) return;
nsTArray<uint32_t> empty;
rv = UpdateData(mTraitFile, true, empty, empty);
if (NS_FAILED(rv)) {
NS_WARNING("failed to read training data.");
MOZ_LOG(BayesianFilterLogModule, LogLevel::Error,
("failed to read training data."));
}
return;
}
nsresult CorpusStore::resetTrainingData() {
// clear out our in memory training tokens...
if (countTokens()) clearTokens();
uint32_t length = mMessageCounts.Length();
for (uint32_t index = 0; index < length; index++) mMessageCounts[index] = 0;
if (mTrainingFile) mTrainingFile->Remove(false);
if (mTraitFile) mTraitFile->Remove(false);
return NS_OK;
}
inline CorpusToken* CorpusStore::get(const char* word) {
return static_cast<CorpusToken*>(TokenHash::get(word));
}
nsresult CorpusStore::updateTrait(CorpusToken* token, uint32_t aTraitId,
int32_t aCountChange) {
NS_ENSURE_ARG_POINTER(token);
uint32_t nextLink = token->mTraitLink;
uint32_t lastLink = 0;
uint32_t linkCount, maxLinks = 100; // sanity check
for (linkCount = 0; nextLink && linkCount < maxLinks; linkCount++) {
TraitPerToken& traitPT = mTraitStore[nextLink];
if (traitPT.mId == aTraitId) {
// be careful with signed versus unsigned issues here
if (static_cast<int32_t>(traitPT.mCount) + aCountChange > 0)
traitPT.mCount += aCountChange;
else
traitPT.mCount = 0;
// we could delete zero count traits here, but let's not. It's rare
// anyway.
return NS_OK;
}
lastLink = nextLink;
nextLink = traitPT.mNextLink;
}
if (linkCount >= maxLinks) return NS_ERROR_FAILURE;
// trait does not exist, so add it
if (aCountChange > 0) // don't set a negative count
{
TraitPerToken traitPT(aTraitId, aCountChange);
if (mTraitStore.Length() == mNextTraitIndex)
mTraitStore.InsertElementAt(mNextTraitIndex, traitPT);
else if (mTraitStore.Length() > mNextTraitIndex)
mTraitStore.ReplaceElementsAt(mNextTraitIndex, 1, traitPT);
else
return NS_ERROR_FAILURE;
if (lastLink)
// the token had a parent, so update it
mTraitStore[lastLink].mNextLink = mNextTraitIndex;
else
// need to update the token's root link
token->mTraitLink = mNextTraitIndex;
mNextTraitIndex++;
}
return NS_OK;
}
uint32_t CorpusStore::getTraitCount(CorpusToken* token, uint32_t aTraitId) {
uint32_t nextLink;
if (!token || !(nextLink = token->mTraitLink)) return 0;
uint32_t linkCount, maxLinks = 100; // sanity check
for (linkCount = 0; nextLink && linkCount < maxLinks; linkCount++) {
TraitPerToken& traitPT = mTraitStore[nextLink];
if (traitPT.mId == aTraitId) return traitPT.mCount;
nextLink = traitPT.mNextLink;
}
NS_ASSERTION(linkCount < maxLinks, "Corrupt trait count store");
// trait not found (or error), so count is zero
return 0;
}
CorpusToken* CorpusStore::add(const char* word, uint32_t aTraitId,
uint32_t aCount) {
CorpusToken* token = static_cast<CorpusToken*>(TokenHash::add(word));
if (token) {
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug,
("adding word to corpus store: %s (Trait=%d) (deltaCount=%d)", word,
aTraitId, aCount));
updateTrait(token, aTraitId, aCount);
}
return token;
}
void CorpusStore::remove(const char* word, uint32_t aTraitId, uint32_t aCount) {
MOZ_LOG(BayesianFilterLogModule, LogLevel::Debug,
("remove word: %s (TraitId=%d) (Count=%d)", word, aTraitId, aCount));
CorpusToken* token = get(word);
if (token) updateTrait(token, aTraitId, -static_cast<int32_t>(aCount));
}
uint32_t CorpusStore::getMessageCount(uint32_t aTraitId) {
size_t index = mMessageCountsId.IndexOf(aTraitId);
if (index == mMessageCountsId.NoIndex) return 0;
return mMessageCounts.ElementAt(index);
}
void CorpusStore::setMessageCount(uint32_t aTraitId, uint32_t aCount) {
size_t index = mMessageCountsId.IndexOf(aTraitId);
if (index == mMessageCountsId.NoIndex) {
mMessageCounts.AppendElement(aCount);
mMessageCountsId.AppendElement(aTraitId);
} else {
mMessageCounts[index] = aCount;
}
}
nsresult CorpusStore::UpdateData(nsIFile* aFile, bool aIsAdd,
const nsTArray<uint32_t>& aFromTraits,
const nsTArray<uint32_t>& aToTraits) {
NS_ENSURE_ARG_POINTER(aFile);
MOZ_ASSERT(aFromTraits.Length() == aToTraits.Length());
int64_t fileSize;
nsresult rv = aFile->GetFileSize(&fileSize);
NS_ENSURE_SUCCESS(rv, rv);
FILE* stream;
rv = aFile->OpenANSIFileDesc("rb", &stream);
NS_ENSURE_SUCCESS(rv, rv);
bool error;
do // break on error or done
{
char cookie[4];
if ((error = (fread(cookie, sizeof(cookie), 1, stream) != 1))) break;
if ((error = memcmp(cookie, kTraitCookie, sizeof(cookie)))) break;
uint32_t fileTrait;
while (!(error = (readUInt32(stream, &fileTrait) != 1)) && fileTrait) {
uint32_t count;
if ((error = (readUInt32(stream, &count) != 1))) break;
uint32_t localTrait = fileTrait;
// remap the trait
for (uint32_t i = 0; i < aFromTraits.Length(); i++) {
if (aFromTraits[i] == fileTrait) localTrait = aToTraits[i];
}
uint32_t messageCount = getMessageCount(localTrait);
if (aIsAdd)
messageCount += count;
else if (count > messageCount)
messageCount = 0;
else
messageCount -= count;
setMessageCount(localTrait, messageCount);
if ((error = !readTokens(stream, fileSize, localTrait, aIsAdd))) break;
}
break;
} while (0);
fclose(stream);
if (error) return NS_ERROR_FAILURE;
return NS_OK;
}
nsresult CorpusStore::ClearTrait(uint32_t aTrait) {
// clear message counts
setMessageCount(aTrait, 0);
TokenEnumeration tokens = getTokens();
while (tokens.hasMoreTokens()) {
CorpusToken* token = static_cast<CorpusToken*>(tokens.nextToken());
int32_t wordCount = static_cast<int32_t>(getTraitCount(token, aTrait));
updateTrait(token, aTrait, -wordCount);
}
return NS_OK;
}