<commit_before>package obj

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"

	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"

	"github.com/Azure/azure-sdk-for-go/storage"
)

type microsoftClient struct {
	blobClient storage.BlobStorageClient
	container  string
}

func newMicrosoftClient(container string, accountName string, accountKey string) (*microsoftClient, error) {
	client, err := storage.NewBasicClient(
		accountName,
		accountKey,
	)
	if err != nil {
		return nil, err
	}

	return &microsoftClient{
		blobClient: client.GetBlobService(),
		container:  container,
	}, nil
}

func (c *microsoftClient) Writer(name string) (io.WriteCloser, error) {
	writer, err := newMicrosoftWriter(c, name)
	if err != nil {
		return nil, err
	}
	return newBackoffWriteCloser(c, writer), nil
}

func (c *microsoftClient) Reader(name string, offset uint64, size uint64) (io.ReadCloser, error) {
	byteRange := byteRange(offset, size)
	var reader io.ReadCloser
	var err error
	if byteRange == "" {
		reader, err = c.blobClient.GetBlob(c.container, name)
	} else {
		reader, err = c.blobClient.GetBlobRange(c.container, name, byteRange, nil)
	}

	if err != nil {
		return nil, err
	}
	return newBackoffReadCloser(c, reader), nil
}

func (c *microsoftClient) Delete(name string) error {
	return c.blobClient.DeleteBlob(c.container, name, nil)
}

func (c *microsoftClient) Walk(name string, fn func(name string) error) error {
	// See Azure docs for what `marker` does:
	// https://docs.microsoft.com/en-us/rest/api/storageservices/List-Blobs?redirectedfrom=MSDN
	var marker string
	for {
		blobList, err := c.blobClient.ListBlobs(c.container, storage.ListBlobsParameters{
			Prefix: name,
			Marker: marker,
		})
		if err != nil {
			return err
		}

		for _, file := range blobList.Blobs {
			if err := fn(file.Name); err != nil {
				return err
			}
		}

		// NextMarker is empty when all results have been returned
		if blobList.NextMarker == "" {
			break
		}
		marker = blobList.NextMarker
	}
	return nil
}

func (c *microsoftClient) Exists(name string) bool {
	exists, _ := c.blobClient.BlobExists(c.container, name)
	return exists
}

func (c *microsoftClient) IsRetryable(err error) (ret bool) {
	microsoftErr, ok := err.(storage.AzureStorageServiceError)
	if !ok {
		return false
	}
	return microsoftErr.StatusCode >= 500
}

func (c *microsoftClient) IsNotExist(err error) bool {
	microsoftErr, ok := err.(storage.AzureStorageServiceError)
	if !ok {
		return false
	}
	return microsoftErr.StatusCode == 404
}

func (c *microsoftClient) IsIgnorable(err error) bool {
	return false
}

type microsoftWriter struct {
	container  string
	blob       string
	blobClient storage.BlobStorageClient
}

func newMicrosoftWriter(client *microsoftClient, name string) (*microsoftWriter, error) {
	// create container
	_, err := client.blobClient.CreateContainerIfNotExists(client.container, storage.ContainerAccessTypePrivate)
	if err != nil {
		return nil, err
	}

	// create blob
	err = client.blobClient.CreateBlockBlob(client.container, name)
	if err != nil {
		return nil, err
	}

	return &microsoftWriter{
		container:  client.container,
		blob:       name,
		blobClient: client.blobClient,
	}, nil
}

func (w *microsoftWriter) Write(b []byte) (int, error) {
	blockList, err := w.blobClient.GetBlockList(w.container, w.blob, storage.BlockListTypeAll)
	if err != nil {
		return 0, err
	}

	blocksLen := len(blockList.CommittedBlocks)
	amendList := []storage.Block{}
	for _, v := range blockList.CommittedBlocks {
		amendList = append(amendList, storage.Block{v.Name, storage.BlockStatusCommitted})
	}

	inputSourceReader := bytes.NewReader(b)
	chunk := grpcutil.GetBuffer()
	defer grpcutil.PutBuffer(chunk)
	for {
		n, err := inputSourceReader.Read(chunk)
		if err == io.EOF {
			break
		}
		if err != nil {
			return 0, err
		}

		blockID := base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%011d\n", blocksLen)))
		data := chunk[:n]
		err = w.blobClient.PutBlock(w.container, w.blob, blockID, data)
		if err != nil {
			return 0, fmt.Errorf("BlobStorageClient.PutBlock: %v", err)
		}
		// add current uncommitted block to temporary block list
		amendList = append(amendList, storage.Block{blockID, storage.BlockStatusUncommitted})
		blocksLen++
	}

	// update block list to blob committed block list.
	err = w.blobClient.PutBlockList(w.container, w.blob, amendList)
	if err != nil {
		return 0, fmt.Errorf("BlobStorageClient.PutBlockList: %v", err)
	}
	return len(b), nil
}

func (w *microsoftWriter) Close() error {
	return nil
}
<commit_msg>Use append blobs instead of blobk blobs.<commit_after>package obj

import (
	"io"

	"github.com/Azure/azure-sdk-for-go/storage"
)

const (
	maxBlockSize = 4 * 1024 * 1024 // 4MB (according to: https://docs.microsoft.com/en-us/rest/api/storageservices/understanding-block-blobs--append-blobs--and-page-blobs#about-append-blobs)
)

type microsoftClient struct {
	blobClient storage.BlobStorageClient
	container  string
}

func newMicrosoftClient(container string, accountName string, accountKey string) (*microsoftClient, error) {
	client, err := storage.NewBasicClient(
		accountName,
		accountKey,
	)
	if err != nil {
		return nil, err
	}

	return &microsoftClient{
		blobClient: client.GetBlobService(),
		container:  container,
	}, nil
}

func (c *microsoftClient) Writer(name string) (io.WriteCloser, error) {
	writer, err := newMicrosoftWriter(c, name)
	if err != nil {
		return nil, err
	}
	return newBackoffWriteCloser(c, writer), nil
}

func (c *microsoftClient) Reader(name string, offset uint64, size uint64) (io.ReadCloser, error) {
	byteRange := byteRange(offset, size)
	var reader io.ReadCloser
	var err error
	if byteRange == "" {
		reader, err = c.blobClient.GetBlob(c.container, name)
	} else {
		reader, err = c.blobClient.GetBlobRange(c.container, name, byteRange, nil)
	}

	if err != nil {
		return nil, err
	}
	return newBackoffReadCloser(c, reader), nil
}

func (c *microsoftClient) Delete(name string) error {
	return c.blobClient.DeleteBlob(c.container, name, nil)
}

func (c *microsoftClient) Walk(name string, fn func(name string) error) error {
	// See Azure docs for what `marker` does:
	// https://docs.microsoft.com/en-us/rest/api/storageservices/List-Blobs?redirectedfrom=MSDN
	var marker string
	for {
		blobList, err := c.blobClient.ListBlobs(c.container, storage.ListBlobsParameters{
			Prefix: name,
			Marker: marker,
		})
		if err != nil {
			return err
		}

		for _, file := range blobList.Blobs {
			if err := fn(file.Name); err != nil {
				return err
			}
		}

		// NextMarker is empty when all results have been returned
		if blobList.NextMarker == "" {
			break
		}
		marker = blobList.NextMarker
	}
	return nil
}

func (c *microsoftClient) Exists(name string) bool {
	exists, _ := c.blobClient.BlobExists(c.container, name)
	return exists
}

func (c *microsoftClient) IsRetryable(err error) (ret bool) {
	microsoftErr, ok := err.(storage.AzureStorageServiceError)
	if !ok {
		return false
	}
	return microsoftErr.StatusCode >= 500
}

func (c *microsoftClient) IsNotExist(err error) bool {
	microsoftErr, ok := err.(storage.AzureStorageServiceError)
	if !ok {
		return false
	}
	return microsoftErr.StatusCode == 404
}

func (c *microsoftClient) IsIgnorable(err error) bool {
	return false
}

type microsoftWriter struct {
	container  string
	blob       string
	blobClient storage.BlobStorageClient
}

func newMicrosoftWriter(client *microsoftClient, name string) (*microsoftWriter, error) {
	// create container
	if _, err := client.blobClient.CreateContainerIfNotExists(client.container, storage.ContainerAccessTypePrivate); err != nil {
		return nil, err
	}

	// create blob
	if err := client.blobClient.PutAppendBlob(client.container, name, nil); err != nil {
		return nil, err
	}

	return &microsoftWriter{
		container:  client.container,
		blob:       name,
		blobClient: client.blobClient,
	}, nil
}

func (w *microsoftWriter) Write(b []byte) (int, error) {
	nBytes := 0
	for i := 0; i < len(b); i += maxBlockSize {
		var chunk []byte
		if i+maxBlockSize >= len(b) {
			chunk = b[i:]
		} else {
			chunk = b[i : i+maxBlockSize]
		}
		if err := w.blobClient.AppendBlock(w.container, w.blob, chunk, nil); err != nil {
			return nBytes, err
		}
		nBytes += len(chunk)
	}
	return nBytes, nil
}

func (w *microsoftWriter) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package chunk

import (
	"bytes"
	"compress/gzip"
	"context"
	"io"
	"path"

	"github.com/chmduquesne/rollinghash/buzhash64"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
	"github.com/pachyderm/pachyderm/src/server/pkg/storage/hash"
	"golang.org/x/sync/errgroup"
	"golang.org/x/sync/semaphore"
)

const (
	// MB is Megabytes.
	MB = 1024 * 1024
	// WindowSize is the size of the rolling hash window.
	WindowSize = 64
	// (bryce) this should be configurable.
	bufSize = 50 * MB
)

// initialWindow is the set of bytes used to initialize the window
// of the rolling hash function.
var initialWindow = make([]byte, WindowSize)

// WriterFunc is a callback that returns a data reference to the next chunk and the annotations within the chunk.
type WriterFunc func(*DataRef, []*Annotation) error

// byteSet is a unit of work for the workers.
// A worker will roll the rolling hash function across the bytes
// while processing the associated annotations.
type byteSet struct {
	data        []byte
	annotations []*Annotation
	// nextBytes is used for the edge case where no split point is found in an assigned byte set.
	// A worker that did not find a chunk in its assigned byte set will pass it to the prior
	// worker in the chain with nextBytes set to the next worker's bytes channel. This allows shuffling
	// of byte sets between workers until a split point is found.
	nextBytes chan *byteSet
}

// chanSet is a group of channels used to setup the daisy chain and shuffle data between the workers.
// How these channels are used by a worker depends on whether they are associated with
// the previous or next worker in the chain.
type chanSet struct {
	bytes chan *byteSet
	done  chan struct{}
}

type worker struct {
	ctx         context.Context
	objC        obj.Client
	hash        *buzhash64.Buzhash64
	splitMask   uint64
	first       bool
	buf         *bytes.Buffer
	annotations []*Annotation
	f           WriterFunc
	fs          []func() error
	prev, next  *chanSet
	stats       *stats
}

func (w *worker) run(byteSet *byteSet) error {
	// Roll through the assigned byte set.
	if err := w.rollByteSet(byteSet); err != nil {
		return err
	}
	// No split point found.
	if w.prev != nil && w.first {
		byteSet.nextBytes = w.next.bytes
		select {
		case w.prev.bytes <- byteSet:
		case <-w.ctx.Done():
			return w.ctx.Err()
		}
	} else {
		// Wait for the next byte set to roll.
		nextBytes := w.next.bytes
		for nextBytes != nil {
			select {
			case byteSet, more := <-nextBytes:
				// The next bytes channel is closed for the last worker,
				// so it uploads the last buffer as a chunk.
				if !more {
					if err := w.put(); err != nil {
						return err
					}
					nextBytes = nil
					break
				} else if err := w.rollByteSet(byteSet); err != nil {
					return err
				}
				nextBytes = byteSet.nextBytes
			case <-w.ctx.Done():
				return w.ctx.Err()
			}
		}
	}
	// Execute the writer function for the chunks that were found.
	return w.executeFuncs()
}

