From 59ea320d960ff857a00690d9c8639ab89cfff900 Mon Sep 17 00:00:00 2001
From: Bartlomiej Plotka <bwplotka@gmail.com>
Date: Wed, 8 Jan 2020 20:39:38 +0000
Subject: [PATCH] tsdb: Fixed Symbol Lookup edge case; Added tests.

Signed-off-by: Bartlomiej Plotka <bwplotka@gmail.com>
---
 tsdb/index/index.go      |  9 +++-----
 tsdb/index/index_test.go | 48 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/tsdb/index/index.go b/tsdb/index/index.go
index 81d4a808b..116ae95b4 100644
--- a/tsdb/index/index.go
+++ b/tsdb/index/index.go
@@ -1234,17 +1234,13 @@ func NewSymbols(bs ByteSlice, version int, off int) (*Symbols, error) {
 		version: version,
 		off:     off,
 	}
-	if off == 0 {
-		// Only happens in some tests.
-		return nil, nil
-	}
 	d := encoding.NewDecbufAt(bs, off, castagnoliTable)
 	var (
 		origLen = d.Len()
 		cnt     = d.Be32int()
 		basePos = off + 4
 	)
-	s.offsets = make([]int, 0, cnt/symbolFactor)
+	s.offsets = make([]int, 0, 1+cnt/symbolFactor)
 	for d.Err() == nil && s.seen < cnt {
 		if s.seen%symbolFactor == 0 {
 			s.offsets = append(s.offsets, basePos+origLen-d.Len())
@@ -1262,8 +1258,9 @@ func (s Symbols) Lookup(o uint32) (string, error) {
 	d := encoding.Decbuf{
 		B: s.bs.Range(0, s.bs.Len()),
 	}
+
 	if s.version == FormatV2 {
-		if int(o) > s.seen {
+		if int(o) >= s.seen {
 			return "", errors.Errorf("unknown symbol offset %d", o)
 		}
 		d.Skip(s.offsets[int(o/symbolFactor)])
diff --git a/tsdb/index/index_test.go b/tsdb/index/index_test.go
index 5868184ac..a7d38a88b 100644
--- a/tsdb/index/index_test.go
+++ b/tsdb/index/index_test.go
@@ -16,6 +16,7 @@ package index
 import (
 	"context"
 	"fmt"
+	"hash/crc32"
 	"io/ioutil"
 	"math/rand"
 	"os"
@@ -505,3 +506,50 @@ func TestNewFileReaderErrorNoOpenFiles(t *testing.T) {
 	// dir.Close will fail on Win if idxName fd is not closed on error path.
 	dir.Close()
 }
+
+func TestSymbols(t *testing.T) {
+	buf := encoding.Encbuf{}
+
+	// Add prefix to the buffer to simulate symbols as part of larger buffer.
+	buf.PutUvarintStr("something")
+
+	symbolsStart := buf.Len()
+	buf.PutBE32int(204) // Length of symbols table.
+	buf.PutBE32int(100) // Number of symbols.
+	for i := 0; i < 100; i++ {
+		// i represents index in unicode characters table.
+		buf.PutUvarintStr(string(i)) // Symbol.
+	}
+	checksum := crc32.Checksum(buf.Get()[symbolsStart+4:], castagnoliTable)
+	buf.PutBE32(checksum) // Check sum at the end.
+
+	s, err := NewSymbols(realByteSlice(buf.Get()), FormatV2, symbolsStart)
+	testutil.Ok(t, err)
+
+	// We store only 4 offsets to symbols.
+	testutil.Equals(t, 32, s.Size())
+
+	for i := 99; i >= 0; i-- {
+		s, err := s.Lookup(uint32(i))
+		testutil.Ok(t, err)
+		testutil.Equals(t, string(i), s)
+	}
+	_, err = s.Lookup(100)
+	testutil.NotOk(t, err)
+
+	for i := 99; i >= 0; i-- {
+		r, err := s.ReverseLookup(string(i))
+		testutil.Ok(t, err)
+		testutil.Equals(t, uint32(i), r)
+	}
+	_, err = s.ReverseLookup(string(100))
+	testutil.NotOk(t, err)
+
+	iter := s.Iter()
+	i := 0
+	for iter.Next() {
+		testutil.Equals(t, string(i), iter.At())
+		i++
+	}
+	testutil.Ok(t, iter.Err())
+}