Skip to content

Commit

Permalink
Fixed ImapTokenCache to properly equate non-ASCII tokens
Browse files Browse the repository at this point in the history
  • Loading branch information
@jstedfast
jstedfast committed May 18, 2024
1 parent 71319f6 commit d82f45c
Showing 1 changed file with 47 additions and 64 deletions.
111 changes: 47 additions & 64 deletions MailKit/Net/Imap/ImapTokenCache.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,9 @@

using System;
using System.Text;
using System.Collections.Generic;
using System.Buffers;
using System.Diagnostics;
using System.Collections.Generic;

namespace MailKit.Net.Imap
{
Expand All @@ -40,7 +40,7 @@ class ImapTokenCache
readonly LinkedList<ImapTokenItem> list;
readonly ImapTokenKey lookupKey;
readonly Decoder[] decoders;
readonly char[] chars;
char[] charBuffer;

public ImapTokenCache ()
{
Expand All @@ -54,13 +54,13 @@ public ImapTokenCache ()
TextEncodings.Latin1.GetDecoder ()
};

chars = new char[128];
charBuffer = ArrayPool<char>.Shared.Rent (256);
}

public ImapToken AddOrGet (ImapTokenType type, ByteArrayBuilder builder)
{
// lookupKey is a pre-allocated key used for lookups
lookupKey.Init (decoders, chars, type, builder.GetBuffer (), builder.Length, out var decoder, out int charsNeeded);
lookupKey.Init (decoders, ref charBuffer, type, builder.GetBuffer (), builder.Length, out int charsNeeded);

if (cache.TryGetValue (lookupKey, out var node)) {
// move the node to the head of the list
Expand All @@ -71,26 +71,7 @@ public ImapToken AddOrGet (ImapTokenType type, ByteArrayBuilder builder)
return node.Value.Token;
}

string value;

if (charsNeeded <= chars.Length) {
// If the number of needed chars is <= the length of our temp buffer, then it should all be contained.
value = new string (chars, 0, charsNeeded);
} else {
var buffer = ArrayPool<char>.Shared.Rent (charsNeeded);
try {
// Note: This conversion should go flawlessly, so we'll just Debug.Assert() our expectations.
decoder.Convert (builder.GetBuffer (), 0, builder.Length, buffer, 0, buffer.Length, true, out var bytesUsed, out var charsUsed, out var completed);
Debug.Assert (bytesUsed == builder.Length);
Debug.Assert (charsUsed == charsNeeded);
Debug.Assert (completed);
value = new string (buffer, 0, charsUsed);
} finally {
ArrayPool<char>.Shared.Return (buffer);
decoder.Reset ();
}
}

var value = new string (charBuffer, 0, charsNeeded);
var token = new ImapToken (type, value);

if (cache.Count >= capacity) {
Expand All @@ -100,10 +81,10 @@ public ImapToken AddOrGet (ImapTokenType type, ByteArrayBuilder builder)
cache.Remove (node.Value.Key);

// re-use the node, item and key to avoid allocations
node.Value.Key.Init (type, (string) token.Value);
node.Value.Key.Init (type, value, lookupKey);
node.Value.Token = token;
} else {
var key = new ImapTokenKey (type, (string) token.Value);
var key = new ImapTokenKey (type, value, lookupKey);
var item = new ImapTokenItem (key, token);

node = new LinkedListNode<ImapTokenItem> (item);
Expand All @@ -118,47 +99,50 @@ public ImapToken AddOrGet (ImapTokenType type, ByteArrayBuilder builder)
class ImapTokenKey
{
ImapTokenType type;
byte[] byteArrayKey;
char[] charBuffer;
string stringKey;
int length;
int hashCode;
int length;

public ImapTokenKey ()
{
}

public ImapTokenKey (ImapTokenType type, string key)
public ImapTokenKey (ImapTokenType type, string value, ImapTokenKey key)
{
Init (type, key);
Init (type, value, key);
}

public void Init (Decoder[] decoders, char[] chars, ImapTokenType type, byte[] key, int length, out Decoder correctDecoder, out int charsNeeded)
public void Init (Decoder[] decoders, ref char[] charBuffer, ImapTokenType type, byte[] key, int length, out int charsNeeded)
{
this.type = type;
this.byteArrayKey = key;
this.stringKey = null;
this.length = length;

var hash = new HashCode ();
hash.Add ((int) type);

correctDecoder = null;
charsNeeded = 0;

// Make sure the char buffer is at least as large as the key.
if (charBuffer.Length < length) {
ArrayPool<char>.Shared.Return (charBuffer);
charBuffer = ArrayPool<char>.Shared.Rent (length);
}

foreach (var decoder in decoders) {
bool completed;
int index = 0;

correctDecoder = decoder;

do {
try {
decoder.Convert (key, index, length - index, chars, 0, chars.Length, true, out var bytesUsed, out var charsUsed, out completed);
decoder.Convert (key, index, length - index, charBuffer, charsNeeded, charBuffer.Length - charsNeeded, true, out var bytesUsed, out var charsUsed, out completed);
charsNeeded += charsUsed;
index += bytesUsed;

for (int i = 0; i < charsUsed; i++)
hash.Add (chars[i]);
hash.Add (charBuffer[i]);

if (completed)
break;
} catch (DecoderFallbackException) {
// Restart the hash...
hash = new HashCode ();
Expand All @@ -167,36 +151,39 @@ public void Init (Decoder[] decoders, char[] chars, ImapTokenType type, byte[] k
charsNeeded = 0;
break;
}
} while (!completed);

// The char buffer was not large enough to contain the full token. Resize it and try again.
var newBuffer = ArrayPool<char>.Shared.Rent (charBuffer.Length + (length - index));
charBuffer.AsSpan (0, charsNeeded).CopyTo (newBuffer);
ArrayPool<char>.Shared.Return (charBuffer);
charBuffer = newBuffer;
} while (true);

decoder.Reset ();

if (completed)
break;
}

this.charBuffer = charBuffer;
this.length = charsNeeded;

this.hashCode = hash.ToHashCode ();
}

public void Init (ImapTokenType type, string key)
public void Init (ImapTokenType type, string value, ImapTokenKey key)
{
this.type = type;
this.byteArrayKey = null;
this.stringKey = key;
this.length = key.Length;

var hash = new HashCode ();
hash.Add ((int) type);
for (int i = 0; i < length; i++)
hash.Add (key[i]);

this.hashCode = hash.ToHashCode ();
this.charBuffer = null;
this.stringKey = value;
this.length = value.Length;
this.hashCode = key.hashCode;
}

static bool Equals (string str, byte[] bytes)
static bool Equals (string str, char[] chars)
{
for (int i = 0; i < str.Length; i++) {
if (str[i] != (char) bytes[i])
if (str[i] != chars[i])
return false;
}

Expand All @@ -208,22 +195,18 @@ static bool Equals (ImapTokenKey self, ImapTokenKey other)
if (self.type != other.type || self.length != other.length)
return false;

// Note: At most, only one of the ImapTokenKeys will use a charBuffer and that ImapTokenKey will be the lookup key.
if (self.stringKey != null) {
if (other.stringKey != null)
return self.stringKey.Equals (other.stringKey, StringComparison.Ordinal);

return Equals (self.stringKey, other.byteArrayKey);
}

if (other.stringKey != null)
return Equals (other.stringKey, self.byteArrayKey);
return Equals (self.stringKey, other.charBuffer);
} else {
// Note: 'self' MUST be the lookup key.
Debug.Assert (self.charBuffer != null);

for (int i = 0; i < self.length; i++) {
if (self.byteArrayKey[i] != other.byteArrayKey[i])
return false;
return Equals (other.stringKey, self.charBuffer);
}

return true;
}

public override bool Equals (object obj)
Expand All @@ -238,7 +221,7 @@ public override int GetHashCode ()

public override string ToString ()
{
return string.Format ("{0}: {1}", type, stringKey ?? Encoding.UTF8.GetString (byteArrayKey, 0, length));
return string.Format ("{0}: {1}", type, stringKey ?? new string (charBuffer, 0, length));
}
}

Expand Down

0 comments on commit d82f45c

Please sign in to comment.