func (w *worker) rollByteSet(byteSet *byteSet) error {
	// Roll across the byte set.
	for i, a := range byteSet.annotations {
		var data []byte
		if i == len(byteSet.annotations)-1 {
			data = byteSet.data[a.Offset:len(byteSet.data)]
		} else {
			data = byteSet.data[a.Offset:byteSet.annotations[i+1].Offset]
		}
		// Convert from byte set offset to chunk offset.
		a.Offset = int64(w.buf.Len())
		w.annotations = joinAnnotations(w.annotations, a)
		if err := w.roll(data); err != nil {
			return err
		}
		// Reset hash between annotations.
		w.resetHash()
	}
	return nil
}

func joinAnnotations(as []*Annotation, a *Annotation) []*Annotation {
	// If the annotation being added is the same as the
	// last, then they are merged.
	if as != nil && as[len(as)-1].Meta == a.Meta {
		return as
	}
	return append(as, a)
}

func (w *worker) roll(data []byte) error {
	offset := 0
	for i, b := range data {
		w.hash.Roll(b)
		if w.hash.Sum64()&w.splitMask == 0 {
			w.buf.Write(data[offset : i+1])
			offset = i + 1
			if w.prev != nil && w.first {
				// We do not consider chunk split points within WindowSize bytes
				// of the start of the byte set.
				if w.buf.Len() < WindowSize {
					continue
				}
				byteSet := &byteSet{
					data:        w.buf.Bytes(),
					annotations: w.annotations,
				}
				w.prev.bytes <- byteSet
				w.first = false
			} else if err := w.put(); err != nil {
				return err
			}
			w.buf = &bytes.Buffer{}
			w.annotations = splitAnnotations(w.annotations)
		}
	}
	w.buf.Write(data[offset:])
	return nil
}

func (w *worker) put() error {
	chunk := &Chunk{Hash: hash.EncodeHash(hash.Sum(w.buf.Bytes()))}
	path := path.Join(prefix, chunk.Hash)
	// If the chunk does not exist, upload it.
	if !w.objC.Exists(w.ctx, path) {
		if err := w.upload(path); err != nil {
			return err
		}
	}
	chunkRef := &DataRef{
		Chunk:     chunk,
		SizeBytes: int64(w.buf.Len()),
	}
	// Update the annotations for the current chunk.
	updateAnnotations(chunkRef, w.buf.Bytes(), w.annotations)
	annotations := w.annotations
	w.fs = append(w.fs, func() error {
		return w.f(chunkRef, annotations)
	})
	return nil
}

func updateAnnotations(chunkRef *DataRef, buf []byte, annotations []*Annotation) {
	// (bryce) add check for no buf and greater than one annotation.
	// Fast path for next data reference being full chunk.
	if len(annotations) == 1 && annotations[0].NextDataRef != nil {
		annotations[0].NextDataRef = chunkRef
	}
	for i, a := range annotations {
		// (bryce) probably a better way to communicate whether to compute datarefs for an annotation.
		if a.NextDataRef != nil {
			a.NextDataRef.Chunk = chunkRef.Chunk
			var data []byte
			if i == len(annotations)-1 {
				data = buf[a.Offset:len(buf)]
			} else {
				data = buf[a.Offset:annotations[i+1].Offset]
			}
			a.NextDataRef.Hash = hash.EncodeHash(hash.Sum(data))
			a.NextDataRef.OffsetBytes = a.Offset
			a.NextDataRef.SizeBytes = int64(len(data))
		}
	}
}

func splitAnnotations(as []*Annotation) []*Annotation {
	if len(as) == 0 {
		return nil
	}
	// Copy the last annotation.
	lastA := as[len(as)-1]
	copyA := &Annotation{Meta: lastA.Meta}
	if lastA.NextDataRef != nil {
		copyA.NextDataRef = &DataRef{}
	}
	return []*Annotation{copyA}
}

func (w *worker) resetHash() {
	w.hash.Reset()
	w.hash.Write(initialWindow)
}

func (w *worker) upload(path string) error {
	objW, err := w.objC.Writer(w.ctx, path)
	if err != nil {
		return err
	}
	defer objW.Close()
	gzipW, err := gzip.NewWriterLevel(objW, gzip.BestSpeed)
	if err != nil {
		return err
	}
	defer gzipW.Close()
	// (bryce) Encrypt?
	_, err = io.Copy(gzipW, bytes.NewReader(w.buf.Bytes()))
	return err
}

func (w *worker) executeFuncs() error {
	// Wait for the prior worker in the chain to signal
	// that it is done.
	if w.prev != nil {
		select {
		case <-w.prev.done:
		case <-w.ctx.Done():
			return w.ctx.Err()
		}
	}
	// Execute the writer function for each chunk.
	for _, f := range w.fs {
		w.stats.chunkCount++
		if err := f(); err != nil {
			return err
		}
	}
	// Signal to the next worker in the chain that this worker is done.
	w.next.done <- struct{}{}
	return nil
}

type stats struct {
	chunkCount         int64
	annotatedBytesSize int64
}

// Writer splits a byte stream into content defined chunks that are hashed and deduplicated/uploaded to object storage.
// Chunk split points are determined by a bit pattern in a rolling hash function (buzhash64 at https://github.com/chmduquesne/rollinghash).
// The byte stream is split into byte sets for parallel processing. Workers roll the rolling hash function and perform the execution
// of the writer function on these byte sets. The workers are daisy chained such that split points across byte sets can be resolved by shuffling
// bytes between workers in the chain and the writer function is executed on the sequential ordering of the chunks in the byte stream.
type Writer struct {
	ctx, cancelCtx context.Context
	buf            *bytes.Buffer
	annotations    []*Annotation
	memoryLimiter  *semaphore.Weighted
	eg             *errgroup.Group
	newWorkerFunc  func(context.Context, *chanSet, *chanSet) *worker
	prev           *chanSet
	f              WriterFunc
	stats          *stats
}

func newWriter(ctx context.Context, objC obj.Client, memoryLimiter *semaphore.Weighted, averageBits int, f WriterFunc, seed int64) *Writer {
	stats := &stats{}
	newWorkerFunc := func(ctx context.Context, prev, next *chanSet) *worker {
		w := &worker{
			ctx:       ctx,
			objC:      objC,
			hash:      buzhash64.NewFromUint64Array(buzhash64.GenerateHashes(seed)),
			splitMask: (1 << uint64(averageBits)) - 1,
			first:     true,
			buf:       &bytes.Buffer{},
			prev:      prev,
			next:      next,
			f:         f,
			stats:     stats,
		}
		w.resetHash()
		return w
	}
	eg, cancelCtx := errgroup.WithContext(ctx)
	w := &Writer{
		ctx:           ctx,
		cancelCtx:     cancelCtx,
		buf:           &bytes.Buffer{},
		memoryLimiter: memoryLimiter,
		eg:            eg,
		newWorkerFunc: newWorkerFunc,
		f:             f,
		stats:         stats,
	}
	return w
}

// Annotate associates an annotation with the current byte set.
func (w *Writer) Annotate(a *Annotation) {
	a.Offset = int64(w.buf.Len())
	if a.Offset == 0 {
		w.annotations = nil
	}
	w.annotations = append(w.annotations, a)
	w.stats.annotatedBytesSize = 0
}

// AnnotatedBytesSize returns the size of the bytes for the current annotation.
func (w *Writer) AnnotatedBytesSize() int64 {
	return w.stats.annotatedBytesSize
}

// Flush flushes the buffered data.
func (w *Writer) Flush() error {
	// Write out the last buffer.
	if w.buf.Len() > 0 {
		w.writeByteSet()
	}
	// Signal to the last worker that it is last.
	if w.prev != nil {
		close(w.prev.bytes)
		w.prev = nil
	}
	// Wait for the workers to finish.
	if err := w.eg.Wait(); err != nil {
		return err
	}
	w.eg, w.cancelCtx = errgroup.WithContext(w.ctx)
	return nil
}

// Reset resets the buffer and annotations.
func (w *Writer) Reset() {
	// (bryce) should cancel all workers.
	w.buf = &bytes.Buffer{}
	w.annotations = nil
	w.stats.annotatedBytesSize = 0
}

// ChunkCount returns a count of the number of chunks created/referenced by
// the writer.
func (w *Writer) ChunkCount() int64 {
	return w.stats.chunkCount
}

// Write rolls through the data written, calling c.f when a chunk is found.
// Note: If making changes to this function, be wary of the performance
// implications (check before and after performance with chunker benchmarks).
func (w *Writer) Write(data []byte) (int, error) {
	var written int
	for w.buf.Len()+len(data) >= bufSize {
		i := bufSize - w.buf.Len()
		w.buf.Write(data[:i])
		w.writeByteSet()
		written += i
		data = data[i:]
	}
	w.buf.Write(data)
	written += len(data)
	w.stats.annotatedBytesSize += int64(written)
	return written, nil
}

func (w *Writer) writeByteSet() {
	w.memoryLimiter.Acquire(w.ctx, bufSize)
	prev := w.prev
	next := &chanSet{
		bytes: make(chan *byteSet, 1),
		done:  make(chan struct{}, 1),
	}
	byteSet := &byteSet{
		data:        w.buf.Bytes(),
		annotations: w.annotations,
	}
	w.eg.Go(func() error {
		defer w.memoryLimiter.Release(bufSize)
		return w.newWorkerFunc(w.cancelCtx, prev, next).run(byteSet)
	})
	w.prev = next
	w.buf = &bytes.Buffer{}
	w.annotations = splitAnnotations(w.annotations)
}

// Copy does a cheap copy from a reader to a writer.
func (w *Writer) Copy(r *Reader, n ...int64) error {
	c, err := r.ReadCopy(n...)
	if err != nil {
		return err
	}
	return w.WriteCopy(c)
}

// WriteCopy writes copy data to the writer.
func (w *Writer) WriteCopy(c *Copy) error {
	if _, err := io.Copy(w, c.before); err != nil {
		return err
	}
	for _, chunkRef := range c.chunkRefs {
		w.stats.chunkCount++
		// (bryce) might want to double check if this is correct.
		w.stats.annotatedBytesSize += chunkRef.SizeBytes
		updateAnnotations(chunkRef, nil, w.annotations)
		if err := w.f(chunkRef, w.annotations); err != nil {
			return err
		}
	}
	_, err := io.Copy(w, c.after)
	return err
}

// Close closes the writer.
func (w *Writer) Close() error {
	return w.Flush()
}
<commit_msg>Use channel closing to indicate a worker is done in the chain<commit_after>package chunk

import (
	"bytes"
	"compress/gzip"
	"context"
	"io"
	"path"

	"github.com/chmduquesne/rollinghash/buzhash64"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
	"github.com/pachyderm/pachyderm/src/server/pkg/storage/hash"
	"golang.org/x/sync/errgroup"
	"golang.org/x/sync/semaphore"
)

const (
	// MB is Megabytes.
	MB = 1024 * 1024
	// WindowSize is the size of the rolling hash window.
	WindowSize = 64
	// (bryce) this should be configurable.
	bufSize = 50 * MB
)

// initialWindow is the set of bytes used to initialize the window
// of the rolling hash function.
var initialWindow = make([]byte, WindowSize)

// WriterFunc is a callback that returns a data reference to the next chunk and the annotations within the chunk.
type WriterFunc func(*DataRef, []*Annotation) error

// byteSet is a unit of work for the workers.
// A worker will roll the rolling hash function across the bytes
// while processing the associated annotations.
type byteSet struct {
	data        []byte
	annotations []*Annotation
	// nextBytes is used for the edge case where no split point is found in an assigned byte set.
	// A worker that did not find a chunk in its assigned byte set will pass it to the prior
	// worker in the chain with nextBytes set to the next worker's bytes channel. This allows shuffling
	// of byte sets between workers until a split point is found.
	nextBytes chan *byteSet
}

// chanSet is a group of channels used to setup the daisy chain and shuffle data between the workers.
// How these channels are used by a worker depends on whether they are associated with
// the previous or next worker in the chain.
type chanSet struct {
	bytes chan *byteSet
	done  chan struct{}
}

type worker struct {
	ctx         context.Context
	objC        obj.Client
	hash        *buzhash64.Buzhash64
	splitMask   uint64
	first       bool
	buf         *bytes.Buffer
	annotations []*Annotation
	f           WriterFunc
	fs          []func() error
	prev, next  *chanSet
	stats       *stats
}

func (w *worker) run(byteSet *byteSet) error {
	// Roll through the assigned byte set.
	if err := w.rollByteSet(byteSet); err != nil {
		return err
	}
	// No split point found.
	if w.prev != nil && w.first {
		byteSet.nextBytes = w.next.bytes
		select {
		case w.prev.bytes <- byteSet:
		case <-w.ctx.Done():
			return w.ctx.Err()
		}
	} else {
		// Wait for the next byte set to roll.
		nextBytes := w.next.bytes
		for nextBytes != nil {
			select {
			case byteSet, more := <-nextBytes:
				// The next bytes channel is closed for the last worker,
				// so it uploads the last buffer as a chunk.
				if !more {
					if err := w.put(); err != nil {
						return err
					}
					nextBytes = nil
					break
				} else if err := w.rollByteSet(byteSet); err != nil {
					return err
				}
				nextBytes = byteSet.nextBytes
			case <-w.ctx.Done():
				return w.ctx.Err()
			}
		}
	}
	// Execute the writer function for the chunks that were found.
	return w.executeFuncs()
}

func (w *worker) rollByteSet(byteSet *byteSet) error {
	// Roll across the byte set.
	for i, a := range byteSet.annotations {
		var data []byte
		if i == len(byteSet.annotations)-1 {
			data = byteSet.data[a.Offset:len(byteSet.data)]
		} else {
			data = byteSet.data[a.Offset:byteSet.annotations[i+1].Offset]
		}
		// Convert from byte set offset to chunk offset.
		a.Offset = int64(w.buf.Len())
		w.annotations = joinAnnotations(w.annotations, a)
		if err := w.roll(data); err != nil {
			return err
		}
		// Reset hash between annotations.
		w.resetHash()
	}
	return nil
}

func joinAnnotations(as []*Annotation, a *Annotation) []*Annotation {
	// If the annotation being added is the same as the
	// last, then they are merged.
	if as != nil && as[len(as)-1].Meta == a.Meta {
		return as
	}
	return append(as, a)
}

func (w *worker) roll(data []byte) error {
	offset := 0
	for i, b := range data {
		w.hash.Roll(b)
		if w.hash.Sum64()&w.splitMask == 0 {
			w.buf.Write(data[offset : i+1])
			offset = i + 1
			if w.prev != nil && w.first {
				// We do not consider chunk split points within WindowSize bytes
				// of the start of the byte set.
				if w.buf.Len() < WindowSize {
					continue
				}
				byteSet := &byteSet{
					data:        w.buf.Bytes(),
					annotations: w.annotations,
				}
				w.prev.bytes <- byteSet
				w.first = false
			} else if err := w.put(); err != nil {
				return err
			}
			w.buf = &bytes.Buffer{}
			w.annotations = splitAnnotations(w.annotations)
		}
	}
	w.buf.Write(data[offset:])
	return nil
}

func (w *worker) put() error {
	chunk := &Chunk{Hash: hash.EncodeHash(hash.Sum(w.buf.Bytes()))}
	path := path.Join(prefix, chunk.Hash)
	// If the chunk does not exist, upload it.
	if !w.objC.Exists(w.ctx, path) {
		if err := w.upload(path); err != nil {
			return err
		}
	}
	chunkRef := &DataRef{
		Chunk:     chunk,
		SizeBytes: int64(w.buf.Len()),
	}
	// Update the annotations for the current chunk.
	updateAnnotations(chunkRef, w.buf.Bytes(), w.annotations)
	annotations := w.annotations
	w.fs = append(w.fs, func() error {
		return w.f(chunkRef, annotations)
	})
	return nil
}

func updateAnnotations(chunkRef *DataRef, buf []byte, annotations []*Annotation) {
	// (bryce) add check for no buf and greater than one annotation.
	// Fast path for next data reference being full chunk.
	if len(annotations) == 1 && annotations[0].NextDataRef != nil {
		annotations[0].NextDataRef = chunkRef
	}
	for i, a := range annotations {
		// (bryce) probably a better way to communicate whether to compute datarefs for an annotation.
		if a.NextDataRef != nil {
			a.NextDataRef.Chunk = chunkRef.Chunk
			var data []byte
			if i == len(annotations)-1 {
				data = buf[a.Offset:len(buf)]
			} else {
				data = buf[a.Offset:annotations[i+1].Offset]
			}
			a.NextDataRef.Hash = hash.EncodeHash(hash.Sum(data))
			a.NextDataRef.OffsetBytes = a.Offset
			a.NextDataRef.SizeBytes = int64(len(data))
		}
	}
}

func splitAnnotations(as []*Annotation) []*Annotation {
	if len(as) == 0 {
		return nil
	}
	// Copy the last annotation.
	lastA := as[len(as)-1]
	copyA := &Annotation{Meta: lastA.Meta}
	if lastA.NextDataRef != nil {
		copyA.NextDataRef = &DataRef{}
	}
	return []*Annotation{copyA}
}

func (w *worker) resetHash() {
	w.hash.Reset()
	w.hash.Write(initialWindow)
}

func (w *worker) upload(path string) error {
	objW, err := w.objC.Writer(w.ctx, path)
	if err != nil {
		return err
	}
	defer objW.Close()
	gzipW, err := gzip.NewWriterLevel(objW, gzip.BestSpeed)
	if err != nil {
		return err
	}
	defer gzipW.Close()
	// (bryce) Encrypt?
	_, err = io.Copy(gzipW, bytes.NewReader(w.buf.Bytes()))
	return err
}

func (w *worker) executeFuncs() error {
	// Wait for the prior worker in the chain to signal
	// that it is done.
	if w.prev != nil {
		select {
		case <-w.prev.done:
		case <-w.ctx.Done():
			return w.ctx.Err()
		}
	}
	// Execute the writer function for each chunk.
	for _, f := range w.fs {
		w.stats.chunkCount++
		if err := f(); err != nil {
			return err
		}
	}
	// Signal to the next worker in the chain that this worker is done.
	close(w.next.done)
	return nil
}

type stats struct {
	chunkCount         int64
	annotatedBytesSize int64
}

// Writer splits a byte stream into content defined chunks that are hashed and deduplicated/uploaded to object storage.
// Chunk split points are determined by a bit pattern in a rolling hash function (buzhash64 at https://github.com/chmduquesne/rollinghash).
// The byte stream is split into byte sets for parallel processing. Workers roll the rolling hash function and perform the execution
// of the writer function on these byte sets. The workers are daisy chained such that split points across byte sets can be resolved by shuffling
// bytes between workers in the chain and the writer function is executed on the sequential ordering of the chunks in the byte stream.
type Writer struct {
	ctx, cancelCtx context.Context
	buf            *bytes.Buffer
	annotations    []*Annotation
	memoryLimiter  *semaphore.Weighted
	eg             *errgroup.Group
	newWorkerFunc  func(context.Context, *chanSet, *chanSet) *worker
	prev           *chanSet
	f              WriterFunc
	stats          *stats
}

func newWriter(ctx context.Context, objC obj.Client, memoryLimiter *semaphore.Weighted, averageBits int, f WriterFunc, seed int64) *Writer {
	stats := &stats{}
	newWorkerFunc := func(ctx context.Context, prev, next *chanSet) *worker {
		w := &worker{
			ctx:       ctx,
			objC:      objC,
			hash:      buzhash64.NewFromUint64Array(buzhash64.GenerateHashes(seed)),
			splitMask: (1 << uint64(averageBits)) - 1,
			first:     true,
			buf:       &bytes.Buffer{},
			prev:      prev,
			next:      next,
			f:         f,
			stats:     stats,
		}
		w.resetHash()
		return w
	}
	eg, cancelCtx := errgroup.WithContext(ctx)
	w := &Writer{
		ctx:           ctx,
		cancelCtx:     cancelCtx,
		buf:           &bytes.Buffer{},
		memoryLimiter: memoryLimiter,
		eg:            eg,
		newWorkerFunc: newWorkerFunc,
		f:             f,
		stats:         stats,
	}
	return w
}

// Annotate associates an annotation with the current byte set.
func (w *Writer) Annotate(a *Annotation) {
	a.Offset = int64(w.buf.Len())
	if a.Offset == 0 {
		w.annotations = nil
	}
	w.annotations = append(w.annotations, a)
	w.stats.annotatedBytesSize = 0
}

// AnnotatedBytesSize returns the size of the bytes for the current annotation.
func (w *Writer) AnnotatedBytesSize() int64 {
	return w.stats.annotatedBytesSize
}

// Flush flushes the buffered data.
func (w *Writer) Flush() error {
	// Write out the last buffer.
	if w.buf.Len() > 0 {
		w.writeByteSet()
	}
	// Signal to the last worker that it is last.
	if w.prev != nil {
		close(w.prev.bytes)
		w.prev = nil
	}
	// Wait for the workers to finish.
	if err := w.eg.Wait(); err != nil {
		return err
	}
	w.eg, w.cancelCtx = errgroup.WithContext(w.ctx)
	return nil
}

// Reset resets the buffer and annotations.
func (w *Writer) Reset() {
	// (bryce) should cancel all workers.
	w.buf = &bytes.Buffer{}
	w.annotations = nil
	w.stats.annotatedBytesSize = 0
}

// ChunkCount returns a count of the number of chunks created/referenced by
// the writer.
func (w *Writer) ChunkCount() int64 {
	return w.stats.chunkCount
}

// Write rolls through the data written, calling c.f when a chunk is found.
// Note: If making changes to this function, be wary of the performance
// implications (check before and after performance with chunker benchmarks).
func (w *Writer) Write(data []byte) (int, error) {
	var written int
	for w.buf.Len()+len(data) >= bufSize {
		i := bufSize - w.buf.Len()
		w.buf.Write(data[:i])
		w.writeByteSet()
		written += i
		data = data[i:]
	}
	w.buf.Write(data)
	written += len(data)
	w.stats.annotatedBytesSize += int64(written)
	return written, nil
}

func (w *Writer) writeByteSet() {
	w.memoryLimiter.Acquire(w.ctx, bufSize)
	prev := w.prev
	next := &chanSet{
		bytes: make(chan *byteSet, 1),
		done:  make(chan struct{}),
	}
	byteSet := &byteSet{
		data:        w.buf.Bytes(),
		annotations: w.annotations,
	}
	w.eg.Go(func() error {
		defer w.memoryLimiter.Release(bufSize)
		return w.newWorkerFunc(w.cancelCtx, prev, next).run(byteSet)
	})
	w.prev = next
	w.buf = &bytes.Buffer{}
	w.annotations = splitAnnotations(w.annotations)
}

// Copy does a cheap copy from a reader to a writer.
func (w *Writer) Copy(r *Reader, n ...int64) error {
	c, err := r.ReadCopy(n...)
	if err != nil {
		return err
	}
	return w.WriteCopy(c)
}

// WriteCopy writes copy data to the writer.
func (w *Writer) WriteCopy(c *Copy) error {
	if _, err := io.Copy(w, c.before); err != nil {
		return err
	}
	for _, chunkRef := range c.chunkRefs {
		w.stats.chunkCount++
		// (bryce) might want to double check if this is correct.
		w.stats.annotatedBytesSize += chunkRef.SizeBytes
		updateAnnotations(chunkRef, nil, w.annotations)
		if err := w.f(chunkRef, w.annotations); err != nil {
			return err
		}
	}
	_, err := io.Copy(w, c.after)
	return err
}

// Close closes the writer.
func (w *Writer) Close() error {
	return w.Flush()
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 The Syncthing Authors.
//
// This program is free software: you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the Free
// Software Foundation, either version 3 of the License, or (at your option)
// any later version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
// more details.
//
// You should have received a copy of the GNU General Public License along
// with this program. If not, see <http://www.gnu.org/licenses/>.

// +build integration

package integration

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"time"
)

var env = []string{
	"HOME=.",
	"STGUIAPIKEY=" + apiKey,
	"STNORESTART=1",
}

type syncthingProcess struct {
	instance  string
	argv      []string
	port      int
	apiKey    string
	csrfToken string
	lastEvent int

	cmd   *exec.Cmd
	logfd *os.File
}

func (p *syncthingProcess) start() error {
	if p.logfd == nil {
		logfd, err := os.Create("logs/" + getTestName() + "-" + p.instance + ".out")
		if err != nil {
			return err
		}
		p.logfd = logfd
	}

	binary := "../bin/syncthing"

	// We check to see if there's an instance specific binary we should run,
	// for example if we are running integration tests between different
	// versions. If there isn't, we just go with the default.
	if _, err := os.Stat(binary + "-" + p.instance); err == nil {
		binary = binary + "-" + p.instance
	}
	if _, err := os.Stat(binary + "-" + p.instance + ".exe"); err == nil {
		binary = binary + "-" + p.instance + ".exe"
	}

	cmd := exec.Command(binary, p.argv...)
	cmd.Stdout = p.logfd
	cmd.Stderr = p.logfd
	cmd.Env = append(os.Environ(), env...)

	err := cmd.Start()
	if err != nil {
		return err
	}
	p.cmd = cmd

	for {
		resp, err := p.get("/")
		if err == nil {
			resp.Body.Close()
			return nil
		}
		time.Sleep(250 * time.Millisecond)
	}
}

func (p *syncthingProcess) stop() error {
	p.cmd.Process.Signal(os.Kill)
	p.cmd.Wait()

	fd, err := os.Open(p.instance + ".out")
	if err != nil {
		return err
	}
	defer fd.Close()

	raceConditionStart := []byte("WARNING: DATA RACE")
	raceConditionSep := []byte("==================")
	sc := bufio.NewScanner(fd)
	race := false
	for sc.Scan() {
		line := sc.Bytes()
		if race {
			fmt.Printf("%s\n", line)
			if bytes.Contains(line, raceConditionSep) {
				race = false
			}
		} else if bytes.Contains(line, raceConditionStart) {
			fmt.Printf("%s\n", raceConditionSep)
			fmt.Printf("%s\n", raceConditionStart)
			race = true
			if err == nil {
				err = errors.New("Race condition detected")
			}
		}
	}
	return err
}

func (p *syncthingProcess) get(path string) (*http.Response, error) {
	client := &http.Client{
		Timeout: 30 * time.Second,
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	req, err := http.NewRequest("GET", fmt.Sprintf("http://127.0.0.1:%d%s", p.port, path), nil)
	if err != nil {
		return nil, err
	}
	if p.apiKey != "" {
		req.Header.Add("X-API-Key", p.apiKey)
	}
	if p.csrfToken != "" {
		req.Header.Add("X-CSRF-Token", p.csrfToken)
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

func (p *syncthingProcess) post(path string, data io.Reader) (*http.Response, error) {
	client := &http.Client{
		Timeout: 600 * time.Second,
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	req, err := http.NewRequest("POST", fmt.Sprintf("http://127.0.0.1:%d%s", p.port, path), data)
	if err != nil {
		return nil, err
	}
	if p.apiKey != "" {
		req.Header.Add("X-API-Key", p.apiKey)
	}
	if p.csrfToken != "" {
		req.Header.Add("X-CSRF-Token", p.csrfToken)
	}
	req.Header.Add("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

func (p *syncthingProcess) peerCompletion() (map[string]int, error) {
	resp, err := p.get("/rest/debug/peerCompletion")
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	comp := map[string]int{}
	err = json.NewDecoder(resp.Body).Decode(&comp)
	return comp, err
}

type model struct {
	GlobalBytes   int
	GlobalDeleted int
	GlobalFiles   int
	InSyncBytes   int
	InSyncFiles   int
	Invalid       string
	LocalBytes    int
	LocalDeleted  int
	LocalFiles    int
	NeedBytes     int
	NeedFiles     int
	State         string
	StateChanged  time.Time
	Version       int
}

func (p *syncthingProcess) model(folder string) (model, error) {
	resp, err := p.get("/rest/model?folder=" + folder)
	if err != nil {
		return model{}, err
	}

	var res model
	err = json.NewDecoder(resp.Body).Decode(&res)
	if err != nil {
		return model{}, err
	}

	return res, nil
}

type event struct {
	ID   int
	Time time.Time
	Type string
	Data interface{}
}

func (p *syncthingProcess) events() ([]event, error) {
	resp, err := p.get(fmt.Sprintf("/rest/events?since=%d", p.lastEvent))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var evs []event
	err = json.NewDecoder(resp.Body).Decode(&evs)
	if err != nil {
		return nil, err
	}
	p.lastEvent = evs[len(evs)-1].ID
	return evs, err
}

type versionResp struct {
	Version string
}

func (p *syncthingProcess) version() (string, error) {
	resp, err := p.get("/rest/version")
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	var v versionResp
	err = json.NewDecoder(resp.Body).Decode(&v)
	if err != nil {
		return "", err
	}
	return v.Version, nil
}
<commit_msg>Fix integration tests<commit_after>// Copyright (C) 2014 The Syncthing Authors.
//
// This program is free software: you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the Free
// Software Foundation, either version 3 of the License, or (at your option)
// any later version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
// more details.
//
// You should have received a copy of the GNU General Public License along
// with this program. If not, see <http://www.gnu.org/licenses/>.

// +build integration

package integration

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"time"
)

var env = []string{
	"HOME=.",
	"STGUIAPIKEY=" + apiKey,
	"STNORESTART=1",
}

type syncthingProcess struct {
	instance  string
	argv      []string
	port      int
	apiKey    string
	csrfToken string
	lastEvent int

	cmd   *exec.Cmd
	logfd *os.File
}

func (p *syncthingProcess) start() error {
	if p.logfd == nil {
		logfd, err := os.Create("logs/" + getTestName() + "-" + p.instance + ".out")
		if err != nil {
			return err
		}
		p.logfd = logfd
	}

	binary := "../bin/syncthing"

	// We check to see if there's an instance specific binary we should run,
	// for example if we are running integration tests between different
	// versions. If there isn't, we just go with the default.
	if _, err := os.Stat(binary + "-" + p.instance); err == nil {
		binary = binary + "-" + p.instance
	}
	if _, err := os.Stat(binary + "-" + p.instance + ".exe"); err == nil {
		binary = binary + "-" + p.instance + ".exe"
	}

	cmd := exec.Command(binary, p.argv...)
	cmd.Stdout = p.logfd
	cmd.Stderr = p.logfd
	cmd.Env = append(os.Environ(), env...)

	err := cmd.Start()
	if err != nil {
		return err
	}
	p.cmd = cmd

	for {
		resp, err := p.get("/")
		if err == nil {
			resp.Body.Close()
			return nil
		}
		time.Sleep(250 * time.Millisecond)
	}
}

func (p *syncthingProcess) stop() error {
	p.cmd.Process.Signal(os.Kill)
	p.cmd.Wait()

	fd, err := os.Open(p.logfd.Name())
	if err != nil {
		return err
	}
	defer fd.Close()

	raceConditionStart := []byte("WARNING: DATA RACE")
	raceConditionSep := []byte("==================")
	sc := bufio.NewScanner(fd)
	race := false
	for sc.Scan() {
		line := sc.Bytes()
		if race {
			fmt.Printf("%s\n", line)
			if bytes.Contains(line, raceConditionSep) {
				race = false
			}
		} else if bytes.Contains(line, raceConditionStart) {
			fmt.Printf("%s\n", raceConditionSep)
			fmt.Printf("%s\n", raceConditionStart)
			race = true
			if err == nil {
				err = errors.New("Race condition detected")
			}
		}
	}
	return err
}

func (p *syncthingProcess) get(path string) (*http.Response, error) {
	client := &http.Client{
		Timeout: 30 * time.Second,
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	req, err := http.NewRequest("GET", fmt.Sprintf("http://127.0.0.1:%d%s", p.port, path), nil)
	if err != nil {
		return nil, err
	}
	if p.apiKey != "" {
		req.Header.Add("X-API-Key", p.apiKey)
	}
	if p.csrfToken != "" {
		req.Header.Add("X-CSRF-Token", p.csrfToken)
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

func (p *syncthingProcess) post(path string, data io.Reader) (*http.Response, error) {
	client := &http.Client{
		Timeout: 600 * time.Second,
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	req, err := http.NewRequest("POST", fmt.Sprintf("http://127.0.0.1:%d%s", p.port, path), data)
	if err != nil {
		return nil, err
	}
	if p.apiKey != "" {
		req.Header.Add("X-API-Key", p.apiKey)
	}
	if p.csrfToken != "" {
		req.Header.Add("X-CSRF-Token", p.csrfToken)
	}
	req.Header.Add("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

func (p *syncthingProcess) peerCompletion() (map[string]int, error) {
	resp, err := p.get("/rest/debug/peerCompletion")
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	comp := map[string]int{}
	err = json.NewDecoder(resp.Body).Decode(&comp)
	return comp, err
}

type model struct {
	GlobalBytes   int
	GlobalDeleted int
	GlobalFiles   int
	InSyncBytes   int
	InSyncFiles   int
	Invalid       string
	LocalBytes    int
	LocalDeleted  int
	LocalFiles    int
	NeedBytes     int
	NeedFiles     int
	State         string
	StateChanged  time.Time
	Version       int
}

func (p *syncthingProcess) model(folder string) (model, error) {
	resp, err := p.get("/rest/model?folder=" + folder)
	if err != nil {
		return model{}, err
	}

	var res model
	err = json.NewDecoder(resp.Body).Decode(&res)
	if err != nil {
		return model{}, err
	}

	return res, nil
}

type event struct {
	ID   int
	Time time.Time
	Type string
	Data interface{}
}

func (p *syncthingProcess) events() ([]event, error) {
	resp, err := p.get(fmt.Sprintf("/rest/events?since=%d", p.lastEvent))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var evs []event
	err = json.NewDecoder(resp.Body).Decode(&evs)
	if err != nil {
		return nil, err
	}
	p.lastEvent = evs[len(evs)-1].ID
	return evs, err
}

type versionResp struct {
	Version string
}

func (p *syncthingProcess) version() (string, error) {
	resp, err := p.get("/rest/version")
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	var v versionResp
	err = json.NewDecoder(resp.Body).Decode(&v)
	if err != nil {
		return "", err
	}
	return v.Version, nil
}
<|endoftext|>


<commit_before>package parser

import (
	"archive/zip"
	"encoding/csv"
	"errors"
	"io"
	"log"
	"net"
	"regexp"
	"strconv"

	"github.com/m-lab/annotation-service/loader"
)

const (
	ipNumColumnsGlite2        = 10
	locationNumColumnsGlite2  = 13
	gLite2Prefix              = "GeoLite2-City"
	geoLite2BlocksFilenameIP4 = "GeoLite2-City-Blocks-IPv4.csv"  // Filename of ipv4 blocks file
	geoLite2BlocksFilenameIP6 = "GeoLite2-City-Blocks-IPv6.csv"  // Filename of ipv6 blocks file
	geoLite2LocationsFilename = "GeoLite2-City-Locations-en.csv" // Filename of locations file
)

func LoadGeoLite2(zip *zip.Reader) (*GeoDataset, error) {
	locations, err := loader.FindFile(geoLite2LocationsFilename, zip)
	if err != nil {
		return nil, err
	}
	// geoidMap is just a temporary map that will be discarded once the blocks are parsed
	locationNode, geoidMap, err := LoadLocListGLite2(locations)
	if err != nil {
		return nil, err
	}
	blocks4, err := loader.FindFile(geoLite2BlocksFilenameIP4, zip)
	if err != nil {
		return nil, err
	}
	ipNodes4, err := LoadIPListGLite2(blocks4, geoidMap)
	if err != nil {
		return nil, err
	}
	blocks6, err := loader.FindFile(geoLite2BlocksFilenameIP6, zip)
	if err != nil {
		return nil, err
	}
	ipNodes6, err := LoadIPListGLite2(blocks6, geoidMap)
	if err != nil {
		return nil, err
	}
	return &GeoDataset{IP4Nodes: ipNodes4, IP6Nodes: ipNodes6, LocationNodes: locationNode}, nil
}

// Finds the smallest and largest net.IP from a CIDR range
// Example: "1.0.0.0/24" -> 1.0.0.0 , 1.0.0.255
func rangeCIDR(cidr string) (net.IP, net.IP, error) {
	ip, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		return nil, nil, errors.New("Invalid CIDR IP range")
	}
	lowIp := make(net.IP, len(ip))
	copy(lowIp, ip)
	mask := ipnet.Mask
	for x, _ := range ip {
		if len(mask) == 4 {
			if x < 12 {
				ip[x] |= 0
			} else {
				ip[x] |= ^mask[x-12]
			}
		} else {
			ip[x] |= ^mask[x]
		}
	}
	return lowIp, ip, nil
}

// Create Location list for GLite2 databases
// TODO This code is a bit fragile.  Should probably parse the header and
// use that to guide the parsing of the rows.
func LoadLocListGLite2(reader io.Reader) ([]LocationNode, map[int]int, error) {
	idMap := make(map[int]int, mapMax)
	list := []LocationNode{}
	r := csv.NewReader(reader)
	// Skip the first line
	_, err := r.Read()
	if err == io.EOF {
		log.Println("Empty input data")
		return nil, nil, errors.New("Empty input data")
	}
	// FieldsPerRecord is the expected column length
	r.FieldsPerRecord = locationNumColumnsGlite2
	for {
		record, err := r.Read()
		if err != nil {
			if err == io.EOF {
				break
			}
			if len(record) != r.FieldsPerRecord {
				log.Println("Incorrect number of columns in IP list got: ", len(record), " wanted: ", r.FieldsPerRecord)
				log.Println(record)
				return nil, nil, errors.New("Corrupted Data: wrong number of columns")

			} else {
				log.Println(err, ": ", record)
				return nil, nil, err
			}
		}
		var lNode LocationNode
		lNode.GeonameID, err = strconv.Atoi(record[0])
		if err != nil {
			if len(record[0]) > 0 {
				log.Println("GeonameID should be a number ", record[0])
				return nil, nil, errors.New("Corrupted Data: GeonameID should be a number")
			}
		}
		lNode.ContinentCode, err = checkCaps(record[2], "Continent code")
		if err != nil {
			return nil, nil, err
		}
		lNode.CountryCode, err = checkCaps(record[4], "Country code")
		if err != nil {
			return nil, nil, err
		}
		match, _ := regexp.MatchString(`^[^0-9]*$`, record[5])
		if match {
			lNode.CountryName = record[5]
		} else {
			log.Println("Country name should be letters only : ", record[5])
			return nil, nil, errors.New("Corrupted Data: country name should be letters")
		}
		// TODO - should probably do some validation.
		lNode.RegionCode = record[6]
		lNode.RegionName = record[7]
		lNode.MetroCode, err = strconv.ParseInt(record[11], 10, 64)
		if err != nil {
			if len(record[11]) > 0 {
				log.Println("MetroCode should be a number")
				return nil, nil, errors.New("Corrupted Data: metrocode should be a number")
			}
		}
		lNode.CityName = record[10]
		list = append(list, lNode)
		idMap[lNode.GeonameID] = len(list) - 1
	}
	return list, idMap, nil
}

// LoadIPListGLite2 creates a List of IPNodes from a GeoLite2 reader.
// TODO(gfr) Update to use recursion instead of stack.
func LoadIPListGLite2(reader io.Reader, idMap map[int]int) ([]IPNode, error) {
	list := []IPNode{}
	r := csv.NewReader(reader)
	stack := []IPNode{}
	// Skip first line
	_, err := r.Read()
	if err == io.EOF {
		log.Println("Empty input data")
		return nil, errors.New("Empty input data")
	}
	for {
		var newNode IPNode
		// Example:
		// GLite2 : record = [2a04:97c0::/29,2658434,2658434,0,0,47,8,100]
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		err = checkNumColumns(record, ipNumColumnsGlite2)
		if err != nil {
			return nil, err
		}
		lowIp, highIp, err := rangeCIDR(record[0])
		if err != nil {
			return nil, err
		}
		newNode.IPAddressLow = lowIp
		newNode.IPAddressHigh = highIp
		// Look for GeoId within idMap and return index
		index, err := lookupGeoId(record[1], idMap)
		if err != nil {
			if backupIndex, err := lookupGeoId(record[2], idMap); err == nil {
				index = backupIndex
			} else {
				log.Println("Couldn't get a valid Geoname id!", record)
				//TODO: Add a prometheus metric here
			}

		}
		newNode.LocationIndex = index
		newNode.PostalCode = record[6]
		newNode.Latitude, err = stringToFloat(record[7], "Latitude")
		if err != nil {
			return nil, err
		}
		newNode.Longitude, err = stringToFloat(record[8], "Longitude")
		if err != nil {
			return nil, err
		}
		stack, list = handleStack(stack, list, newNode)
	}
	var pop IPNode
	pop, stack = stack[len(stack)-1], stack[:len(stack)-1]
	for ; len(stack) > 0; pop, stack = stack[len(stack)-1], stack[:len(stack)-1] {
		peek := stack[len(stack)-1]
		peek.IPAddressLow = PlusOne(pop.IPAddressHigh)
		list = append(list, peek)
	}
	return list, nil
}
<commit_msg>allow record length >= 13<commit_after>package parser

import (
	"archive/zip"
	"encoding/csv"
	"errors"
	"io"
	"log"
	"net"
	"regexp"
	"strconv"

	"github.com/m-lab/annotation-service/loader"
)

const (
	ipNumColumnsGlite2        = 10
	locationNumColumnsGlite2  = 13
	gLite2Prefix              = "GeoLite2-City"
	geoLite2BlocksFilenameIP4 = "GeoLite2-City-Blocks-IPv4.csv"  // Filename of ipv4 blocks file
	geoLite2BlocksFilenameIP6 = "GeoLite2-City-Blocks-IPv6.csv"  // Filename of ipv6 blocks file
	geoLite2LocationsFilename = "GeoLite2-City-Locations-en.csv" // Filename of locations file
)

func LoadGeoLite2(zip *zip.Reader) (*GeoDataset, error) {
	locations, err := loader.FindFile(geoLite2LocationsFilename, zip)
	if err != nil {
		return nil, err
	}
	// geoidMap is just a temporary map that will be discarded once the blocks are parsed
	locationNode, geoidMap, err := LoadLocListGLite2(locations)
	if err != nil {
		return nil, err
	}
	blocks4, err := loader.FindFile(geoLite2BlocksFilenameIP4, zip)
	if err != nil {
		return nil, err
	}
	ipNodes4, err := LoadIPListGLite2(blocks4, geoidMap)
	if err != nil {
		return nil, err
	}
	blocks6, err := loader.FindFile(geoLite2BlocksFilenameIP6, zip)
	if err != nil {
		return nil, err
	}
	ipNodes6, err := LoadIPListGLite2(blocks6, geoidMap)
	if err != nil {
		return nil, err
	}
	return &GeoDataset{IP4Nodes: ipNodes4, IP6Nodes: ipNodes6, LocationNodes: locationNode}, nil
}

// Finds the smallest and largest net.IP from a CIDR range
// Example: "1.0.0.0/24" -> 1.0.0.0 , 1.0.0.255
func rangeCIDR(cidr string) (net.IP, net.IP, error) {
	ip, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		return nil, nil, errors.New("Invalid CIDR IP range")
	}
	lowIp := make(net.IP, len(ip))
	copy(lowIp, ip)
	mask := ipnet.Mask
	for x, _ := range ip {
		if len(mask) == 4 {
			if x < 12 {
				ip[x] |= 0
			} else {
				ip[x] |= ^mask[x-12]
			}
		} else {
			ip[x] |= ^mask[x]
		}
	}
	return lowIp, ip, nil
}

// Create Location list for GLite2 databases
// TODO This code is a bit fragile.  Should probably parse the header and
// use that to guide the parsing of the rows.
func LoadLocListGLite2(reader io.Reader) ([]LocationNode, map[int]int, error) {
	idMap := make(map[int]int, mapMax)
	list := []LocationNode{}
	r := csv.NewReader(reader)
	// Skip the first line
	// TODO - we should parse the first line, instead of skipping it!!
	// This should set r.FieldsPerRecord.
	first, err := r.Read()
	if err == io.EOF {
		log.Println("Empty input data")
		return nil, nil, errors.New("Empty input data")
	}
	// TODO - this is a bit hacky.  May want to improve it.
	// Older geoLite2 have 13 columns, but since 2018/03, they have 14 columns.
	// This will print a log every time it loads a newer location file.
	if len(first) != locationNumColumnsGlite2 {
		log.Println("Incorrect number of columns in header, got: ", len(first), " wanted: ", locationNumColumnsGlite2)
		log.Println(first)
		if len(first) < locationNumColumnsGlite2 {
			return nil, nil, errors.New("Corrupted Data: wrong number of columns")
		}
	}
	// FieldsPerRecord is the expected column length
	// r.FieldsPerRecord = locationNumColumnsGlite2
	for {
		record, err := r.Read()
		if err != nil {
			if err == io.EOF {
				break
			}
			if len(record) != r.FieldsPerRecord {
				log.Println("Incorrect number of columns in IP list got: ", len(record), " wanted: ", r.FieldsPerRecord)
				log.Println(record)
				return nil, nil, errors.New("Corrupted Data: wrong number of columns")

			} else {
				log.Println(err, ": ", record)
				return nil, nil, err
			}
		}
		var lNode LocationNode
		lNode.GeonameID, err = strconv.Atoi(record[0])
		if err != nil {
			if len(record[0]) > 0 {
				log.Println("GeonameID should be a number ", record[0])
				return nil, nil, errors.New("Corrupted Data: GeonameID should be a number")
			}
		}
		lNode.ContinentCode, err = checkCaps(record[2], "Continent code")
		if err != nil {
			return nil, nil, err
		}
		lNode.CountryCode, err = checkCaps(record[4], "Country code")
		if err != nil {
			return nil, nil, err
		}
		match, _ := regexp.MatchString(`^[^0-9]*$`, record[5])
		if match {
			lNode.CountryName = record[5]
		} else {
			log.Println("Country name should be letters only : ", record[5])
			return nil, nil, errors.New("Corrupted Data: country name should be letters")
		}
		// TODO - should probably do some validation.
		lNode.RegionCode = record[6]
		lNode.RegionName = record[7]
		lNode.MetroCode, err = strconv.ParseInt(record[11], 10, 64)
		if err != nil {
			if len(record[11]) > 0 {
				log.Println("MetroCode should be a number")
				return nil, nil, errors.New("Corrupted Data: metrocode should be a number")
			}
		}
		lNode.CityName = record[10]
		list = append(list, lNode)
		idMap[lNode.GeonameID] = len(list) - 1
	}
	return list, idMap, nil
}

// LoadIPListGLite2 creates a List of IPNodes from a GeoLite2 reader.
// TODO(gfr) Update to use recursion instead of stack.
func LoadIPListGLite2(reader io.Reader, idMap map[int]int) ([]IPNode, error) {
	list := []IPNode{}
	r := csv.NewReader(reader)
	stack := []IPNode{}
	// Skip first line
	_, err := r.Read()
	if err == io.EOF {
		log.Println("Empty input data")
		return nil, errors.New("Empty input data")
	}
	for {
		var newNode IPNode
		// Example:
		// GLite2 : record = [2a04:97c0::/29,2658434,2658434,0,0,47,8,100]
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		err = checkNumColumns(record, ipNumColumnsGlite2)
		if err != nil {
			return nil, err
		}
		lowIp, highIp, err := rangeCIDR(record[0])
		if err != nil {
			return nil, err
		}
		newNode.IPAddressLow = lowIp
		newNode.IPAddressHigh = highIp
		// Look for GeoId within idMap and return index
		index, err := lookupGeoId(record[1], idMap)
		if err != nil {
			if backupIndex, err := lookupGeoId(record[2], idMap); err == nil {
				index = backupIndex
			} else {
				log.Println("Couldn't get a valid Geoname id!", record)
				//TODO: Add a prometheus metric here
			}

		}
		newNode.LocationIndex = index
		newNode.PostalCode = record[6]
		newNode.Latitude, err = stringToFloat(record[7], "Latitude")
		if err != nil {
			return nil, err
		}
		newNode.Longitude, err = stringToFloat(record[8], "Longitude")
		if err != nil {
			return nil, err
		}
		stack, list = handleStack(stack, list, newNode)
	}
	var pop IPNode
	pop, stack = stack[len(stack)-1], stack[:len(stack)-1]
	for ; len(stack) > 0; pop, stack = stack[len(stack)-1], stack[:len(stack)-1] {
		peek := stack[len(stack)-1]
		peek.IPAddressLow = PlusOne(pop.IPAddressHigh)
		list = append(list, peek)
	}
	return list, nil
}
<|endoftext|>


<commit_before>package parser

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"os/user"
	"regexp"
	"strconv"
	"strings"

	"github.com/fatih/color"
)

// Initial json Regexp
var jsonReg *regexp.Regexp

// Current accepted text types
// This information is useful to let the parser know
// how to interpret the given file input
var textTypes = []string{
	"plain",
	"json",
}

var (
	// Color functions to help display meaningful input
	success = color.New(color.FgHiGreen, color.Bold).SprintFunc()
	fail    = color.New(color.FgHiWhite, color.BgRed, color.Bold).SprintFunc()
	alert   = color.New(color.FgHiYellow, color.Bold).SprintFunc()
	info    = color.New(color.FgHiMagenta, color.Bold).SprintFunc()
)

// Parser type definition
type parser struct {
	file    string
	text    string
	filter  string
	lines   int
	page    int
	noColor bool
	regex   *regexp.Regexp
}

// Initialize json regexp on package initialization
func init() {
	jsonReg = regexp.MustCompile(`\[?{".*:.*}\]?`)
}

// New returns a new parser object
func New(file, text, filter string, lines, page int, noColor, withRegex bool) *parser {
	// If the file path starts with "~"
	// it means the user is on a Unix based OS
	// and "~" represents the home directory path
	if strings.HasPrefix(file, "~") {
		user, err := user.Current()
		if err != nil {
			log.Fatal(err)
		}
		file = strings.Replace(file, "~", user.HomeDir, 1)
	}
	// Enable regex support is user asks for it
	var regex *regexp.Regexp
	// Attach regexp to parser
	if withRegex {
		// If no filter is provided then regex is useless
		// In this case notify the user
		if filter == "" {
			exitWithError(fail("Error! Regex is present but no filter value was provided!"))
		}
		// Compile regex expression here to be user later in the parser
		re, err := regexp.Compile(filter)
		if err != nil {
			exitWithError(fail(fmt.Sprintf("Regex parse error: %s", err.Error())))
		}
		regex = re
	}

	return &parser{
		file:    file,
		text:    text,
		filter:  filter,
		lines:   lines,
		page:    page,
		noColor: noColor,
		regex:   regex,
	}
}

// Parse parses the file and shows the output to the user
func (p *parser) Parse() {
	// Check if a valid lines value was provided
	if p.lines <= 0 {
		exitWithError(fail("Error! Option flag -lines must be strictly positive"))
	}
	// Check if a valid page value was provided
	if p.page <= 0 {
		exitWithError(fail("Error! Option flag -page must be strictly positive"))
	}
	// Check if a valid test type was provided
	if !stringInSlice(p.text, textTypes) {
		exitWithError(fail(fmt.Sprintf("Error! Accepted text types are: %s", strings.Join(textTypes, ", "))))
	}
	// Disables colorized output
	if p.noColor {
		color.NoColor = true
	}
	// Disables regex support for characters of length 1
	// If regex remains enabled strange out is given
	// Also there is no sense in having a regex with length of 1
	if len(p.filter) == 1 {
		p.regex = nil
	}
	// Check if a file path was provided
	if p.file == "" {
		exitWithError(fail("Error! File is required"))
	}
	// Set current page to what the parser gave us
	currentPage := p.page
	// Count all the lines and provide all page offsets
	// These offsets help us navigate to any page instantly
	// This is where all the "magic" happens
	pageOffsets, err := p.countLines()

	if err != nil {
		log.Fatal(err)
	}
	// Determine total number of pages
	numPages := len(pageOffsets)
	// If nothing can be shown tell the user
	if numPages == 0 {
		fmt.Printf("%s\n", info("Sorry. Nothing to show here!"))
		return
	}
	// The current page cannot be greater than the total number of pages
	if currentPage > numPages {
		fmt.Printf("\n%s\n\n", fail(fmt.Sprintf("Error! Page number cannot be greater than %d", numPages)))
		return
	}
	// Get the first page output and print it
	output := p.getFilePage(pageOffsets[currentPage-1])

	fmt.Println(output)
	// If no more pages are to be shown stop here
	// If means we only have 1 page which we already displayed
	if numPages == 1 {
		return
	}
	// Show a message telling the user at which page we are right now and prompt to navigate to whatever page
	fmt.Print(alert(fmt.Sprintf("Page %d/%d. Enter page number to navigate or press Ctrl+C to quit:", currentPage, numPages)), " ")
	// Start a new input scanner
	in := bufio.NewScanner(os.Stdin)
	// Scan for incoming input
	for in.Scan() {
		// We only accept page numbers here
		// This is how the parser knows where to navigate next
		inputPage, err := strconv.Atoi(in.Text())
		if err != nil {
			fmt.Printf("\n%s\n\n", fail("Error! A valid number is required"))
			fmt.Print(alert(fmt.Sprintf("Page %d/%d. Enter page number to navigate or press Ctrl+C to quit:", currentPage, numPages)), " ")
		} else {
			switch {
			case inputPage > numPages:
				fmt.Printf("\n%s\n\n", fail(fmt.Sprintf("Error! Page number cannot be greater than %d", numPages)))
				fmt.Print(alert(fmt.Sprintf("Page %d/%d. Enter page number to navigate or press Ctrl+C to quit:", currentPage, numPages)), " ")
			case inputPage < 1:
				fmt.Printf("\n%s\n\n", fail("Error! Page number cannot be smaller than 1"))
				fmt.Print(alert(fmt.Sprintf("Page %d/%d. Enter page number to navigate or press Ctrl+C to quit:", currentPage, numPages)), " ")
			default:
				currentPage = inputPage
				output := p.getFilePage(pageOffsets[currentPage-1])
				fmt.Printf("\n%s\n", output)
				fmt.Print(alert(fmt.Sprintf("Page %d/%d. Enter page number to navigate or press Ctrl+C to quit:", currentPage, numPages)), " ")
			}
		}
	}

	if err := in.Err(); err != nil {
		log.Fatal(err)
	}
}

// Gets the output for a new page on the input file
func (p *parser) getFilePage(offset int64) string {
	// Open the file
	f, err := os.Open(p.file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	// Navigate to the given offset
	// This way we skip the part we don't need
	// and avoid parsing unnecessary lines
	_, err = f.Seek(offset, io.SeekStart)
	if err != nil {
		log.Fatal(err)
	}
	// Start a new scanner
	s := bufio.NewScanner(f)
	// Set a larger token size just in case
	s.Buffer(nil, 64*1024*1024)
	// Current line. We need to keep track
	// of all current line number to determine
	// when a page ends
	var current int
	// This will hold the final output to be shown to the user
	// It is reponsable to display only 1 page
	var output bytes.Buffer
	// The text that is obtained by the "getOutput" method
	var text string
	// Scan the file and extract all lines
	// Stop when we reach the number of lines per page
	// that the user specified
	for s.Scan() {
		current++
		// Get the line output and added in the buffer
		text = p.getOutput(s.Text())
		fmt.Fprintf(&output, "%s\n", text)
		if current >= p.lines {
			break
		}
	}

	if err := s.Err(); err != nil {
		log.Fatal(err)
	}

	return output.String()
}

// This is where all the "magic" happens
// Here we count all the file lines
// and extract a slice with all page offsets
func (p *parser) countLines() ([]int64, error) {
	// Open the file
	f, err := os.Open(p.file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	// Get file size the know the position
	// in the file of the last byte
	fi, err := f.Stat()
	if err != nil {
		log.Fatal(err)
	}
	lastPosition := fi.Size()
	// Start a new reader
	r := bufio.NewReader(f)
	// Here we store all page offsets for all pages
	var pageOffsets []int64
	// Here we store page offsets corresponding to filtered text
	var filterOffsets []int64
	// The first line
	// We need to keep track of every line
	// to know when a page ends
	var currentLine int
	// We store the last offset to help us
	// calculate the next offset in the slice
	var lastOffset int64
	// This is the current offset and is
	// calculated line by line
	var offset int64
	// This is used to know if page has been hit
	// by a filter provided by the user
	var pageHit bool
	// We start by adding the first page offset which is 0
	pageOffsets = append(pageOffsets, offset)

	for {
		currentLine++
		// Read the input file line by line
		// This may take a while depending on the file size
		// This is the most time consuming portion of the app
		line, err := r.ReadBytes('\n')
		// If we have a line hit it means we have a page hit
		if p.lineHit(line) {
			pageHit = true
		}
		// Compute the current offset
		offset += int64(len(line))
		if currentLine == p.lines {
			// If we have reached the end of the page
			// Extract the last page offset
			lastOffset = pageOffsets[len(pageOffsets)-1]
			// If we have a page hit (from filtered input)
			// Add it to the filter offsets
			if pageHit {
				filterOffsets = append(filterOffsets, pageOffsets[len(pageOffsets)-1])
			}
			// Add it to the current offset
			// Append this offset to the end of our slice
			offset += lastOffset
			if offset < lastPosition {
				pageOffsets = append(pageOffsets, offset)
			}
			// Reset all counters
			// and also page hit
			currentLine = 0
			offset = 0
			pageHit = false
		}

		switch {
		case err == io.EOF:
			// If we reached the end of file
			// we must take into account also the last lines
			// which were not caught by the page offsets
			if currentLine < p.lines && pageHit {
				filterOffsets = append(filterOffsets, pageOffsets[len(pageOffsets)-1])
			}
			// If the input was filtered return filter page offsets
			// Otherwise return normal page offsets
			if p.filter != "" {
				return filterOffsets, nil
			}
			return pageOffsets, nil

		case err != nil:
			return nil, err
		}
	}
}

// This determines if line matches a given filter
func (p *parser) lineHit(line []byte) bool {
	// If no filter was provided then we do not care about this
	if p.filter == "" {
		return false
	}
	// If regex was enbled, search by regex
	// Otherwise search normally by text bytes
	if p.regex != nil {
		found := p.regex.Find(line)
		return found != nil
	}

	return bytes.Contains(line, []byte(p.filter))
}

// This computes the final output
func (p *parser) getOutput(text string) string {
	// Format input as JSON if needed
	if p.text == "json" {
		jsonMatches := jsonReg.FindAll([]byte(text), -1)
		if jsonMatches != nil {
			for _, m := range jsonMatches {
				text = strings.Replace(text, string(m), formatJSON(m), -1)
			}
		}
	}
	// If no filter was provided give the text as it is
	if p.filter == "" {
		return text
	}
	// For regex/normal search highlight filtered input text
	if p.regex != nil {
		// Find all matches
		matches := p.regex.FindAllString(text, -1)
		if matches == nil {
			return text
		}
		// Highlight every match in given input
		for _, m := range matches {
			text = strings.Replace(text, m, success(m), -1)
		}
		return text
	}

	return strings.Replace(text, p.filter, success(p.filter), -1)
}
<commit_msg>Added some useful comments<commit_after>package parser

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"os/user"
	"regexp"
	"strconv"
	"strings"

	"github.com/fatih/color"
)

// Initial json Regexp
var jsonReg *regexp.Regexp

// Current accepted text types
// This information is useful to let the parser know
// how to interpret the given file input
var textTypes = []string{
	"plain",
	"json",
}

var (
	// Color functions to help display meaningful input
	success = color.New(color.FgHiGreen, color.Bold).SprintFunc()
	fail    = color.New(color.FgHiWhite, color.BgRed, color.Bold).SprintFunc()
	alert   = color.New(color.FgHiYellow, color.Bold).SprintFunc()
	info    = color.New(color.FgHiMagenta, color.Bold).SprintFunc()
)

// Parser type definition
type parser struct {
	file    string
	text    string
	filter  string
	lines   int
	page    int
	noColor bool
	regex   *regexp.Regexp
}

// Initialize json regexp on package initialization
func init() {
	jsonReg = regexp.MustCompile(`\[?{".*:.*}\]?`)
}

// New returns a new parser object
func New(file, text, filter string, lines, page int, noColor, withRegex bool) *parser {
	// If the file path starts with "~"
	// it means the user is probably on a Unix based OS
	// and "~" represents the home directory path
	if strings.HasPrefix(file, "~") {
		user, err := user.Current()
		if err != nil {
			log.Fatal(err)
		}
		file = strings.Replace(file, "~", user.HomeDir, 1)
	}
	// Enable regex support is user asks for it
	var regex *regexp.Regexp
	// Attach regexp to parser
	if withRegex {
		// If no filter is provided then regex is useless
		// In this case notify the user
		if filter == "" {
			exitWithError(fail("Error! Regex is present but no filter value was provided!"))
		}
		// Compile regex expression here to be user later in the parser
		re, err := regexp.Compile(filter)
		if err != nil {
			exitWithError(fail(fmt.Sprintf("Regex parse error: %s", err.Error())))
		}
		regex = re
	}

	return &parser{
		file:    file,
		text:    text,
		filter:  filter,
		lines:   lines,
		page:    page,
		noColor: noColor,
		regex:   regex,
	}
}

// Parse parses the file and shows the output to the user
func (p *parser) Parse() {
	// Check if a valid lines value was provided
	if p.lines <= 0 {
		exitWithError(fail("Error! Option flag -lines must be strictly positive"))
	}
	// Check if a valid page value was provided
	if p.page <= 0 {
		exitWithError(fail("Error! Option flag -page must be strictly positive"))
	}
	// Check if a valid test type was provided
	if !stringInSlice(p.text, textTypes) {
		exitWithError(fail(fmt.Sprintf("Error! Accepted text types are: %s", strings.Join(textTypes, ", "))))
	}
	// Disables colorized output
	if p.noColor {
		color.NoColor = true
	}
	// Disables regex support for characters of length 1
	// If regex remains enabled strange out is given
	// Also there is no sense in having a regex with length of 1
	if len(p.filter) == 1 {
		p.regex = nil
	}
	// Check if a file path was provided
	if p.file == "" {
		exitWithError(fail("Error! File is required"))
	}
	// Set current page to what the parser gave us
	currentPage := p.page
	// Count all the lines and provide all page offsets
	// These offsets help us navigate to any page instantly
	// This is where all the "magic" happens
	pageOffsets, err := p.countLines()

	if err != nil {
		log.Fatal(err)
	}
	// Determine total number of pages
	numPages := len(pageOffsets)
	// If nothing can be shown tell the user
	if numPages == 0 {
		fmt.Printf("%s\n", info("Sorry. Nothing to show here!"))
		return
	}
	// The current page cannot be greater than the total number of pages
	if currentPage > numPages {
		fmt.Printf("\n%s\n\n", fail(fmt.Sprintf("Error! Page number cannot be greater than %d", numPages)))
		return
	}
	// Get the first page output and print it
	output := p.getFilePage(pageOffsets[currentPage-1])

	fmt.Println(output)
	// If no more pages are to be shown stop here
	// If means we only have 1 page which we already displayed
	if numPages == 1 {
		return
	}
	// Show a message telling the user at which page we are right now and prompt to navigate to whatever page
	fmt.Print(alert(fmt.Sprintf("Page %d/%d. Enter page number to navigate or press Ctrl+C to quit:", currentPage, numPages)), " ")
	// Start a new input scanner
	in := bufio.NewScanner(os.Stdin)
	// Scan for incoming input
	for in.Scan() {
		// We only accept page numbers here
		// This is how the parser knows where to navigate next
		inputPage, err := strconv.Atoi(in.Text())
		if err != nil {
			fmt.Printf("\n%s\n\n", fail("Error! A valid number is required"))
			fmt.Print(alert(fmt.Sprintf("Page %d/%d. Enter page number to navigate or press Ctrl+C to quit:", currentPage, numPages)), " ")
		} else {
			switch {
			case inputPage > numPages:
				fmt.Printf("\n%s\n\n", fail(fmt.Sprintf("Error! Page number cannot be greater than %d", numPages)))
				fmt.Print(alert(fmt.Sprintf("Page %d/%d. Enter page number to navigate or press Ctrl+C to quit:", currentPage, numPages)), " ")
			case inputPage < 1:
				fmt.Printf("\n%s\n\n", fail("Error! Page number cannot be smaller than 1"))
				fmt.Print(alert(fmt.Sprintf("Page %d/%d. Enter page number to navigate or press Ctrl+C to quit:", currentPage, numPages)), " ")
			default:
				currentPage = inputPage
				output := p.getFilePage(pageOffsets[currentPage-1])
				fmt.Printf("\n%s\n", output)
				fmt.Print(alert(fmt.Sprintf("Page %d/%d. Enter page number to navigate or press Ctrl+C to quit:", currentPage, numPages)), " ")
			}
		}
	}

	if err := in.Err(); err != nil {
		log.Fatal(err)
	}
}

// Gets the output for a new page on the input file
func (p *parser) getFilePage(offset int64) string {
	// Open the file
	f, err := os.Open(p.file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	// Navigate to the given offset
	// This way we skip the part we don't need
	// and avoid parsing unnecessary lines
	_, err = f.Seek(offset, io.SeekStart)
	if err != nil {
		log.Fatal(err)
	}
	// Start a new scanner
	s := bufio.NewScanner(f)
	// Set a larger token size just in case
	s.Buffer(nil, 64*1024*1024)
	// Current line. We need to keep track
	// of all current line number to determine
	// when a page ends
	var current int
	// This will hold the final output to be shown to the user
	// It is reponsable to display only 1 page
	var output bytes.Buffer
	// The text that is obtained by the "getOutput" method
	var text string
	// Scan the file and extract all lines
	// Stop when we reach the number of lines per page
	// that the user specified
	for s.Scan() {
		current++
		// Get the line output and added in the buffer
		text = p.getOutput(s.Text())
		fmt.Fprintf(&output, "%s\n", text)
		if current >= p.lines {
			break
		}
	}

	if err := s.Err(); err != nil {
		log.Fatal(err)
	}

	return output.String()
}

// This is where all the "magic" happens
// Here we count all the file lines
// and extract a slice with all page offsets
func (p *parser) countLines() ([]int64, error) {
	// Open the file
	f, err := os.Open(p.file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	// Get file size the know the position
	// in the file of the last byte
	fi, err := f.Stat()
	if err != nil {
		log.Fatal(err)
	}
	lastPosition := fi.Size()
	// Start a new reader
	r := bufio.NewReader(f)
	// Here we store all page offsets for all pages
	var pageOffsets []int64
	// Here we store page offsets corresponding to filtered text
	var filterOffsets []int64
	// The first line
	// We need to keep track of every line
	// to know when a page ends
	var currentLine int
	// We store the last offset to help us
	// calculate the next offset in the slice
	var lastOffset int64
	// This is the current offset and is
	// calculated line by line
	var offset int64
	// This is used to know if page has been hit
	// by a filter provided by the user
	var pageHit bool
	// We start by adding the first page offset which is 0
	pageOffsets = append(pageOffsets, offset)

	for {
		currentLine++
		// Read the input file line by line
		// This may take a while depending on the file size
		// This is the most time consuming portion of the app
		line, err := r.ReadBytes('\n')
		// If we have a line hit it means we have a page hit
		if p.lineHit(line) {
			pageHit = true
		}
		// Compute the current offset
		offset += int64(len(line))
		if currentLine == p.lines {
			// If we have reached the end of the page
			// Extract the last page offset
			lastOffset = pageOffsets[len(pageOffsets)-1]
			// If we have a page hit (from filtered input)
			// Add it to the filter offsets
			if pageHit {
				filterOffsets = append(filterOffsets, pageOffsets[len(pageOffsets)-1])
			}
			// Add it to the current offset
			// Append this offset to the end of our slice
			offset += lastOffset
			if offset < lastPosition {
				pageOffsets = append(pageOffsets, offset)
			}
			// Reset all counters
			// and also page hit
			currentLine = 0
			offset = 0
			pageHit = false
		}

		switch {
		case err == io.EOF:
			// If we reached the end of file
			// we must take into account also the last lines
			// which were not caught by the page offsets
			if currentLine < p.lines && pageHit {
				filterOffsets = append(filterOffsets, pageOffsets[len(pageOffsets)-1])
			}
			// If the input was filtered return filter page offsets
			// Otherwise return normal page offsets
			if p.filter != "" {
				return filterOffsets, nil
			}
			return pageOffsets, nil

		case err != nil:
			return nil, err
		}
	}
}

// This determines if line matches a given filter
func (p *parser) lineHit(line []byte) bool {
	// If no filter was provided then we do not care about this
	if p.filter == "" {
		return false
	}
	// If regex was enbled, search by regex
	// Otherwise search normally by text bytes
	if p.regex != nil {
		found := p.regex.Find(line)
		return found != nil
	}

	return bytes.Contains(line, []byte(p.filter))
}

// This computes the final output
func (p *parser) getOutput(text string) string {
	// Format input as JSON if needed
	if p.text == "json" {
		jsonMatches := jsonReg.FindAll([]byte(text), -1)
		if jsonMatches != nil {
			for _, m := range jsonMatches {
				text = strings.Replace(text, string(m), formatJSON(m), -1)
			}
		}
	}
	// If no filter was provided give the text as it is
	if p.filter == "" {
		return text
	}
	// For regex/normal search highlight filtered input text
	if p.regex != nil {
		// Find all matches
		matches := p.regex.FindAllString(text, -1)
		if matches == nil {
			return text
		}
		// Highlight every match in given input
		for _, m := range matches {
			text = strings.Replace(text, m, success(m), -1)
		}
		return text
	}

	return strings.Replace(text, p.filter, success(p.filter), -1)
}
<|endoftext|>


<commit_before>package parser

import (
	"errors"
	"fmt"
	"github.com/twtiger/go-seccomp/tree"
	"go/ast"
	"go/parser"
	"go/token"
	"regexp"
	"strconv"
)

var tokenTypes = make(map[token.Token]string)
var arithmeticOps = make(map[token.Token]tree.ArithmeticType)
var comparisonOps = make(map[token.Token]tree.ComparisonType)

func buildArithmeticOps() {
	arithmeticOps[token.ADD] = tree.PLUS
	arithmeticOps[token.SUB] = tree.MINUS
	arithmeticOps[token.MUL] = tree.MULT
	arithmeticOps[token.QUO] = tree.DIV
	arithmeticOps[token.REM] = tree.MOD
	arithmeticOps[token.AND] = tree.BINAND
	arithmeticOps[token.OR] = tree.BINOR
	arithmeticOps[token.XOR] = tree.BINXOR
	arithmeticOps[token.SHL] = tree.LSH
	arithmeticOps[token.SHR] = tree.RSH
}

func buildComparisonOps() {
	comparisonOps[token.EQL] = tree.EQL
	comparisonOps[token.LSS] = tree.LT
	comparisonOps[token.GTR] = tree.GT
	comparisonOps[token.NEQ] = tree.NEQL
	comparisonOps[token.LEQ] = tree.LTE
	comparisonOps[token.GEQ] = tree.GTE
	comparisonOps[token.AND] = tree.BIT
}

func init() {
	buildArithmeticOps()
	buildComparisonOps()

	tokenTypes[token.LOR] = "booleanArguments"
	tokenTypes[token.LAND] = "booleanArguments"

	for k := range arithmeticOps {
		tokenTypes[k] = "integerArguments"
	}

	for k := range comparisonOps {
		tokenTypes[k] = "numericArguments"
	}
}

func surround(s string) string {
	return "func() { " + s + "}"
}

func parseExpression(expr string) (tree.Expression, error) {
	// fs := token.NewFileSet()
	tr, _ := parser.ParseExpr(surround(expr))
	// ast.Print(fs, tr)
	parsedtree, err := unwrapToplevel(tr)
	return parsedtree, err
}

func unwrapToplevel(x ast.Node) (tree.Expression, error) {
	switch f := x.(type) {
	case *ast.FuncLit:
		return unwrapToplevel(f.Body)
	case *ast.BlockStmt:
		return unwrapToplevel(f.List[0])
	case *ast.ExprStmt:
		return unwrapBooleanExpression(f.X)
	default:
		// panicWithInfo(x)
	}
	return nil, errors.New("Expression is invalid. Unable to parse.")
}

var argRegexpRE = regexp.MustCompile(`^arg([0-5])$`)

func identExpression(f *ast.Ident) (tree.Numeric, error) {
	if match := argRegexpRE.FindStringSubmatch(f.Name); match != nil {
		ix, _ := strconv.Atoi(match[1])
		return tree.Argument{ix}, nil
	}
	switch f.Name {
	case "true":
		return tree.BooleanLiteral{true}, nil
	case "false":
		return tree.BooleanLiteral{false}, nil
	// Handle other cases here
	default:
		return tree.Variable{f.Name}, nil
	}
	return tree.Variable{f.Name}, nil
}

func unwrapNumericExpression(x ast.Node) (tree.Numeric, error) {
	switch f := x.(type) {
	case *ast.Ident:
		// Ensure ident doesn't contain stupidness like packages and stuff
		return identExpression(f)
	case *ast.BasicLit:
		// TODO: errors here
		i, _ := strconv.Atoi(f.Value)
		return tree.NumericLiteral{uint32(i)}, nil
	case *ast.BinaryExpr:
		left, err := unwrapNumericExpression(f.X)
		right, err := unwrapNumericExpression(f.Y)
		op := arithmeticOps[f.Op]
		// TODO: handle operators we don't support here
		if err != nil {
			return nil, err
		}
		return tree.Arithmetic{Left: left, Right: right, Op: op}, nil
	case *ast.ParenExpr:
		return unwrapNumericExpression(f.X)
	case *ast.UnaryExpr:
		operand, err := unwrapNumericExpression(f.X)
		if err != nil {
			return nil, err
		}
		if f.Op == token.XOR {
			return tree.BinaryNegation{operand}, nil
		}
		// TODO: Fail in a good way here
	default:
		// panicWithInfo(x)
	}
	return nil, errors.New("Expression is invalid. Unable to parse.")
}

func panicWithInfo(x interface{}) {
	panic(fmt.Sprintf("sadness: %#v", x))
}

func takesBooleanArguments(f *ast.BinaryExpr) bool {
	return tokenTypes[f.Op] == "booleanArguments"
}

func takesNumericArguments(f *ast.BinaryExpr) bool {
	return tokenTypes[f.Op] == "numericArguments"
}

func booleanArgExpression(f *ast.BinaryExpr) (tree.Boolean, error) {
	left, err := unwrapBooleanExpression(f.X)
	right, err := unwrapBooleanExpression(f.Y)
	switch f.Op {
	case token.LOR:
		return tree.Or{Left: left, Right: right}, nil
	case token.LAND:
		return tree.And{Left: left, Right: right}, nil
	}
	return nil, err
}

func numericArgExpression(f *ast.BinaryExpr) (tree.Boolean, error) {
	cmp := comparisonOps[f.Op]
	// TODO: handle incorrect thingy here
	left, err := unwrapNumericExpression(f.X)
	right, err := unwrapNumericExpression(f.Y)
	if err != nil {
		return nil, err
	}
	return tree.Comparison{Left: left, Right: right, Op: cmp}, nil
}

func inclusionExpression(f *ast.CallExpr) (tree.Boolean, error) {
	var pos bool
	var left tree.Numeric
	var err error
	var right []tree.Numeric
	var val tree.Numeric

	switch p := f.Fun.(type) {
	case *ast.Ident:
		if p.Name == "in" {
			pos = true
		}
		if p.Name == "notIn" {
			pos = false
		}
	}

	switch p := f.Args[0].(type) {
	case *ast.Ident:
		left, err = identExpression(p)
	case *ast.BasicLit:
		left, err = unwrapNumericExpression(p)
	}

	for _, e := range f.Args {
		switch y := e.(type) {
		case *ast.BasicLit:
			val, err = unwrapNumericExpression(y)
			right = append(right, val)
		}
	}

	if err != nil {
		return nil, err
	}

	return tree.Inclusion{pos, left, right}, nil
}

func unwrapBooleanExpression(x ast.Node) (tree.Boolean, error) {
	switch f := x.(type) {
	case *ast.BasicLit:
		switch f.Value {
		// TODO: Handle other values here
		case "1":
			return tree.BooleanLiteral{true}, nil
		case "0":
			return tree.BooleanLiteral{false}, nil
		}
		// TODO: handle failure here
	case *ast.BinaryExpr:
		if takesBooleanArguments(f) {
			return booleanArgExpression(f)
		} else if takesNumericArguments(f) {
			return numericArgExpression(f)
		}
	case *ast.ParenExpr:
		return unwrapBooleanExpression(f.X)
	case *ast.UnaryExpr:
		operand, err := unwrapBooleanExpression(f.X)
		if err == nil {
			if f.Op == token.NOT {
				return tree.Negation{operand}, nil
			}
		}
	case *ast.CallExpr:
		return inclusionExpression(f)
	case *ast.Ident:
		return identExpression(f)
		// TODO: Fail in a good way here
	default:
		// panicWithInfo(x)
	}
	return nil, errors.New("Expression is invalid. Unable to parse.")
}
<commit_msg>Uh okay goimports<commit_after>package parser

import (
	"errors"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"regexp"
	"strconv"

	"github.com/twtiger/go-seccomp/tree"
)

var tokenTypes = make(map[token.Token]string)
var arithmeticOps = make(map[token.Token]tree.ArithmeticType)
var comparisonOps = make(map[token.Token]tree.ComparisonType)

func buildArithmeticOps() {
	arithmeticOps[token.ADD] = tree.PLUS
	arithmeticOps[token.SUB] = tree.MINUS
	arithmeticOps[token.MUL] = tree.MULT
	arithmeticOps[token.QUO] = tree.DIV
	arithmeticOps[token.REM] = tree.MOD
	arithmeticOps[token.AND] = tree.BINAND
	arithmeticOps[token.OR] = tree.BINOR
	arithmeticOps[token.XOR] = tree.BINXOR
	arithmeticOps[token.SHL] = tree.LSH
	arithmeticOps[token.SHR] = tree.RSH
}

func buildComparisonOps() {
	comparisonOps[token.EQL] = tree.EQL
	comparisonOps[token.LSS] = tree.LT
	comparisonOps[token.GTR] = tree.GT
	comparisonOps[token.NEQ] = tree.NEQL
	comparisonOps[token.LEQ] = tree.LTE
	comparisonOps[token.GEQ] = tree.GTE
	comparisonOps[token.AND] = tree.BIT
}

func init() {
	buildArithmeticOps()
	buildComparisonOps()

	tokenTypes[token.LOR] = "booleanArguments"
	tokenTypes[token.LAND] = "booleanArguments"

	for k := range arithmeticOps {
		tokenTypes[k] = "integerArguments"
	}

	for k := range comparisonOps {
		tokenTypes[k] = "numericArguments"
	}
}

func surround(s string) string {
	return "func() { " + s + "}"
}

func parseExpression(expr string) (tree.Expression, error) {
	// fs := token.NewFileSet()
	tr, _ := parser.ParseExpr(surround(expr))
	// ast.Print(fs, tr)
	parsedtree, err := unwrapToplevel(tr)
	return parsedtree, err
}

func unwrapToplevel(x ast.Node) (tree.Expression, error) {
	switch f := x.(type) {
	case *ast.FuncLit:
		return unwrapToplevel(f.Body)
	case *ast.BlockStmt:
		return unwrapToplevel(f.List[0])
	case *ast.ExprStmt:
		return unwrapBooleanExpression(f.X)
	default:
		// panicWithInfo(x)
	}
	return nil, errors.New("Expression is invalid. Unable to parse.")
}

var argRegexpRE = regexp.MustCompile(`^arg([0-5])$`)

func identExpression(f *ast.Ident) (tree.Numeric, error) {
	if match := argRegexpRE.FindStringSubmatch(f.Name); match != nil {
		ix, _ := strconv.Atoi(match[1])
		return tree.Argument{ix}, nil
	}
	switch f.Name {
	case "true":
		return tree.BooleanLiteral{true}, nil
	case "false":
		return tree.BooleanLiteral{false}, nil
	// Handle other cases here
	default:
		return tree.Variable{f.Name}, nil
	}
	return tree.Variable{f.Name}, nil
}

func unwrapNumericExpression(x ast.Node) (tree.Numeric, error) {
	switch f := x.(type) {
	case *ast.Ident:
		// Ensure ident doesn't contain stupidness like packages and stuff
		return identExpression(f)
	case *ast.BasicLit:
		// TODO: errors here
		i, _ := strconv.Atoi(f.Value)
		return tree.NumericLiteral{uint32(i)}, nil
	case *ast.BinaryExpr:
		left, err := unwrapNumericExpression(f.X)
		right, err := unwrapNumericExpression(f.Y)
		op := arithmeticOps[f.Op]
		// TODO: handle operators we don't support here
		if err != nil {
			return nil, err
		}
		return tree.Arithmetic{Left: left, Right: right, Op: op}, nil
	case *ast.ParenExpr:
		return unwrapNumericExpression(f.X)
	case *ast.UnaryExpr:
		operand, err := unwrapNumericExpression(f.X)
		if err != nil {
			return nil, err
		}
		if f.Op == token.XOR {
			return tree.BinaryNegation{operand}, nil
		}
		// TODO: Fail in a good way here
	default:
		// panicWithInfo(x)
	}
	return nil, errors.New("Expression is invalid. Unable to parse.")
}

func panicWithInfo(x interface{}) {
	panic(fmt.Sprintf("sadness: %#v", x))
}

func takesBooleanArguments(f *ast.BinaryExpr) bool {
	return tokenTypes[f.Op] == "booleanArguments"
}

func takesNumericArguments(f *ast.BinaryExpr) bool {
	return tokenTypes[f.Op] == "numericArguments"
}

func booleanArgExpression(f *ast.BinaryExpr) (tree.Boolean, error) {
	left, err := unwrapBooleanExpression(f.X)
	right, err := unwrapBooleanExpression(f.Y)
	switch f.Op {
	case token.LOR:
		return tree.Or{Left: left, Right: right}, nil
	case token.LAND:
		return tree.And{Left: left, Right: right}, nil
	}
	return nil, err
}

func numericArgExpression(f *ast.BinaryExpr) (tree.Boolean, error) {
	cmp := comparisonOps[f.Op]
	// TODO: handle incorrect thingy here
	left, err := unwrapNumericExpression(f.X)
	right, err := unwrapNumericExpression(f.Y)
	if err != nil {
		return nil, err
	}
	return tree.Comparison{Left: left, Right: right, Op: cmp}, nil
}

func inclusionExpression(f *ast.CallExpr) (tree.Boolean, error) {
	var pos bool
	var left tree.Numeric
	var err error
	var right []tree.Numeric
	var val tree.Numeric

	switch p := f.Fun.(type) {
	case *ast.Ident:
		if p.Name == "in" {
			pos = true
		}
		if p.Name == "notIn" {
			pos = false
		}
	}

	switch p := f.Args[0].(type) {
	case *ast.Ident:
		left, err = identExpression(p)
	case *ast.BasicLit:
		left, err = unwrapNumericExpression(p)
	}

	for _, e := range f.Args {
		switch y := e.(type) {
		case *ast.BasicLit:
			val, err = unwrapNumericExpression(y)
			right = append(right, val)
		}
	}

	if err != nil {
		return nil, err
	}

	return tree.Inclusion{pos, left, right}, nil
}

func unwrapBooleanExpression(x ast.Node) (tree.Boolean, error) {
	switch f := x.(type) {
	case *ast.BasicLit:
		switch f.Value {
		// TODO: Handle other values here
		case "1":
			return tree.BooleanLiteral{true}, nil
		case "0":
			return tree.BooleanLiteral{false}, nil
		}
		// TODO: handle failure here
	case *ast.BinaryExpr:
		if takesBooleanArguments(f) {
			return booleanArgExpression(f)
		} else if takesNumericArguments(f) {
			return numericArgExpression(f)
		}
	case *ast.ParenExpr:
		return unwrapBooleanExpression(f.X)
	case *ast.UnaryExpr:
		operand, err := unwrapBooleanExpression(f.X)
		if err == nil {
			if f.Op == token.NOT {
				return tree.Negation{operand}, nil
			}
		}
	case *ast.CallExpr:
		return inclusionExpression(f)
	case *ast.Ident:
		return identExpression(f)
		// TODO: Fail in a good way here
	default:
		// panicWithInfo(x)
	}
	return nil, errors.New("Expression is invalid. Unable to parse.")
}
<|endoftext|>