<commit_before>package gtreap

import (
	"errors"
	"os"
)

// A persistent store holding collections of ordered keys & values.
// The persistence is append-only based on immutable, copy-on-write
// treaps.  This implementation is single-threaded, so users should
// serialize their accesses.
//
// TODO: use atomic.CAS and unsafe.Pointers for safe snapshot'ability.
// TODO: allow read-only snapshots.
//
type Store struct {
	coll map[string]*PTreap
	file *os.File
}

func NewStore(file *os.File) (*Store, error) {
	if file == nil { // Return a memory-only Store.
		return &Store{coll: make(map[string]*PTreap)}, nil
	}
	return nil, errors.New("not implemented yet")
}

func (s *Store) AddCollection(name string, compare KeyCompare) *PTreap {
	if s.coll[name] == nil {
		s.coll[name] = &PTreap{store: s, compare: compare}
	}
	return s.coll[name]
}

func (s *Store) GetCollection(name string) *PTreap {
	return s.coll[name]
}

func (s *Store) RemoveCollection(name string) {
	delete(s.coll, name)
}

// User-supplied key comparison func should return 0 if a == b,
// -1 if a < b, and +1 if a > b.
type KeyCompare func(a, b []byte) int

// A persisted treap.
type PTreap struct {
	store   *Store
	compare KeyCompare
	root    pnodeLoc
}

// A persisted node and its persistence location.
type pnodeLoc struct {
	loc  *ploc  // Can be nil if node is dirty (not yet persisted).
	node *pnode // Can be nil if node is not fetched into memory yet.
}

func (nloc *pnodeLoc) isEmpty() bool {
	return nloc == nil || (nloc.loc == nil && nloc.node == nil)
}

var empty = &pnodeLoc{}

// A persisted node.
type pnode struct {
	item        PItemLoc
	left, right pnodeLoc
}

// A persisted item.
type PItem struct {
	Key, Val []byte // Val may be nil if not fetched into memory yet.
	Priority int32
}

// A persisted item and its persistence location.
type PItemLoc struct {
	loc  *ploc  // Can be nil if item is dirty (not yet persisted).
	item *PItem // Can be nil if item is not fetched into memory yet.
}

// Offset/location of persisted range of bytes.
type ploc struct {
	offset int64  // Usable for os.Seek/ReadAt/WriteAt() at file offset 0.
	length uint32 // Number of bytes.
}

func (t *PTreap) Get(key []byte, withValue bool) (*PItem, error) {
	n, err := t.store.loadNodeLoc(&t.root)
	for {
		if err != nil || n.isEmpty() {
			break
		}
		i, err := t.store.loadItemLoc(&n.node.item, false)
		if err != nil {
			return nil, err
		}
		if i == nil || i.item == nil || i.item.Key == nil {
			panic("no item after loadMetaItemLoc() in get()")
		}
		c := t.compare(key, i.item.Key)
		if c < 0 {
			n, err = t.store.loadNodeLoc(&n.node.left)
		} else if c > 0 {
			n, err = t.store.loadNodeLoc(&n.node.right)
		} else {
			if withValue {
				t.store.loadItemLoc(i, true)
			}
			return i.item, nil
		}
	}
	return nil, err
}

// Replace or insert an item of a given key.
func (t *PTreap) Upsert(item *PItem) error {
	r, err := t.store.union(t, &t.root,
		&pnodeLoc{node: &pnode{item: PItemLoc{item: &PItem{
			Key:      item.Key,
			Val:      item.Val,
			Priority: item.Priority,
		}}}})
	if err == nil {
		t.root = *r
	}
	return err
}

func (t *PTreap) Delete(key []byte) error {
	left, _, right, err := t.store.split(t, &t.root, key)
	if err == nil {
		r, err := t.store.join(left, right)
		if err == nil {
			t.root = *r
		}
	}
	return err
}

func (t *PTreap) Min(withValue bool) (*PItem, error) {
	return t.store.edge(t, withValue, func(n *pnode) *pnodeLoc { return &n.left })
}

func (t *PTreap) Max(withValue bool) (*PItem, error) {
	return t.store.edge(t, withValue, func(n *pnode) *pnodeLoc { return &n.right })
}

type PItemVisitor func(i *PItem) bool

// Visit items greater-than-or-equal to the target.
func (t *PTreap) VisitAscend(target []byte, withValue bool, visitor PItemVisitor) error {
	_, err := t.store.visitAscendNode(t, &t.root, target, withValue, visitor)
	return err
}

func (o *Store) loadNodeLoc(nloc *pnodeLoc) (*pnodeLoc, error) {
	if nloc != nil && nloc.node == nil && nloc.loc != nil {
		// TODO.
	}
	return nloc, nil
}

func (o *Store) loadItemLoc(iloc *PItemLoc, withValue bool) (*PItemLoc, error) {
	if iloc != nil && iloc.item == nil && iloc.loc != nil {
		// TODO.
	}
	return iloc, nil
}

func (o *Store) union(t *PTreap, this *pnodeLoc, that *pnodeLoc) (*pnodeLoc, error) {
	thisNode, err := o.loadNodeLoc(this)
	if err != nil {
		return empty, err
	}
	thatNode, err := o.loadNodeLoc(that)
	if err != nil {
		return empty, err
	}

	if thisNode.isEmpty() {
		return that, nil
	}
	if thatNode.isEmpty() {
		return this, nil
	}

	thisItem, err := o.loadItemLoc(&thisNode.node.item, false)
	if err != nil {
		return empty, err
	}
	thatItem, err := o.loadItemLoc(&thatNode.node.item, false)
	if err != nil {
		return empty, err
	}

	if thisItem.item.Priority > thatItem.item.Priority {
		left, middle, right, err := o.split(t, that, thisItem.item.Key)
		if err != nil {
			return empty, err
		}
		if middle.isEmpty() {
			newLeft, err := o.union(t, &thisNode.node.left, left)
			if err != nil {
				return empty, err
			}
			newRight, err := o.union(t, &thisNode.node.right, right)
			if err != nil {
				return nil, err
			}
			return &pnodeLoc{node: &pnode{
				item:  *thisItem,
				left:  *newLeft,
				right: *newRight,
			}}, nil
		}

		newLeft, err := o.union(t, &thisNode.node.left, left)
		if err != nil {
			return empty, err
		}
		newRight, err := o.union(t, &thisNode.node.right, right)
		if err != nil {
			return empty, err
		}
		return &pnodeLoc{node: &pnode{
			item:  middle.node.item,
			left:  *newLeft,
			right: *newRight,
		}}, nil
	}

	// We don't use middle because the "that" node has precendence.
	left, _, right, err := o.split(t, this, thatItem.item.Key)
	if err != nil {
		return empty, err
	}
	newLeft, err := o.union(t, left, &thatNode.node.left)
	if err != nil {
		return empty, err
	}
	newRight, err := o.union(t, right, &thatNode.node.right)
	if err != nil {
		return empty, err
	}
	return &pnodeLoc{node: &pnode{
		item:  *thatItem,
		left:  *newLeft,
		right: *newRight,
	}}, nil
}

// Splits a treap into two treaps based on a split key "s".  The
// result is (left, middle, right), where left treap has keys < s,
// right treap has keys > s, and middle is either...
// * empty/nil - meaning key s was not in the original treap.
// * non-empty - returning the pnodeLoc that had item s.
func (o *Store) split(t *PTreap, n *pnodeLoc, s []byte) (
	*pnodeLoc, *pnodeLoc, *pnodeLoc, error) {
	nNode, err := o.loadNodeLoc(n)
	if err != nil || nNode.isEmpty() {
		return empty, empty, empty, err
	}
	nItem, err := o.loadItemLoc(&nNode.node.item, false)
	if err != nil {
		return empty, empty, empty, err
	}

	c := t.compare(s, nItem.item.Key)
	if c == 0 {
		return &nNode.node.left, n, &nNode.node.right, nil
	}

	if c < 0 {
		left, middle, right, err := o.split(t, &nNode.node.left, s)
		if err != nil {
			return empty, empty, empty, err
		}
		return left, middle, &pnodeLoc{node: &pnode{
			item:  *nItem,
			left:  *right,
			right: nNode.node.right,
		}}, nil
	}

	left, middle, right, err := o.split(t, &nNode.node.right, s)
	if err != nil {
		return empty, empty, empty, err
	}
	return &pnodeLoc{node: &pnode{
		item:  *nItem,
		left:  nNode.node.left,
		right: *left,
	}}, middle, right, nil
}

// All the keys from this are < keys from that.
func (o *Store) join(this *pnodeLoc, that *pnodeLoc) (*pnodeLoc, error) {
	thisNode, err := o.loadNodeLoc(this)
	if err != nil {
		return empty, err
	}
	thatNode, err := o.loadNodeLoc(that)
	if err != nil {
		return empty, err
	}

	if thisNode.isEmpty() {
		return that, nil
	}
	if thatNode.isEmpty() {
		return this, nil
	}

	thisItem, err := o.loadItemLoc(&thisNode.node.item, false)
	if err != nil {
		return empty, err
	}
	thatItem, err := o.loadItemLoc(&thatNode.node.item, false)
	if err != nil {
		return empty, err
	}

	if thisItem.item.Priority > thatItem.item.Priority {
		newRight, err := o.join(&thisNode.node.right, that)
		if err != nil {
			return empty, err
		}
		return &pnodeLoc{node: &pnode{
			item:  *thisItem,
			left:  thisNode.node.left,
			right: *newRight,
		}}, nil
	}

	newLeft, err := o.join(this, &thatNode.node.left)
	if err != nil {
		return empty, err
	}
	return &pnodeLoc{node: &pnode{
		item:  *thatItem,
		left:  *newLeft,
		right: thatNode.node.right,
	}}, nil
}

func (o *Store) edge(t *PTreap, withValue bool, cfn func(*pnode) *pnodeLoc) (
	*PItem, error) {
	n, err := o.loadNodeLoc(&t.root)
	if err != nil || n.isEmpty() {
		return nil, err
	}
	for {
		child, err := o.loadNodeLoc(cfn(n.node))
		if err != nil {
			return nil, err
		}
		if child.isEmpty() {
			i, err := o.loadItemLoc(&n.node.item, false)
			if err == nil {
				if withValue {
					i, err = o.loadItemLoc(i, true)
				}
				if err == nil {
					return i.item, nil
				}
			}
			return nil, err
		}
		n = child
	}
	return nil, nil
}

func (o *Store) visitAscendNode(t *PTreap, n *pnodeLoc, target []byte,
	withValue bool, visitor PItemVisitor) (bool, error) {
	nNode, err := o.loadNodeLoc(n)
	if err != nil {
		return false, err
	}
	if nNode.isEmpty() {
		return true, nil
	}
	nItem, err := o.loadItemLoc(&nNode.node.item, false)
	if err != nil {
		return false, err
	}
	if t.compare(target, nItem.item.Key) <= 0 {
		keepGoing, err := o.visitAscendNode(t, &nNode.node.left, target, withValue, visitor)
		if err != nil || !keepGoing {
			return false, err
		}
		if withValue {
			nItem, err = o.loadItemLoc(nItem, true)
			if err != nil {
				return false, err
			}
		}
		if !visitor(nItem.item) {
			return false, nil
		}
	}
	return o.visitAscendNode(t, &nNode.node.right, target, withValue, visitor)
}
<commit_msg>More compact code.<commit_after>package gtreap

import (
	"errors"
	"os"
)

// A persistent store holding collections of ordered keys & values.
// The persistence is append-only based on immutable, copy-on-write
// treaps.  This implementation is single-threaded, so users should
// serialize their accesses.
//
// TODO: use atomic.CAS and unsafe.Pointers for safe snapshot'ability.
// TODO: allow read-only snapshots.
//
type Store struct {
	coll map[string]*PTreap
	file *os.File
}

func NewStore(file *os.File) (*Store, error) {
	if file == nil { // Return a memory-only Store.
		return &Store{coll: make(map[string]*PTreap)}, nil
	}
	return nil, errors.New("not implemented yet")
}

func (s *Store) AddCollection(name string, compare KeyCompare) *PTreap {
	if s.coll[name] == nil {
		s.coll[name] = &PTreap{store: s, compare: compare}
	}
	return s.coll[name]
}

func (s *Store) GetCollection(name string) *PTreap {
	return s.coll[name]
}

func (s *Store) RemoveCollection(name string) {
	delete(s.coll, name)
}

// User-supplied key comparison func should return 0 if a == b,
// -1 if a < b, and +1 if a > b.
type KeyCompare func(a, b []byte) int

// A persisted treap.
type PTreap struct {
	store   *Store
	compare KeyCompare
	root    pnodeLoc
}

// A persisted node and its persistence location.
type pnodeLoc struct {
	loc  *ploc  // Can be nil if node is dirty (not yet persisted).
	node *pnode // Can be nil if node is not fetched into memory yet.
}

func (nloc *pnodeLoc) isEmpty() bool {
	return nloc == nil || (nloc.loc == nil && nloc.node == nil)
}

var empty = &pnodeLoc{}

// A persisted node.
type pnode struct {
	item        PItemLoc
	left, right pnodeLoc
}

// A persisted item.
type PItem struct {
	Key, Val []byte // Val may be nil if not fetched into memory yet.
	Priority int32
}

// A persisted item and its persistence location.
type PItemLoc struct {
	loc  *ploc  // Can be nil if item is dirty (not yet persisted).
	item *PItem // Can be nil if item is not fetched into memory yet.
}

// Offset/location of persisted range of bytes.
type ploc struct {
	offset int64  // Usable for os.Seek/ReadAt/WriteAt() at file offset 0.
	length uint32 // Number of bytes.
}

func (t *PTreap) Get(key []byte, withValue bool) (*PItem, error) {
	n, err := t.store.loadNodeLoc(&t.root)
	for {
		if err != nil || n.isEmpty() {
			break
		}
		i, err := t.store.loadItemLoc(&n.node.item, false)
		if err != nil {
			return nil, err
		}
		if i == nil || i.item == nil || i.item.Key == nil {
			panic("no item after loadMetaItemLoc() in get()")
		}
		c := t.compare(key, i.item.Key)
		if c < 0 {
			n, err = t.store.loadNodeLoc(&n.node.left)
		} else if c > 0 {
			n, err = t.store.loadNodeLoc(&n.node.right)
		} else {
			if withValue {
				t.store.loadItemLoc(i, true)
			}
			return i.item, nil
		}
	}
	return nil, err
}

// Replace or insert an item of a given key.
func (t *PTreap) Upsert(item *PItem) error {
	r, err := t.store.union(t, &t.root,
		&pnodeLoc{node: &pnode{item: PItemLoc{item: &PItem{
			Key:      item.Key,
			Val:      item.Val,
			Priority: item.Priority,
		}}}})
	if err == nil {
		t.root = *r
	}
	return err
}

func (t *PTreap) Delete(key []byte) error {
	left, _, right, err := t.store.split(t, &t.root, key)
	if err == nil {
		r, err := t.store.join(left, right)
		if err == nil {
			t.root = *r
		}
	}
	return err
}

func (t *PTreap) Min(withValue bool) (*PItem, error) {
	return t.store.edge(t, withValue, func(n *pnode) *pnodeLoc { return &n.left })
}

func (t *PTreap) Max(withValue bool) (*PItem, error) {
	return t.store.edge(t, withValue, func(n *pnode) *pnodeLoc { return &n.right })
}

type PItemVisitor func(i *PItem) bool

// Visit items greater-than-or-equal to the target.
func (t *PTreap) VisitAscend(target []byte, withValue bool, visitor PItemVisitor) error {
	_, err := t.store.visitAscendNode(t, &t.root, target, withValue, visitor)
	return err
}

func (o *Store) loadNodeLoc(nloc *pnodeLoc) (*pnodeLoc, error) {
	if nloc != nil && nloc.node == nil && nloc.loc != nil {
		// TODO.
	}
	return nloc, nil
}

func (o *Store) loadItemLoc(iloc *PItemLoc, withValue bool) (*PItemLoc, error) {
	if iloc != nil && iloc.item == nil && iloc.loc != nil {
		// TODO.
	}
	return iloc, nil
}

func (o *Store) union(t *PTreap, this *pnodeLoc, that *pnodeLoc) (*pnodeLoc, error) {
	thisNode, err := o.loadNodeLoc(this)
	if err != nil {
		return empty, err
	}
	thatNode, err := o.loadNodeLoc(that)
	if err != nil {
		return empty, err
	}

	if thisNode.isEmpty() {
		return that, nil
	}
	if thatNode.isEmpty() {
		return this, nil
	}

	thisItem, err := o.loadItemLoc(&thisNode.node.item, false)
	if err != nil {
		return empty, err
	}
	thatItem, err := o.loadItemLoc(&thatNode.node.item, false)
	if err != nil {
		return empty, err
	}

	if thisItem.item.Priority > thatItem.item.Priority {
		left, middle, right, err := o.split(t, that, thisItem.item.Key)
		if err != nil {
			return empty, err
		}
		if middle.isEmpty() {
			newLeft, err := o.union(t, &thisNode.node.left, left)
			if err != nil {
				return empty, err
			}
			newRight, err := o.union(t, &thisNode.node.right, right)
			if err != nil {
				return nil, err
			}
			return &pnodeLoc{node: &pnode{
				item:  *thisItem,
				left:  *newLeft,
				right: *newRight,
			}}, nil
		}

		newLeft, err := o.union(t, &thisNode.node.left, left)
		if err != nil {
			return empty, err
		}
		newRight, err := o.union(t, &thisNode.node.right, right)
		if err != nil {
			return empty, err
		}
		return &pnodeLoc{node: &pnode{
			item:  middle.node.item,
			left:  *newLeft,
			right: *newRight,
		}}, nil
	}

	// We don't use middle because the "that" node has precendence.
	left, _, right, err := o.split(t, this, thatItem.item.Key)
	if err != nil {
		return empty, err
	}
	newLeft, err := o.union(t, left, &thatNode.node.left)
	if err != nil {
		return empty, err
	}
	newRight, err := o.union(t, right, &thatNode.node.right)
	if err != nil {
		return empty, err
	}
	return &pnodeLoc{node: &pnode{
		item:  *thatItem,
		left:  *newLeft,
		right: *newRight,
	}}, nil
}

// Splits a treap into two treaps based on a split key "s".  The
// result is (left, middle, right), where left treap has keys < s,
// right treap has keys > s, and middle is either...
// * empty/nil - meaning key s was not in the original treap.
// * non-empty - returning the pnodeLoc that had item s.
func (o *Store) split(t *PTreap, n *pnodeLoc, s []byte) (
	*pnodeLoc, *pnodeLoc, *pnodeLoc, error) {
	nNode, err := o.loadNodeLoc(n)
	if err != nil || nNode.isEmpty() {
		return empty, empty, empty, err
	}
	nItem, err := o.loadItemLoc(&nNode.node.item, false)
	if err != nil {
		return empty, empty, empty, err
	}

	c := t.compare(s, nItem.item.Key)
	if c == 0 {
		return &nNode.node.left, n, &nNode.node.right, nil
	}

	if c < 0 {
		left, middle, right, err := o.split(t, &nNode.node.left, s)
		if err != nil {
			return empty, empty, empty, err
		}
		return left, middle, &pnodeLoc{node: &pnode{
			item:  *nItem,
			left:  *right,
			right: nNode.node.right,
		}}, nil
	}

	left, middle, right, err := o.split(t, &nNode.node.right, s)
	if err != nil {
		return empty, empty, empty, err
	}
	return &pnodeLoc{node: &pnode{
		item:  *nItem,
		left:  nNode.node.left,
		right: *left,
	}}, middle, right, nil
}

// All the keys from this are < keys from that.
func (o *Store) join(this *pnodeLoc, that *pnodeLoc) (*pnodeLoc, error) {
	thisNode, err := o.loadNodeLoc(this)
	if err == nil {
		thatNode, err := o.loadNodeLoc(that)
		if err == nil {
			if thisNode.isEmpty() {
				return that, nil
			}
			if thatNode.isEmpty() {
				return this, nil
			}
			thisItem, err := o.loadItemLoc(&thisNode.node.item, false)
			if err == nil {
				thatItem, err := o.loadItemLoc(&thatNode.node.item, false)
				if err == nil {
					if thisItem.item.Priority > thatItem.item.Priority {
						newRight, err := o.join(&thisNode.node.right, that)
						if err == nil {
							return &pnodeLoc{node: &pnode{
								item:  *thisItem,
								left:  thisNode.node.left,
								right: *newRight,
							}}, nil
						}
					} else {
						newLeft, err := o.join(this, &thatNode.node.left)
						if err == nil {
							return &pnodeLoc{node: &pnode{
								item:  *thatItem,
								left:  *newLeft,
								right: thatNode.node.right,
							}}, nil
						}
					}
				}
			}
		}
	}
	return empty, err
}

func (o *Store) edge(t *PTreap, withValue bool, cfn func(*pnode) *pnodeLoc) (
	*PItem, error) {
	n, err := o.loadNodeLoc(&t.root)
	if err != nil || n.isEmpty() {
		return nil, err
	}
	for {
		child, err := o.loadNodeLoc(cfn(n.node))
		if err != nil {
			return nil, err
		}
		if child.isEmpty() {
			i, err := o.loadItemLoc(&n.node.item, false)
			if err == nil {
				if withValue {
					i, err = o.loadItemLoc(i, true)
				}
				if err == nil {
					return i.item, nil
				}
			}
			return nil, err
		}
		n = child
	}
	return nil, nil
}

func (o *Store) visitAscendNode(t *PTreap, n *pnodeLoc, target []byte,
	withValue bool, visitor PItemVisitor) (bool, error) {
	nNode, err := o.loadNodeLoc(n)
	if err != nil {
		return false, err
	}
	if nNode.isEmpty() {
		return true, nil
	}
	nItem, err := o.loadItemLoc(&nNode.node.item, false)
	if err != nil {
		return false, err
	}
	if t.compare(target, nItem.item.Key) <= 0 {
		keepGoing, err := o.visitAscendNode(t, &nNode.node.left, target, withValue, visitor)
		if err != nil || !keepGoing {
			return false, err
		}
		if withValue {
			nItem, err = o.loadItemLoc(nItem, true)
			if err != nil {
				return false, err
			}
		}
		if !visitor(nItem.item) {
			return false, nil
		}
	}
	return o.visitAscendNode(t, &nNode.node.right, target, withValue, visitor)
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, Christopher Patton.
// All rights reserved.
package store

// TODO(me) Missing features:
//  * PubStore) <--> protobuf
//  * NewPrivStore(K []byte, params *StoreParams) (*PrivStore, error)

/*
// The next line gets things going on Mac:
#cgo CPPFLAGS: -I/usr/local/opt/openssl/include
#cgo LDFLAGS: -lstruct -lcrypto
#include <struct/const.h>
#include <struct/dict.h>
#include "string.h"

char **new_str_list(int len) {
	return calloc(sizeof(char *), len);
}

int *new_int_list(int len) {
	return calloc(sizeof(int), len);
}

void set_str_list(char **list, int idx, char *val) {
	list[idx] = val;
}

void set_int_list(int *list, int idx, int val) {
	list[idx] = val;
}

int get_int_list(int *list, int idx) {
	return list[idx];
}

void free_str_list(char **list, int len) {
	int i;
	for (i = 0; i < len; i++) {
		if (list[i] != NULL) {
			free(list[i]);
		}
	}
	free(list);
}

void free_int_list(int *list) {
	free(list);
}

char *get_row_ptr(char *table, int row, int row_bytes) {
	return &table[row * row_bytes];
}
*/
import "C"
import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"unsafe"

	"golang.org/x/crypto/pbkdf2"
)

// Number of bytes to use for the salt. The salt is a random string used to
// construct the table. It is prepended to the input of each HMAC call.
const SaltBytes = 8

// Number of row bytes allocated for the tag.
const TagBytes = 3

// The maximum length of the row. In general, the length of the row depends on
// the length of the longest output in the map. HASH_BYTES is defined in
// c/const.h.
const MaxRowBytes = C.HASH_BYTES

// The maximum length of the outputs. 1 byte of each row is allocated for
// padding the output string.
const MaxOutputBytes = MaxRowBytes - TagBytes - 1

// Length of the HMAC key. HMAC_KEY_BYTES is defined in c/const.h.
const KeyBytes = C.HMAC_KEY_BYTES

type Error string

func (err Error) Error() string {
	return string(err)
}

// Returned by Get() and priv.GetValue() if the input was not found in the
// dictionary.
const ItemNotFound = Error("item not found")

// Returned by GetRow() in case idx not in the table index.
const ErrorIdx = Error("index out of range")

// CError propagates an error from the internal C code.
func CError(fn string, errNo C.int) Error {
	return Error(fmt.Sprintf("%s returns error %d", fn, errNo))
}

// Returned by New(), Get(), priv.GetIdx(), or priv.GetValue() if the C
// implementation of HMAC returns an error.
const ErrorHMAC = Error("HMAC failed")

// The public parameters of Store, needed by both the client and server.
type StoreParams struct {
	TableLen       int
	MaxOutputBytes int
	RowBytes       int
	TagBytes       int
	Salt           []byte
}

// The public representation of the map.
type PubStore struct {
	dict *C.cdict_t
}

// The private state required for evaluation queries.
type PrivStore struct {
	tinyCtx *C.tiny_ctx
	params  C.dict_params_t
}

// GenerateKey generates a fresh, random key and returns it.
func GenerateKey() []byte {
	K := make([]byte, KeyBytes)
	_, err := rand.Read(K)
	if err != nil {
		return nil
	}
	return K
}

// DeriveKeyFromPassword derives a key from a password and (optional) salt and
// returns it.
//
// NOTE The salt is not the same as StoreParams.Salt. StoreParams.Salt is
// generated by New(), which in turn depends on the key.
func DeriveKeyFromPassword(password, salt []byte) []byte {
	return pbkdf2.Key(password, salt, 4096, KeyBytes, sha256.New)
}

// New generates a new structure (pub, priv) for the map M and key K.
//
// NOTE You must call pub.Free() and priv.Free() before these variables go out
// of scope. These structures contain C types that were allocated on the heap
// and must be freed before losing a reference to them.
func New(K []byte, M map[string]string) (*PubStore, *PrivStore, error) {

	// Check that K is the right length.
	if len(K) != KeyBytes {
		return nil, nil, fmt.Errorf("len(K) = %d, expected %d", len(K), KeyBytes)
	}

	pub := new(PubStore)
	priv := new(PrivStore)

	// Copy input/output pairs into C land.
	itemCt := C.int(len(M))
	inputs := C.new_str_list(itemCt)
	inputBytes := C.new_int_list(itemCt)
	outputs := C.new_str_list(itemCt)
	outputBytes := C.new_int_list(itemCt)
	defer C.free_str_list(inputs, itemCt)
	defer C.free_str_list(outputs, itemCt)
	defer C.free_int_list(inputBytes)
	defer C.free_int_list(outputBytes)

	maxOutputueBytes := 0
	i := C.int(0)
	for input, output := range M {
		if len(output) > maxOutputueBytes {
			maxOutputueBytes = len(output)
		}
		// NOTE C.CString() copies all the bytes of its input, even if it
		// encounters a null byte.
		C.set_str_list(inputs, i, C.CString(input))
		C.set_int_list(inputBytes, i, C.int(len(input)))
		C.set_str_list(outputs, i, C.CString(output))
		C.set_int_list(outputBytes, i, C.int(len(output)))
		i++
	}

	tableLen := C.dict_compute_table_length(C.int(len(M)))
	dict := C.dict_new(
		tableLen,
		C.int(maxOutputueBytes),
		C.int(TagBytes),
		C.int(SaltBytes))
	if dict == nil {
		return nil, nil, Error(fmt.Sprintf("maxOutputBytes > %d", MaxOutputBytes))
	}
	defer C.dict_free(dict)

	priv.tinyCtx = C.tinyprf_new(tableLen)
	if priv.tinyCtx == nil {
		return nil, nil, Error("tableLen < 2")
	}

	cK := C.CString(string(K))
	defer C.free(unsafe.Pointer(cK))
	errNo := C.tinyprf_init(priv.tinyCtx, cK)
	if errNo != C.OK {
		priv.Free()
		return nil, nil, CError("tinyprf_init", errNo)
	}

	// Create the dictionary.
	errNo = C.dict_create(
		dict, priv.tinyCtx, inputs, inputBytes, outputs, outputBytes, itemCt)
	if errNo != C.OK {
		priv.Free()
		return nil, nil, CError("dict_create", errNo)
	}

	// Create compressed representation (no 0 rows).
	pub.dict = C.dict_compress(dict)

	// Copy parameters to priv.
	priv.params.table_length = pub.dict.params.table_length
	priv.params.max_value_bytes = pub.dict.params.max_value_bytes
	priv.params.tag_bytes = pub.dict.params.tag_bytes
	priv.params.row_bytes = pub.dict.params.row_bytes
	priv.params.salt_bytes = pub.dict.params.salt_bytes
	priv.params.salt = (*C.char)(C.malloc(C.size_t(pub.dict.params.salt_bytes + 1)))
	C.memcpy(unsafe.Pointer(priv.params.salt),
		unsafe.Pointer(pub.dict.params.salt),
		C.size_t(priv.params.salt_bytes))

	return pub, priv, nil
}

// Get queries input on the structure (pub, priv). The result is M[input] =
// output, where M is the map represented by (pub, priv).
func Get(pub *PubStore, priv *PrivStore, input string) (string, error) {
	cInput := C.CString(input)
	// FIXME(me) Better way to do the following?
	cOutput := C.CString(string(make([]byte, pub.dict.params.max_value_bytes)))
	cOutputBytes := C.int(0)
	defer C.free(unsafe.Pointer(cInput))
	defer C.free(unsafe.Pointer(cOutput))
	errNo := C.cdict_get(
		pub.dict, priv.tinyCtx, cInput, C.int(len(input)), cOutput, &cOutputBytes)
	if errNo == C.ERR_DICT_BAD_KEY {
		return "", ItemNotFound
	} else if errNo != C.OK {
		return "", CError("cdict_get", errNo)
	}
	return C.GoStringN(cOutput, cOutputBytes), nil
}

// GetRow returns the row of the table associated with idx.
func (pub *PubStore) GetRow(idx int) ([]byte, error) {
	if idx < 0 || idx >= int(pub.dict.params.table_length) {
		return nil, ErrorIdx
	}
	realIdx := C.cdict_binsearch(pub.dict, C.int(idx), 0,
		pub.dict.compressed_table_length)
	return pub.getRealRow(realIdx), nil
}

// GetTable copies the table to a new [][]byte and returns it.
func (pub *PubStore) GetTable() [][]byte {
	table := make([][]byte, pub.dict.compressed_table_length)
	for i := 0; i < int(pub.dict.compressed_table_length); i++ {
		table[i] = pub.getRealRow(C.int(i))
	}
	return table
}

// GetTableIdx copies the table index to a new []int and returns it.
func (pub *PubStore) GetTableIdx() []int {
	tableIdx := make([]int, pub.dict.compressed_table_length)
	for i := 0; i < int(pub.dict.compressed_table_length); i++ {
		tableIdx[i] = int(C.get_int_list(pub.dict.idx, C.int(i)))
	}
	return tableIdx
}

// ToString returns a string representation of the table.
func (pub *PubStore) ToString() string {
	table := pub.GetTable()
	idx := pub.GetTableIdx()
	str := ""
	for i := 0; i < len(table); i++ {
		str += fmt.Sprintf("%-3d %s\n", idx[i], hex.EncodeToString(table[i]))
	}
	return str
}

// GetParams returns the public parameters of the data structure.
func (pub *PubStore) GetParams() *StoreParams {
	return cParamsToStoreParams(&pub.dict.params)
}

// Free deallocates memory associated with the underlying C implementation of
// the data structure.
func (pub *PubStore) Free() {
	C.cdict_free(pub.dict)
}

// GetIdx computes the two indices of the table associated with input and
// returns them.
func (priv *PrivStore) GetIdx(input string) (int, int, error) {
	cInput := C.CString(input)
	defer C.free(unsafe.Pointer(cInput))
	var x, y C.int
	errNo := C.dict_compute_rows(
		priv.params, priv.tinyCtx, cInput, C.int(len(input)), &x, &y)
	if errNo != C.OK {
		return 0, 0, CError("dict_compute_rows", errNo)
	}
	return int(x), int(y), nil
}

// GetValue computes the output associated with the input and the table rows.
func (priv *PrivStore) GetValue(input string, rows [][]byte) (string, error) {
	cInput := C.CString(input)
	// FIXME(me) Better way to do the following?
	cOutput := C.CString(string(make([]byte, priv.params.max_value_bytes)))
	defer C.free(unsafe.Pointer(cInput))
	defer C.free(unsafe.Pointer(cOutput))
	cOutputBytes := C.int(0)

	xRow := C.CString(string(rows[0]))
	yRow := C.CString(string(rows[1]))
	defer C.free(unsafe.Pointer(xRow))
	defer C.free(unsafe.Pointer(yRow))

	errNo := C.dict_compute_value(priv.params, priv.tinyCtx, cInput,
		C.int(len(input)), xRow, yRow, cOutput, &cOutputBytes)

	if errNo == C.ERR_DICT_BAD_KEY {
		return "", ItemNotFound
	} else if errNo != C.OK {
		return "", CError("dict_compute_value", errNo)
	}
	return C.GoStringN(cOutput, cOutputBytes), nil
}

// GetParams returns the public parameters of the data structure.
func (priv *PrivStore) GetParams() *StoreParams {
	return cParamsToStoreParams(&priv.params)
}

// Free deallocates moemory associated with the C implementation of the
// underlying data structure.
func (priv *PrivStore) Free() {
	C.free(unsafe.Pointer(priv.params.salt))
	C.tinyprf_free(priv.tinyCtx)
}

// Returns true if the first saltBytes of *a and *b are equal.
func cBytesToString(str *C.char, bytes C.int) string {
	return C.GoStringN(str, bytes)
}

// cParamsToStoreParams creates *StoreParams from a *C.dict_params_t, making a
// deep copy of the salt.
//
// Called by pub.GetParams() and priv.GetParams().
func cParamsToStoreParams(cParams *C.dict_params_t) *StoreParams {
	params := new(StoreParams)
	params.TableLen = int(cParams.table_length)
	params.MaxOutputBytes = int(cParams.max_value_bytes)
	params.RowBytes = int(cParams.row_bytes)
	params.TagBytes = int(cParams.tag_bytes)
	params.Salt = C.GoBytes(unsafe.Pointer(cParams.salt), cParams.salt_bytes)
	return params
}

// getRealRow copies a row of the table and returns it.
func (pub *PubStore) getRealRow(idx C.int) []byte {
	rowPtr := C.get_row_ptr(pub.dict.table, idx, pub.dict.params.row_bytes)
	return C.GoBytes(unsafe.Pointer(rowPtr), pub.dict.params.row_bytes)
}
<commit_msg>Modify pub.GetTable().<commit_after>// Copyright (c) 2017, Christopher Patton.
// All rights reserved.
package store

// TODO(me) Missing features:
//  * PubStore) <--> protobuf
//  * NewPrivStore(K []byte, params *StoreParams) (*PrivStore, error)

/*
// The next line gets things going on Mac:
#cgo CPPFLAGS: -I/usr/local/opt/openssl/include
#cgo LDFLAGS: -lstruct -lcrypto
#include <struct/const.h>
#include <struct/dict.h>
#include "string.h"

char **new_str_list(int len) {
	return calloc(sizeof(char *), len);
}

int *new_int_list(int len) {
	return calloc(sizeof(int), len);
}

void set_str_list(char **list, int idx, char *val) {
	list[idx] = val;
}

void set_int_list(int *list, int idx, int val) {
	list[idx] = val;
}

int get_int_list(int *list, int idx) {
	return list[idx];
}

void free_str_list(char **list, int len) {
	int i;
	for (i = 0; i < len; i++) {
		if (list[i] != NULL) {
			free(list[i]);
		}
	}
	free(list);
}

void free_int_list(int *list) {
	free(list);
}

char *get_row_ptr(char *table, int row, int row_bytes) {
	return &table[row * row_bytes];
}
*/
import "C"
import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"unsafe"

	"golang.org/x/crypto/pbkdf2"
)

// Number of bytes to use for the salt. The salt is a random string used to
// construct the table. It is prepended to the input of each HMAC call.
const SaltBytes = 8

// Number of row bytes allocated for the tag.
const TagBytes = 3

// The maximum length of the row. In general, the length of the row depends on
// the length of the longest output in the map. HASH_BYTES is defined in
// c/const.h.
const MaxRowBytes = C.HASH_BYTES

// The maximum length of the outputs. 1 byte of each row is allocated for
// padding the output string.
const MaxOutputBytes = MaxRowBytes - TagBytes - 1

// Length of the HMAC key. HMAC_KEY_BYTES is defined in c/const.h.
const KeyBytes = C.HMAC_KEY_BYTES

type Error string

func (err Error) Error() string {
	return string(err)
}

// Returned by Get() and priv.GetValue() if the input was not found in the
// dictionary.
const ItemNotFound = Error("item not found")

// Returned by GetRow() in case idx not in the table index.
const ErrorIdx = Error("index out of range")

// CError propagates an error from the internal C code.
func CError(fn string, errNo C.int) Error {
	return Error(fmt.Sprintf("%s returns error %d", fn, errNo))
}

// Returned by New(), Get(), priv.GetIdx(), or priv.GetValue() if the C
// implementation of HMAC returns an error.
const ErrorHMAC = Error("HMAC failed")

// The public parameters of Store, needed by both the client and server.
type StoreParams struct {
	TableLen       int
	MaxOutputBytes int
	RowBytes       int
	TagBytes       int
	Salt           []byte
}

// The public representation of the map.
type PubStore struct {
	dict *C.cdict_t
}

// The private state required for evaluation queries.
type PrivStore struct {
	tinyCtx *C.tiny_ctx
	params  C.dict_params_t
}

// GenerateKey generates a fresh, random key and returns it.
func GenerateKey() []byte {
	K := make([]byte, KeyBytes)
	_, err := rand.Read(K)
	if err != nil {
		return nil
	}
	return K
}

// DeriveKeyFromPassword derives a key from a password and (optional) salt and
// returns it.
//
// NOTE The salt is not the same as StoreParams.Salt. StoreParams.Salt is
// generated by New(), which in turn depends on the key.
func DeriveKeyFromPassword(password, salt []byte) []byte {
	return pbkdf2.Key(password, salt, 4096, KeyBytes, sha256.New)
}

// New generates a new structure (pub, priv) for the map M and key K.
//
// NOTE You must call pub.Free() and priv.Free() before these variables go out
// of scope. These structures contain C types that were allocated on the heap
// and must be freed before losing a reference to them.
func New(K []byte, M map[string]string) (*PubStore, *PrivStore, error) {

	// Check that K is the right length.
	if len(K) != KeyBytes {
		return nil, nil, fmt.Errorf("len(K) = %d, expected %d", len(K), KeyBytes)
	}

	pub := new(PubStore)
	priv := new(PrivStore)

	// Copy input/output pairs into C land.
	itemCt := C.int(len(M))
	inputs := C.new_str_list(itemCt)
	inputBytes := C.new_int_list(itemCt)
	outputs := C.new_str_list(itemCt)
	outputBytes := C.new_int_list(itemCt)
	defer C.free_str_list(inputs, itemCt)
	defer C.free_str_list(outputs, itemCt)
	defer C.free_int_list(inputBytes)
	defer C.free_int_list(outputBytes)

	maxOutputueBytes := 0
	i := C.int(0)
	for input, output := range M {
		if len(output) > maxOutputueBytes {
			maxOutputueBytes = len(output)
		}
		// NOTE C.CString() copies all the bytes of its input, even if it
		// encounters a null byte.
		C.set_str_list(inputs, i, C.CString(input))
		C.set_int_list(inputBytes, i, C.int(len(input)))
		C.set_str_list(outputs, i, C.CString(output))
		C.set_int_list(outputBytes, i, C.int(len(output)))
		i++
	}

	tableLen := C.dict_compute_table_length(C.int(len(M)))
	dict := C.dict_new(
		tableLen,
		C.int(maxOutputueBytes),
		C.int(TagBytes),
		C.int(SaltBytes))
	if dict == nil {
		return nil, nil, Error(fmt.Sprintf("maxOutputBytes > %d", MaxOutputBytes))
	}
	defer C.dict_free(dict)

	priv.tinyCtx = C.tinyprf_new(tableLen)
	if priv.tinyCtx == nil {
		return nil, nil, Error("tableLen < 2")
	}

	cK := C.CString(string(K))
	defer C.free(unsafe.Pointer(cK))
	errNo := C.tinyprf_init(priv.tinyCtx, cK)
	if errNo != C.OK {
		priv.Free()
		return nil, nil, CError("tinyprf_init", errNo)
	}

	// Create the dictionary.
	errNo = C.dict_create(
		dict, priv.tinyCtx, inputs, inputBytes, outputs, outputBytes, itemCt)
	if errNo != C.OK {
		priv.Free()
		return nil, nil, CError("dict_create", errNo)
	}

	// Create compressed representation (no 0 rows).
	pub.dict = C.dict_compress(dict)

	// Copy parameters to priv.
	priv.params.table_length = pub.dict.params.table_length
	priv.params.max_value_bytes = pub.dict.params.max_value_bytes
	priv.params.tag_bytes = pub.dict.params.tag_bytes
	priv.params.row_bytes = pub.dict.params.row_bytes
	priv.params.salt_bytes = pub.dict.params.salt_bytes
	priv.params.salt = (*C.char)(C.malloc(C.size_t(pub.dict.params.salt_bytes + 1)))
	C.memcpy(unsafe.Pointer(priv.params.salt),
		unsafe.Pointer(pub.dict.params.salt),
		C.size_t(priv.params.salt_bytes))

	return pub, priv, nil
}

// Get queries input on the structure (pub, priv). The result is M[input] =
// output, where M is the map represented by (pub, priv).
func Get(pub *PubStore, priv *PrivStore, input string) (string, error) {
	cInput := C.CString(input)
	// NOTE(me) Better way to do the following?
	cOutput := C.CString(string(make([]byte, pub.dict.params.max_value_bytes)))
	cOutputBytes := C.int(0)
	defer C.free(unsafe.Pointer(cInput))
	defer C.free(unsafe.Pointer(cOutput))
	errNo := C.cdict_get(
		pub.dict, priv.tinyCtx, cInput, C.int(len(input)), cOutput, &cOutputBytes)
	if errNo == C.ERR_DICT_BAD_KEY {
		return "", ItemNotFound
	} else if errNo != C.OK {
		return "", CError("cdict_get", errNo)
	}
	return C.GoStringN(cOutput, cOutputBytes), nil
}

// GetRow returns the row of the table associated with idx.
func (pub *PubStore) GetRow(idx int) ([]byte, error) {
	if idx < 0 || idx >= int(pub.dict.params.table_length) {
		return nil, ErrorIdx
	}
	realIdx := C.cdict_binsearch(pub.dict, C.int(idx), 0,
		pub.dict.compressed_table_length)
	return pub.getRealRow(realIdx), nil
}

// GetTable copies the table to a new []byte and returns it.
func (pub *PubStore) GetTable() []byte {
	rowBytes := int(pub.dict.params.row_bytes)
	tableLen := int(pub.dict.compressed_table_length)
	table := make([]byte, rowBytes*tableLen)
	for i := 0; i < tableLen; i++ {
		copy(table[i*rowBytes:(i+1)*rowBytes], pub.getRealRow(C.int(i)))
	}
	return table
}

// GetTableIdx copies the table index to a new []int and returns it.
func (pub *PubStore) GetTableIdx() []int {
	tableIdx := make([]int, pub.dict.compressed_table_length)
	for i := 0; i < int(pub.dict.compressed_table_length); i++ {
		tableIdx[i] = int(C.get_int_list(pub.dict.idx, C.int(i)))
	}
	return tableIdx
}

// ToString returns a string representation of the table.
func (pub *PubStore) ToString() string {
	rowBytes := int(pub.dict.params.row_bytes)
	tableLen := int(pub.dict.compressed_table_length)
	table := pub.GetTable()
	idx := pub.GetTableIdx()
	str := ""
	for i := 0; i < tableLen; i++ {
		row := table[i*rowBytes : (i+1)*rowBytes]
		str += fmt.Sprintf("%-3d %s\n", idx[i], hex.EncodeToString(row))
	}
	return str
}

// GetParams returns the public parameters of the data structure.
func (pub *PubStore) GetParams() *StoreParams {
	return cParamsToStoreParams(&pub.dict.params)
}

// Free deallocates memory associated with the underlying C implementation of
// the data structure.
func (pub *PubStore) Free() {
	C.cdict_free(pub.dict)
}

// GetIdx computes the two indices of the table associated with input and
// returns them.
func (priv *PrivStore) GetIdx(input string) (int, int, error) {
	cInput := C.CString(input)
	defer C.free(unsafe.Pointer(cInput))
	var x, y C.int
	errNo := C.dict_compute_rows(
		priv.params, priv.tinyCtx, cInput, C.int(len(input)), &x, &y)
	if errNo != C.OK {
		return 0, 0, CError("dict_compute_rows", errNo)
	}
	return int(x), int(y), nil
}

// GetValue computes the output associated with the input and the table rows.
func (priv *PrivStore) GetValue(input string, rows [][]byte) (string, error) {
	cInput := C.CString(input)
	// NOTE(me) Better way to do the following?
	cOutput := C.CString(string(make([]byte, priv.params.max_value_bytes)))
	defer C.free(unsafe.Pointer(cInput))
	defer C.free(unsafe.Pointer(cOutput))
	cOutputBytes := C.int(0)

	xRow := C.CString(string(rows[0]))
	yRow := C.CString(string(rows[1]))
	defer C.free(unsafe.Pointer(xRow))
	defer C.free(unsafe.Pointer(yRow))

	errNo := C.dict_compute_value(priv.params, priv.tinyCtx, cInput,
		C.int(len(input)), xRow, yRow, cOutput, &cOutputBytes)

	if errNo == C.ERR_DICT_BAD_KEY {
		return "", ItemNotFound
	} else if errNo != C.OK {
		return "", CError("dict_compute_value", errNo)
	}
	return C.GoStringN(cOutput, cOutputBytes), nil
}

// GetParams returns the public parameters of the data structure.
func (priv *PrivStore) GetParams() *StoreParams {
	return cParamsToStoreParams(&priv.params)
}

// Free deallocates moemory associated with the C implementation of the
// underlying data structure.
func (priv *PrivStore) Free() {
	C.free(unsafe.Pointer(priv.params.salt))
	C.tinyprf_free(priv.tinyCtx)
}

// Returns true if the first saltBytes of *a and *b are equal.
func cBytesToString(str *C.char, bytes C.int) string {
	return C.GoStringN(str, bytes)
}

// cParamsToStoreParams creates *StoreParams from a *C.dict_params_t, making a
// deep copy of the salt.
//
// Called by pub.GetParams() and priv.GetParams().
func cParamsToStoreParams(cParams *C.dict_params_t) *StoreParams {
	params := new(StoreParams)
	params.TableLen = int(cParams.table_length)
	params.MaxOutputBytes = int(cParams.max_value_bytes)
	params.RowBytes = int(cParams.row_bytes)
	params.TagBytes = int(cParams.tag_bytes)
	params.Salt = C.GoBytes(unsafe.Pointer(cParams.salt), cParams.salt_bytes)
	return params
}

// getRealRow copies a row of the table and returns it.
func (pub *PubStore) getRealRow(idx C.int) []byte {
	rowPtr := C.get_row_ptr(pub.dict.table, idx, pub.dict.params.row_bytes)
	return C.GoBytes(unsafe.Pointer(rowPtr), pub.dict.params.row_bytes)
}
<|endoftext|>


<commit_before>package glukit

import (
	"appengine"
	"appengine/channel"
	"appengine/datastore"
	"appengine/delay"
	"appengine/taskqueue"
	"appengine/urlfetch"
	"bufio"
	"github.com/alexandre-normand/glukit/app/engine"
	"github.com/alexandre-normand/glukit/app/importer"
	"github.com/alexandre-normand/glukit/app/model"
	"github.com/alexandre-normand/glukit/app/store"
	"github.com/alexandre-normand/glukit/app/util"
	"github.com/alexandre-normand/glukit/lib/drive"
	"github.com/alexandre-normand/glukit/lib/goauth2/oauth"
	"os"
	"time"
)

var processFile = delay.Func("processSingleFile", processSingleFile)
var processDemoFile = delay.Func("processDemoFile", processStaticDemoFile)
var refreshUserData = delay.Func(REFRESH_USER_DATA_FUNCTION_NAME, func(context appengine.Context, userEmail string,
	autoScheduleNextRun bool) {
	context.Criticalf("This function purely exists as a workaround to the \"initialization loop\" error that " +
		"shows up because the function calls itself. This implementation defines the same signature as the " +
		"real one which we define in init() to override this implementation!")
})

const (
	REFRESH_USER_DATA_FUNCTION_NAME = "refreshUserData"
)

func updateUserData(context appengine.Context, userEmail string, autoScheduleNextRun bool) {
	// noop
}

// updateUserData is an async task that searches on Google Drive for dexcom files. It handles some high
// watermark of the last import to avoid downloading already imported files (unless they've been updated).
// It also schedules itself to run again the next day unless the token is invalid.
func deprecatedUpdateUserData(context appengine.Context, userEmail string, autoScheduleNextRun bool) {
	glukitUser, userProfileKey, _, err := store.GetUserData(context, userEmail)
	if _, ok := err.(store.StoreError); err != nil && !ok {
		context.Errorf("We're trying to run an update data task for user [%s] that doesn't exist. "+
			"Got error: %v", userEmail, err)
		return
	}

	transport := &oauth.Transport{
		Config: config(),
		Transport: &urlfetch.Transport{
			Context: context,
		},
		Token: &glukitUser.Token,
	}

	// If the token is expired, try to get a fresh one by doing a refresh (which should use the refresh_token
	if glukitUser.Token.Expired() {
		transport.Token.RefreshToken = glukitUser.RefreshToken
		err := transport.Refresh(context)
		if err != nil {
			context.Errorf("Error updating token for user [%s], let's hope he comes back soon so we can "+
				"get a fresh token: %v", userEmail, err)
			return
		}

		// Update the user with the new token
		context.Infof("Token refreshed, updating user [%s] with token [%v]", userEmail, glukitUser.Token)
		store.StoreUserProfile(context, time.Now(), *glukitUser)
	}

	// Next update in one day
	nextUpdate := time.Now().AddDate(0, 0, 1)
	files, err := importer.SearchDataFiles(transport.Client(), glukitUser.MostRecentRead.GetTime())
	if err != nil {
		context.Warningf("Error while searching for files on google drive for user [%s]: %v", userEmail, err)
	} else {
		switch {
		case len(files) == 0:
			context.Infof("No new or updated data found for existing user [%s]", userEmail)
		case len(files) > 0:
			context.Infof("Found new data files for user [%s], downloading and storing...", userEmail)
			processFileSearchResults(&glukitUser.Token, files, context, userEmail, userProfileKey)
		}
	}

	engine.CalculateGlukitScoreBatch(context, glukitUser)

	if autoScheduleNextRun {
		task, err := refreshUserData.Task(userEmail, autoScheduleNextRun)
		if err != nil {
			context.Criticalf("Couldn't schedule the next execution of the data refresh for user [%s]. "+
				"This breaks background updating of user data!: %v", userEmail, err)
		}
		task.ETA = nextUpdate
		taskqueue.Add(context, task, "refresh")

		context.Infof("Scheduled next data update for user [%s] at [%s]", userEmail, nextUpdate.Format(util.TIMEFORMAT))
	} else {
		context.Infof("Not scheduling a the next refresh as requested by autoScheduleNextRun [%t]", autoScheduleNextRun)
	}
}

// processFileSearchResults reads the list of files detected on google drive and kicks off a new queued task
// to process each one
func processFileSearchResults(token *oauth.Token, files []*drive.File, context appengine.Context, userEmail string,
	userProfileKey *datastore.Key) {
	// TODO : Look at recent file import log for that file and skip to the new data. It would be nice to be able to
	// use the Http Range header but that's unlikely to be possible since new event/read data is spreadout in the
	// file
	for i := range files {
		task, err := processFile.Task(token, files[i], userEmail, userProfileKey)
		if err != nil {
			util.Propagate(err)
		}
		taskqueue.Add(context, task, "store")
	}
}

// processSingleFile handles the import of a single file. It deals with:
//    1. Logging the file import operation
//    2. Calculating and updating the new GlukitScore
//    3. Sending a "refresh" message to any connected client
func processSingleFile(context appengine.Context, token *oauth.Token, file *drive.File, userEmail string,
	userProfileKey *datastore.Key) {
	t := &oauth.Transport{
		Config: config(),
		Transport: &urlfetch.Transport{
			Context: context,
		},
		Token: token,
	}

	reader, err := importer.GetFileReader(context, t, file)
	if err != nil {
		context.Infof("Error reading file %s, skipping: [%v]", file.OriginalFilename, err)
	} else {
		// Default to beginning of time
		startTime := util.GLUKIT_EPOCH_TIME
		if lastFileImportLog, err := store.GetFileImportLog(context, userProfileKey, file.Id); err == nil {
			startTime = lastFileImportLog.LastDataProcessed
			context.Infof("Reloading data from file [%s]-[%s] starting at date [%s]...", file.Id,
				file.OriginalFilename, startTime.Format(util.TIMEFORMAT))
		} else if err == datastore.ErrNoSuchEntity {
			context.Debugf("First import of file [%s]-[%s]...", file.Id, file.OriginalFilename)
		} else if err != nil {
			util.Propagate(err)
		}

		lastReadTime := importer.ParseContent(context, reader, importer.IMPORT_BATCH_SIZE, userProfileKey, startTime,
			store.StoreDaysOfReads, store.StoreDaysOfCarbs, store.StoreDaysOfInjections, store.StoreDaysOfExercises)
		store.LogFileImport(context, userProfileKey, model.FileImportLog{Id: file.Id, Md5Checksum: file.Md5Checksum,
			LastDataProcessed: lastReadTime})
		reader.Close()

		if glukitUser, err := store.GetUserProfile(context, userProfileKey); err != nil {
			context.Warningf("Error getting retrieving GlukitUser [%s], this needs attention: [%v]", userEmail, err)
		} else {
			// Calculate Glukit Score batch here for the newly imported data
			err := engine.CalculateGlukitScoreBatch(context, glukitUser)
			if err != nil {
				context.Warningf("Error starting batch calculation of GlukitScores for [%s], this needs attention: [%v]", userEmail, err)
			}
		}
	}
	channel.Send(context, userEmail, "Refresh")
}

// processStaticDemoFile imports the static resource included with the app for the demo user
func processStaticDemoFile(context appengine.Context, userProfileKey *datastore.Key) {

	// open input file
	fi, err := os.Open("data.xml")
	if err != nil {
		panic(err)
	}
	// close fi on exit and check for its returned error
	defer func() {
		if fi.Close() != nil {
			panic(err)
		}
	}()
	// make a read buffer
	reader := bufio.NewReader(fi)

	lastReadTime := importer.ParseContent(context, reader, importer.IMPORT_BATCH_SIZE, userProfileKey, util.GLUKIT_EPOCH_TIME,
		store.StoreDaysOfReads, store.StoreDaysOfCarbs, store.StoreDaysOfInjections, store.StoreDaysOfExercises)
	store.LogFileImport(context, userProfileKey, model.FileImportLog{Id: "demo", Md5Checksum: "dummychecksum",
		LastDataProcessed: lastReadTime})

	if userProfile, err := store.GetUserProfile(context, userProfileKey); err != nil {
		context.Warningf("Error while persisting score for %s: %v", DEMO_EMAIL, err)
	} else {
		if err := engine.CalculateGlukitScoreBatch(context, userProfile); err != nil {
			context.Warningf("Error while starting batch calculation of glukit scores for %s: %v", DEMO_EMAIL, err)
		}
	}

	channel.Send(context, DEMO_EMAIL, "Refresh")
}
<commit_msg>Temporarily Reenable Legacy Imports.<commit_after>package glukit

import (
	"appengine"
	"appengine/channel"
	"appengine/datastore"
	"appengine/delay"
	"appengine/taskqueue"
	"appengine/urlfetch"
	"bufio"
	"github.com/alexandre-normand/glukit/app/engine"
	"github.com/alexandre-normand/glukit/app/importer"
	"github.com/alexandre-normand/glukit/app/model"
	"github.com/alexandre-normand/glukit/app/store"
	"github.com/alexandre-normand/glukit/app/util"
	"github.com/alexandre-normand/glukit/lib/drive"
	"github.com/alexandre-normand/glukit/lib/goauth2/oauth"
	"os"
	"time"
)

var processFile = delay.Func("processSingleFile", processSingleFile)
var processDemoFile = delay.Func("processDemoFile", processStaticDemoFile)
var refreshUserData = delay.Func(REFRESH_USER_DATA_FUNCTION_NAME, func(context appengine.Context, userEmail string,
	autoScheduleNextRun bool) {
	context.Criticalf("This function purely exists as a workaround to the \"initialization loop\" error that " +
		"shows up because the function calls itself. This implementation defines the same signature as the " +
		"real one which we define in init() to override this implementation!")
})

const (
	REFRESH_USER_DATA_FUNCTION_NAME = "refreshUserData"
)

func disabledUpdateUserData(context appengine.Context, userEmail string, autoScheduleNextRun bool) {
	// noop
}

// updateUserData is an async task that searches on Google Drive for dexcom files. It handles some high
// watermark of the last import to avoid downloading already imported files (unless they've been updated).
// It also schedules itself to run again the next day unless the token is invalid.
func updateUserData(context appengine.Context, userEmail string, autoScheduleNextRun bool) {
	glukitUser, userProfileKey, _, err := store.GetUserData(context, userEmail)
	if _, ok := err.(store.StoreError); err != nil && !ok {
		context.Errorf("We're trying to run an update data task for user [%s] that doesn't exist. "+
			"Got error: %v", userEmail, err)
		return
	}

	transport := &oauth.Transport{
		Config: config(),
		Transport: &urlfetch.Transport{
			Context: context,
		},
		Token: &glukitUser.Token,
	}

	// If the token is expired, try to get a fresh one by doing a refresh (which should use the refresh_token
	if glukitUser.Token.Expired() {
		transport.Token.RefreshToken = glukitUser.RefreshToken
		err := transport.Refresh(context)
		if err != nil {
			context.Errorf("Error updating token for user [%s], let's hope he comes back soon so we can "+
				"get a fresh token: %v", userEmail, err)
			return
		}

		// Update the user with the new token
		context.Infof("Token refreshed, updating user [%s] with token [%v]", userEmail, glukitUser.Token)
		store.StoreUserProfile(context, time.Now(), *glukitUser)
	}

	// Next update in one day
	nextUpdate := time.Now().AddDate(0, 0, 1)
	files, err := importer.SearchDataFiles(transport.Client(), glukitUser.MostRecentRead.GetTime())
	if err != nil {
		context.Warningf("Error while searching for files on google drive for user [%s]: %v", userEmail, err)
	} else {
		switch {
		case len(files) == 0:
			context.Infof("No new or updated data found for existing user [%s]", userEmail)
		case len(files) > 0:
			context.Infof("Found new data files for user [%s], downloading and storing...", userEmail)
			processFileSearchResults(&glukitUser.Token, files, context, userEmail, userProfileKey)
		}
	}

	engine.CalculateGlukitScoreBatch(context, glukitUser)

	if autoScheduleNextRun {
		task, err := refreshUserData.Task(userEmail, autoScheduleNextRun)
		if err != nil {
			context.Criticalf("Couldn't schedule the next execution of the data refresh for user [%s]. "+
				"This breaks background updating of user data!: %v", userEmail, err)
		}
		task.ETA = nextUpdate
		taskqueue.Add(context, task, "refresh")

		context.Infof("Scheduled next data update for user [%s] at [%s]", userEmail, nextUpdate.Format(util.TIMEFORMAT))
	} else {
		context.Infof("Not scheduling a the next refresh as requested by autoScheduleNextRun [%t]", autoScheduleNextRun)
	}
}

// processFileSearchResults reads the list of files detected on google drive and kicks off a new queued task
// to process each one
func processFileSearchResults(token *oauth.Token, files []*drive.File, context appengine.Context, userEmail string,
	userProfileKey *datastore.Key) {
	// TODO : Look at recent file import log for that file and skip to the new data. It would be nice to be able to
	// use the Http Range header but that's unlikely to be possible since new event/read data is spreadout in the
	// file
	for i := range files {
		task, err := processFile.Task(token, files[i], userEmail, userProfileKey)
		if err != nil {
			util.Propagate(err)
		}
		taskqueue.Add(context, task, "store")
	}
}

// processSingleFile handles the import of a single file. It deals with:
//    1. Logging the file import operation
//    2. Calculating and updating the new GlukitScore
//    3. Sending a "refresh" message to any connected client
func processSingleFile(context appengine.Context, token *oauth.Token, file *drive.File, userEmail string,
	userProfileKey *datastore.Key) {
	t := &oauth.Transport{
		Config: config(),
		Transport: &urlfetch.Transport{
			Context: context,
		},
		Token: token,
	}

	reader, err := importer.GetFileReader(context, t, file)
	if err != nil {
		context.Infof("Error reading file %s, skipping: [%v]", file.OriginalFilename, err)
	} else {
		// Default to beginning of time
		startTime := util.GLUKIT_EPOCH_TIME
		if lastFileImportLog, err := store.GetFileImportLog(context, userProfileKey, file.Id); err == nil {
			startTime = lastFileImportLog.LastDataProcessed
			context.Infof("Reloading data from file [%s]-[%s] starting at date [%s]...", file.Id,
				file.OriginalFilename, startTime.Format(util.TIMEFORMAT))
		} else if err == datastore.ErrNoSuchEntity {
			context.Debugf("First import of file [%s]-[%s]...", file.Id, file.OriginalFilename)
		} else if err != nil {
			util.Propagate(err)
		}

		lastReadTime := importer.ParseContent(context, reader, importer.IMPORT_BATCH_SIZE, userProfileKey, startTime,
			store.StoreDaysOfReads, store.StoreDaysOfCarbs, store.StoreDaysOfInjections, store.StoreDaysOfExercises)
		store.LogFileImport(context, userProfileKey, model.FileImportLog{Id: file.Id, Md5Checksum: file.Md5Checksum,
			LastDataProcessed: lastReadTime})
		reader.Close()

		if glukitUser, err := store.GetUserProfile(context, userProfileKey); err != nil {
			context.Warningf("Error getting retrieving GlukitUser [%s], this needs attention: [%v]", userEmail, err)
		} else {
			// Calculate Glukit Score batch here for the newly imported data
			err := engine.CalculateGlukitScoreBatch(context, glukitUser)
			if err != nil {
				context.Warningf("Error starting batch calculation of GlukitScores for [%s], this needs attention: [%v]", userEmail, err)
			}
		}
	}
	channel.Send(context, userEmail, "Refresh")
}

// processStaticDemoFile imports the static resource included with the app for the demo user
func processStaticDemoFile(context appengine.Context, userProfileKey *datastore.Key) {

	// open input file
	fi, err := os.Open("data.xml")
	if err != nil {
		panic(err)
	}
	// close fi on exit and check for its returned error
	defer func() {
		if fi.Close() != nil {
			panic(err)
		}
	}()
	// make a read buffer
	reader := bufio.NewReader(fi)

	lastReadTime := importer.ParseContent(context, reader, importer.IMPORT_BATCH_SIZE, userProfileKey, util.GLUKIT_EPOCH_TIME,
		store.StoreDaysOfReads, store.StoreDaysOfCarbs, store.StoreDaysOfInjections, store.StoreDaysOfExercises)
	store.LogFileImport(context, userProfileKey, model.FileImportLog{Id: "demo", Md5Checksum: "dummychecksum",
		LastDataProcessed: lastReadTime})

	if userProfile, err := store.GetUserProfile(context, userProfileKey); err != nil {
		context.Warningf("Error while persisting score for %s: %v", DEMO_EMAIL, err)
	} else {
		if err := engine.CalculateGlukitScoreBatch(context, userProfile); err != nil {
			context.Warningf("Error while starting batch calculation of glukit scores for %s: %v", DEMO_EMAIL, err)
		}
	}

	channel.Send(context, DEMO_EMAIL, "Refresh")
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/DECK36/go-gelf/gelf"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"
	"time"
)

/*
	http://www.freedesktop.org/software/systemd/man/systemd.journal-fields.html
	https://github.com/Graylog2/graylog2-docs/wiki/GELF
*/
type SystemdJournalEntry struct {
	Cursor                    string `json:"__CURSOR"`
	Realtime_timestamp        int64  `json:"__REALTIME_TIMESTAMP,string"`
	Monotonic_timestamp       string `json:"__MONOTONIC_TIMESTAMP"`
	Boot_id                   string `json:"_BOOT_ID"`
	Transport                 string `json:"_TRANSPORT"`
	Priority                  int32  `json:"PRIORITY,string"`
	Syslog_facility           string `json:"SYSLOG_FACILITY"`
	Syslog_identifier         string `json:"SYSLOG_IDENTIFIER"`
	Message                   string `json:"MESSAGE"`
	Pid                       string `json:"_PID"`
	Uid                       string `json:"_UID"`
	Gid                       string `json:"_GID"`
	Comm                      string `json:"_COMM"`
	Exe                       string `json:"_EXE"`
	Cmdline                   string `json:"_CMDLINE"`
	Systemd_cgroup            string `json:"_SYSTEMD_CGROUP"`
	Systemd_session           string `json:"_SYSTEMD_SESSION"`
	Systemd_owner_uid         string `json:"_SYSTEMD_OWNER_UID"`
	Systemd_unit              string `json:"_SYSTEMD_UNIT"`
	Source_realtime_timestamp string `json:"_SOURCE_REALTIME_TIMESTAMP"`
	Machine_id                string `json:"_MACHINE_ID"`
	Hostname                  string `json:"_HOSTNAME"`
	FullMessage               string `json:"-"`
}

// Strip date from message-content
var startsWithTimestamp = regexp.MustCompile("^20[0-9][0-9][/\\-][01][0-9][/\\-][0123][0-9] [0-2]?[0-9]:[0-5][0-9]:[0-5][0-9][,0-9]{0,3} ")

func (this *SystemdJournalEntry) toGelf() *gelf.Message {
	var extra = map[string]interface{}{
		"Boot_id":      this.Boot_id,
		"Pid":          this.Pid,
		"Uid":          this.Uid,
		"Systemd_unit": this.Systemd_unit,
	}

	if this.isJsonMessage() {
		if err := json.Unmarshal([]byte(this.Message), &extra); err == nil {
			if m, ok := extra["Message"]; ok {
				this.Message = m.(string)
				delete(extra, "Message")
			}

			if f, ok := extra["FullMessage"]; ok {
				this.FullMessage = f.(string)
				delete(extra, "FullMessage")
			}
		}
	} else if -1 != strings.Index(this.Message, "\n") {
		this.FullMessage = this.Message
		this.Message = strings.Split(this.Message, "\n")[0]
	}

	return &gelf.Message{
		Version:  "1.1",
		Host:     this.Hostname,
		Short:    this.Message,
		Full:     this.FullMessage,
		TimeUnix: float64(this.Realtime_timestamp) / 1000 / 1000,
		Level:    this.Priority,
		Facility: this.Syslog_identifier,
		Extra:    extra,
	}
}

// Custom wrapper to support unprintable chars in message
func (this *SystemdJournalEntry) UnmarshalJSON(data []byte) error {
	// use an alias to prevent recursion
	type entryAlias SystemdJournalEntry
	aux := (*entryAlias)(this)

	if err := json.Unmarshal(data, &aux); err == nil {
		this.Message = startsWithTimestamp.ReplaceAllString(this.Message, "")

		return nil
	} else if ute, ok := err.(*json.UnmarshalTypeError); ok && ute.Field == "MESSAGE" && ute.Value == "array" {
		// Include brackets, which is why we subtract and add by one
		len := int64(strings.Index(string(data[ute.Offset:]), `]`)) + 1

		var message []byte
		if err := json.Unmarshal(data[ute.Offset-1:ute.Offset+len], &message); err != nil {
			return err
		}

		// only the failing field is skipped, so we can still use the rest
		this.Message = string(message)

		return nil
	} else {
		return err
	}
}

func (this *SystemdJournalEntry) send() {
	message := this.toGelf()

	for err := writer.WriteMessage(message); err != nil; err = writer.WriteMessage(message) {
		/*
			UDP is nonblocking, but the OS stores an error which GO will return on the next call.
			This means we've already lost a message, but can keep retrying the current one. Sleep to make this less obtrusive
		*/
		fmt.Fprintln(os.Stderr, "send - processing paused because of: "+err.Error())
		time.Sleep(SLEEP_AFTER_ERROR)
	}
}

func (this *SystemdJournalEntry) isJsonMessage() bool {
	return len(this.Message) > 4 && this.Message[0:2] == `{"`
}

type pendingEntry struct {
	sync.RWMutex
	entry *SystemdJournalEntry
}

func (this *pendingEntry) Push(next SystemdJournalEntry) {
	this.Lock()

	if this.entry != nil {
		this.entry.send()
	}

	this.entry = &next
	this.Unlock()
}

func (this *pendingEntry) Clear() {
	if this.entry == nil {
		return
	}

	this.Lock()
	entry := this.entry
	this.entry = nil
	this.Unlock()

	entry.send()
}

func (this *pendingEntry) ClearEvery(interval time.Duration) {
	for {
		time.Sleep(interval)
		this.Clear()
	}
}

var writer *gelf.Writer

const (
	WRITE_INTERVAL             = 50 * time.Millisecond
	SAMESOURCE_TIME_DIFFERENCE = 100 * 1000
	SLEEP_AFTER_ERROR          = 15 * time.Second
)

func main() {
	if len(os.Args) < 3 {
		panic("usage: SystemdJournal2Gelf SERVER:12201 [JOURNALCTL PARAMETERS]")
	}

	if w, err := gelf.NewWriter(os.Args[1]); err != nil {
		panic("while connecting to Graylog server: " + err.Error())
	} else {
		writer = w
	}

	journalArgs := []string{"--all", "--output=json"}
	journalArgs = append(journalArgs, os.Args[2:]...)
	cmd := exec.Command("journalctl", journalArgs...)

	stderr, _ := cmd.StderrPipe()
	stdout, _ := cmd.StdoutPipe()
	go io.Copy(os.Stderr, stderr)
	d := json.NewDecoder(stdout)

	var pending pendingEntry
	go pending.ClearEvery(WRITE_INTERVAL)
	cmd.Start()

	for {
		var entry SystemdJournalEntry
		if err := d.Decode(&entry); err != nil {
			if err == io.EOF {
				break
			}

			cmd.Process.Kill()
			panic("could not parse journal output: " + err.Error())
		}

		pending.Push(entry)

		// Prevent saturation and throttling
		time.Sleep(1 * time.Millisecond)
	}

	cmd.Wait()

	pending.Clear()
}
<commit_msg>Fix #6 - remove fields we don't use from struct definition<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/DECK36/go-gelf/gelf"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"
	"time"
)

/*
	http://www.freedesktop.org/software/systemd/man/systemd.journal-fields.html
	https://github.com/Graylog2/graylog2-docs/wiki/GELF
*/
type SystemdJournalEntry struct {
	Realtime_timestamp        int64  `json:"__REALTIME_TIMESTAMP,string"`
	Boot_id                   string `json:"_BOOT_ID"`
	Priority                  int32  `json:"PRIORITY,string"`
	Syslog_identifier         string `json:"SYSLOG_IDENTIFIER"`
	Message                   string `json:"MESSAGE"`
	Pid                       string `json:"_PID"`
	Uid                       string `json:"_UID"`
	Systemd_unit              string `json:"_SYSTEMD_UNIT"`
	Hostname                  string `json:"_HOSTNAME"`
	FullMessage               string `json:"-"`
}

// Strip date from message-content
var startsWithTimestamp = regexp.MustCompile("^20[0-9][0-9][/\\-][01][0-9][/\\-][0123][0-9] [0-2]?[0-9]:[0-5][0-9]:[0-5][0-9][,0-9]{0,3} ")

func (this *SystemdJournalEntry) toGelf() *gelf.Message {
	var extra = map[string]interface{}{
		"Boot_id":      this.Boot_id,
		"Pid":          this.Pid,
		"Uid":          this.Uid,
		"Systemd_unit": this.Systemd_unit,
	}

	if this.isJsonMessage() {
		if err := json.Unmarshal([]byte(this.Message), &extra); err == nil {
			if m, ok := extra["Message"]; ok {
				this.Message = m.(string)
				delete(extra, "Message")
			}

			if f, ok := extra["FullMessage"]; ok {
				this.FullMessage = f.(string)
				delete(extra, "FullMessage")
			}
		}
	} else if -1 != strings.Index(this.Message, "\n") {
		this.FullMessage = this.Message
		this.Message = strings.Split(this.Message, "\n")[0]
	}

	return &gelf.Message{
		Version:  "1.1",
		Host:     this.Hostname,
		Short:    this.Message,
		Full:     this.FullMessage,
		TimeUnix: float64(this.Realtime_timestamp) / 1000 / 1000,
		Level:    this.Priority,
		Facility: this.Syslog_identifier,
		Extra:    extra,
	}
}

// Custom wrapper to support unprintable chars in message
func (this *SystemdJournalEntry) UnmarshalJSON(data []byte) error {
	// use an alias to prevent recursion
	type entryAlias SystemdJournalEntry
	aux := (*entryAlias)(this)

	if err := json.Unmarshal(data, &aux); err == nil {
		this.Message = startsWithTimestamp.ReplaceAllString(this.Message, "")

		return nil
	} else if ute, ok := err.(*json.UnmarshalTypeError); ok && ute.Field == "MESSAGE" && ute.Value == "array" {
		// Include brackets, which is why we subtract and add by one
		len := int64(strings.Index(string(data[ute.Offset:]), `]`)) + 1

		var message []byte
		if err := json.Unmarshal(data[ute.Offset-1:ute.Offset+len], &message); err != nil {
			return err
		}

		// only the failing field is skipped, so we can still use the rest
		this.Message = string(message)

		return nil
	} else {
		return err
	}
}

func (this *SystemdJournalEntry) send() {
	message := this.toGelf()

	for err := writer.WriteMessage(message); err != nil; err = writer.WriteMessage(message) {
		/*
			UDP is nonblocking, but the OS stores an error which GO will return on the next call.
			This means we've already lost a message, but can keep retrying the current one. Sleep to make this less obtrusive
		*/
		fmt.Fprintln(os.Stderr, "send - processing paused because of: "+err.Error())
		time.Sleep(SLEEP_AFTER_ERROR)
	}
}

func (this *SystemdJournalEntry) isJsonMessage() bool {
	return len(this.Message) > 4 && this.Message[0:2] == `{"`
}

type pendingEntry struct {
	sync.RWMutex
	entry *SystemdJournalEntry
}

func (this *pendingEntry) Push(next SystemdJournalEntry) {
	this.Lock()

	if this.entry != nil {
		this.entry.send()
	}

	this.entry = &next
	this.Unlock()
}

func (this *pendingEntry) Clear() {
	if this.entry == nil {
		return
	}

	this.Lock()
	entry := this.entry
	this.entry = nil
	this.Unlock()

	entry.send()
}

func (this *pendingEntry) ClearEvery(interval time.Duration) {
	for {
		time.Sleep(interval)
		this.Clear()
	}
}

var writer *gelf.Writer

const (
	WRITE_INTERVAL             = 50 * time.Millisecond
	SAMESOURCE_TIME_DIFFERENCE = 100 * 1000
	SLEEP_AFTER_ERROR          = 15 * time.Second
)

func main() {
	if len(os.Args) < 3 {
		panic("usage: SystemdJournal2Gelf SERVER:12201 [JOURNALCTL PARAMETERS]")
	}

	if w, err := gelf.NewWriter(os.Args[1]); err != nil {
		panic("while connecting to Graylog server: " + err.Error())
	} else {
		writer = w
	}

	journalArgs := []string{"--all", "--output=json"}
	journalArgs = append(journalArgs, os.Args[2:]...)
	cmd := exec.Command("journalctl", journalArgs...)

	stderr, _ := cmd.StderrPipe()
	stdout, _ := cmd.StdoutPipe()
	go io.Copy(os.Stderr, stderr)
	d := json.NewDecoder(stdout)

	var pending pendingEntry
	go pending.ClearEvery(WRITE_INTERVAL)
	cmd.Start()

	for {
		var entry SystemdJournalEntry
		if err := d.Decode(&entry); err != nil {
			if err == io.EOF {
				break
			}

			cmd.Process.Kill()
			panic("could not parse journal output: " + err.Error())
		}

		pending.Push(entry)

		// Prevent saturation and throttling
		time.Sleep(1 * time.Millisecond)
	}

	cmd.Wait()

	pending.Clear()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"github.com/DECK36/go-gelf/gelf"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"time"
)

/*
	http://www.freedesktop.org/software/systemd/man/systemd.journal-fields.html
	https://github.com/Graylog2/graylog2-docs/wiki/GELF
*/
type SystemdJournalEntry struct {
	Cursor                    string `json:"__CURSOR"`
	Realtime_timestamp        int64  `json:"__REALTIME_TIMESTAMP,string"`
	Monotonic_timestamp       string `json:"__MONOTONIC_TIMESTAMP"`
	Boot_id                   string `json:"_BOOT_ID"`
	Transport                 string `json:"_TRANSPORT"`
	Priority                  int32  `json:"PRIORITY,string"`
	Syslog_facility           string `json:"SYSLOG_FACILITY"`
	Syslog_identifier         string `json:"SYSLOG_IDENTIFIER"`
	Message                   string `json:"MESSAGE"`
	Pid                       string `json:"_PID"`
	Uid                       string `json:"_UID"`
	Gid                       string `json:"_GID"`
	Comm                      string `json:"_COMM"`
	Exe                       string `json:"_EXE"`
	Cmdline                   string `json:"_CMDLINE"`
	Systemd_cgroup            string `json:"_SYSTEMD_CGROUP"`
	Systemd_session           string `json:"_SYSTEMD_SESSION"`
	Systemd_owner_uid         string `json:"_SYSTEMD_OWNER_UID"`
	Systemd_unit              string `json:"_SYSTEMD_UNIT"`
	Source_realtime_timestamp string `json:"_SOURCE_REALTIME_TIMESTAMP"`
	Machine_id                string `json:"_MACHINE_ID"`
	Hostname                  string `json:"_HOSTNAME"`
	FullMessage               string
}

// Strip date from message-content. Use named subpatterns to override other fields
var messageReplace = map[string]*regexp.Regexp{
	"*":         regexp.MustCompile("^20[0-9][0-9][/\\-][01][0-9][/\\-][0123][0-9] [0-2]?[0-9]:[0-5][0-9]:[0-5][0-9][,0-9]{0-3} "),
	"nginx":     regexp.MustCompile("\\[(?P<Priority>[a-z]+)\\] "),
	"java":      regexp.MustCompile("(?P<Priority>[A-Z]+): "),
	"mysqld":    regexp.MustCompile("^[0-9]+ \\[(?P<Priority>[A-Z][a-z]+)\\] "),
	"searchd":   regexp.MustCompile("^\\[([A-Z][a-z]{2} ){2} [0-9]+ [0-2][0-9]:[0-5][0-9]:[0-5][0-9]\\.[0-9]{3} 20[0-9][0-9]\\] \\[[ 0-9]+\\] "),
	"jenkins":   regexp.MustCompile("^[A-Z][a-z]{2} [01][0-9], 20[0-9][0-9] [0-2]?[0-9]:[0-5][0-9]:[0-5][0-9] [AP]M "),
	"php-fpm":   regexp.MustCompile("^pool [a-z_0-9]+: "),
	"syncthing": regexp.MustCompile("^\\[[0-9A-Z]{5}\\] [0-2][0-9]:[0-5][0-9]:[0-5][0-9] (?P<Priority>INFO): "),
}

var priorities = map[string]int32{
	"emergency": 0,
	"emerg":     0,
	"alert":     1,
	"critical":  2,
	"crit":      2,
	"error":     3,
	"err":       3,
	"warning":   4,
	"warn":      4,
	"notice":    5,
	"info":      6,
	"debug":     7,
}

func (this *SystemdJournalEntry) toGelf() *gelf.Message {
	var extra = map[string]interface{}{
		"Boot_id": this.Boot_id,
		"Pid":     this.Pid,
		"Uid":     this.Uid,
	}

	// php-fpm refuses to fill identifier
	facility := this.Syslog_identifier
	if "" == facility {
		facility = this.Comm
	}

	if this.isJsonMessage() {
		if err := json.Unmarshal([]byte(this.Message), &extra); err == nil {
			if m, ok := extra["Message"]; ok {
				this.Message = m.(string)
				delete(extra, "Message")
			}

			if f, ok := extra["FullMessage"]; ok {
				this.FullMessage = f.(string)
				delete(extra, "FullMessage")
			}
		}
	} else if -1 != strings.Index(this.Message, "\n") {
		this.FullMessage = this.Message
		this.Message = strings.Split(this.Message, "\n")[0]
	}

	return &gelf.Message{
		Version:  "1.1",
		Host:     this.Hostname,
		Short:    this.Message,
		Full:     this.FullMessage,
		TimeUnix: float64(this.Realtime_timestamp) / 1000 / 1000,
		Level:    this.Priority,
		Facility: facility,
		Extra:    extra,
	}
}

func (this *SystemdJournalEntry) process() {
	// Replace generic timestamp
	this.Message = messageReplace["*"].ReplaceAllString(this.Message, "")

	re := messageReplace[ this.Syslog_identifier ]
	if nil == re {
		return
	}

	m := re.FindStringSubmatch(this.Message)
	if m == nil {
		return
	}

	// Store subpatterns in fields
	for idx, key := range re.SubexpNames() {
		if "Priority" == key {
			this.Priority = priorities[strings.ToLower(m[idx])]
		}
	}

	this.Message = re.ReplaceAllString(this.Message, "")
}

func (this *SystemdJournalEntry) sameSource(message *SystemdJournalEntry) bool {
	if this.Syslog_identifier != message.Syslog_identifier {
		return false
	}

	if this.Priority != message.Priority {
		return false
	}

	if this.Realtime_timestamp-message.Realtime_timestamp > SAMESOURCE_TIME_DIFFERENCE {
		return false
	}

	return true
}

func (this *SystemdJournalEntry) send() {
	message := this.toGelf()

	retryAfter := SEND_ERROR_RETRY_WAIT
	for err := writer.WriteMessage(message); err != nil; err = writer.WriteMessage(message) {
		time.Sleep(retryAfter)

		if retryAfter < 1 * time.Minute {
			// Slowly increase, will take 3.4 minutes (50 + 100 ... + 102400 ms) to reach max of 1.7 minutes (102400 ms)
			retryAfter *= 2
		} else {
			fmt.Fprintf(os.Stderr, "Paused message processing due to %s\n", err)
		}
	}

	if retryAfter > SEND_ERROR_RETRY_WAIT {
		messageTime := time.Unix(int64(message.TimeUnix), 0)
		fmt.Fprintf(os.Stderr, "Resuming message processing after %d seconds\n", time.Now().Sub(messageTime)/time.Second)
	}
}

func (this *SystemdJournalEntry) isJsonMessage() bool {
	return len(this.Message) > 64 && this.Message[0] == '{' && this.Message[1] == '"'
}

func (this *SystemdJournalEntry) extendWith(message *SystemdJournalEntry) {
	if this.FullMessage == "" {
		this.FullMessage = this.Message
	}

	this.FullMessage += "\n" + message.Message
}

var (
	pendingEntry *SystemdJournalEntry
	writer       *gelf.Writer
)

const (
	WRITE_INTERVAL             = 50 * time.Millisecond
	SAMESOURCE_TIME_DIFFERENCE = 100 * 1000
	SEND_ERROR_RETRY_WAIT      = 50 * time.Millisecond
)

func main() {
	if len(os.Args) < 3 {
		fmt.Fprintln(os.Stderr, "Pass server:12201 as first argument and append journalctl parameters to use")
		os.Exit(1)
	}

	if w, err := gelf.NewWriter(os.Args[1]); err != nil {
		fmt.Fprintf(os.Stderr, "While connecting to Graylog server: %s\n", err)
		os.Exit(1)
	} else {
		writer = w
	}

	journalArgs := []string{"--all", "--output=json"}
	journalArgs = append(journalArgs, os.Args[2:]...)
	cmd := exec.Command("journalctl", journalArgs...)

	stderr, _ := cmd.StderrPipe()
	go io.Copy(os.Stderr, stderr)
	stdout, _ := cmd.StdoutPipe()
	s := bufio.NewScanner(stdout)

	go writePendingEntry()

	cmd.Start()

	for s.Scan() {
		line := s.Text()

		var entry = &SystemdJournalEntry{}
		if err := json.Unmarshal([]byte(line), &entry); err != nil {
			//fmt.Fprintf(os.Stderr, "Could not parse line, skipping: %s\n", line)
			continue
		}

		entry.process()

		if pendingEntry == nil {
			pendingEntry = entry
		} else if !pendingEntry.sameSource(entry) || pendingEntry.isJsonMessage() {
			pendingEntry.send()
			pendingEntry = entry
		} else {
			pendingEntry.extendWith(entry)

			// Keeps writePendingEntry waiting longer for us to append even more
			pendingEntry.Realtime_timestamp = entry.Realtime_timestamp
		}

		// Prevent saturation and throttling
		time.Sleep(1 * time.Millisecond)
	}

	if err := s.Err(); err != nil {
		fmt.Fprintf(os.Stderr, "Error from Scanner: %s\n", err)
		cmd.Process.Kill()
	}

	cmd.Wait()
	pendingEntry.send()
}

func writePendingEntry() {
	var entry *SystemdJournalEntry

	for {
		time.Sleep(WRITE_INTERVAL)

		if pendingEntry != nil && (time.Now().UnixNano()/1000-pendingEntry.Realtime_timestamp) > SAMESOURCE_TIME_DIFFERENCE {
			entry = pendingEntry
			pendingEntry = nil

			entry.send()
		}
	}
}
<commit_msg>revert 388c6849e, complete this feature first<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"github.com/DECK36/go-gelf/gelf"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"time"
)

/*
	http://www.freedesktop.org/software/systemd/man/systemd.journal-fields.html
	https://github.com/Graylog2/graylog2-docs/wiki/GELF
*/
type SystemdJournalEntry struct {
	Cursor                    string `json:"__CURSOR"`
	Realtime_timestamp        int64  `json:"__REALTIME_TIMESTAMP,string"`
	Monotonic_timestamp       string `json:"__MONOTONIC_TIMESTAMP"`
	Boot_id                   string `json:"_BOOT_ID"`
	Transport                 string `json:"_TRANSPORT"`
	Priority                  int32  `json:"PRIORITY,string"`
	Syslog_facility           string `json:"SYSLOG_FACILITY"`
	Syslog_identifier         string `json:"SYSLOG_IDENTIFIER"`
	Message                   string `json:"MESSAGE"`
	Pid                       string `json:"_PID"`
	Uid                       string `json:"_UID"`
	Gid                       string `json:"_GID"`
	Comm                      string `json:"_COMM"`
	Exe                       string `json:"_EXE"`
	Cmdline                   string `json:"_CMDLINE"`
	Systemd_cgroup            string `json:"_SYSTEMD_CGROUP"`
	Systemd_session           string `json:"_SYSTEMD_SESSION"`
	Systemd_owner_uid         string `json:"_SYSTEMD_OWNER_UID"`
	Systemd_unit              string `json:"_SYSTEMD_UNIT"`
	Source_realtime_timestamp string `json:"_SOURCE_REALTIME_TIMESTAMP"`
	Machine_id                string `json:"_MACHINE_ID"`
	Hostname                  string `json:"_HOSTNAME"`
	FullMessage               string
}

// Strip date from message-content. Use named subpatterns to override other fields
var messageReplace = map[string]*regexp.Regexp{
	"*":         regexp.MustCompile("^20[0-9][0-9][/\\-][01][0-9][/\\-][0123][0-9] [0-2]?[0-9]:[0-5][0-9]:[0-5][0-9][,0-9]{0-3} "),
	"nginx":     regexp.MustCompile("\\[(?P<Priority>[a-z]+)\\] "),
	"java":      regexp.MustCompile("(?P<Priority>[A-Z]+): "),
	"mysqld":    regexp.MustCompile("^[0-9]+ \\[(?P<Priority>[A-Z][a-z]+)\\] "),
	"searchd":   regexp.MustCompile("^\\[([A-Z][a-z]{2} ){2} [0-9]+ [0-2][0-9]:[0-5][0-9]:[0-5][0-9]\\.[0-9]{3} 20[0-9][0-9]\\] \\[[ 0-9]+\\] "),
	"jenkins":   regexp.MustCompile("^[A-Z][a-z]{2} [01][0-9], 20[0-9][0-9] [0-2]?[0-9]:[0-5][0-9]:[0-5][0-9] [AP]M "),
	"php-fpm":   regexp.MustCompile("^pool [a-z_0-9]+: "),
	"syncthing": regexp.MustCompile("^\\[[0-9A-Z]{5}\\] [0-2][0-9]:[0-5][0-9]:[0-5][0-9] (?P<Priority>INFO): "),
}

var priorities = map[string]int32{
	"emergency": 0,
	"emerg":     0,
	"alert":     1,
	"critical":  2,
	"crit":      2,
	"error":     3,
	"err":       3,
	"warning":   4,
	"warn":      4,
	"notice":    5,
	"info":      6,
	"debug":     7,
}

func (this *SystemdJournalEntry) toGelf() *gelf.Message {
	var extra = map[string]interface{}{
		"Boot_id": this.Boot_id,
		"Pid":     this.Pid,
		"Uid":     this.Uid,
	}

	// php-fpm refuses to fill identifier
	facility := this.Syslog_identifier
	if "" == facility {
		facility = this.Comm
	}

	if this.isJsonMessage() {
		if err := json.Unmarshal([]byte(this.Message), &extra); err == nil {
			if m, ok := extra["Message"]; ok {
				this.Message = m.(string)
				delete(extra, "Message")
			}

			if f, ok := extra["FullMessage"]; ok {
				this.FullMessage = f.(string)
				delete(extra, "FullMessage")
			}
		}
	} else if -1 != strings.Index(this.Message, "\n") {
		this.FullMessage = this.Message
		this.Message = strings.Split(this.Message, "\n")[0]
	}

	return &gelf.Message{
		Version:  "1.1",
		Host:     this.Hostname,
		Short:    this.Message,
		Full:     this.FullMessage,
		TimeUnix: float64(this.Realtime_timestamp) / 1000 / 1000,
		Level:    this.Priority,
		Facility: facility,
		Extra:    extra,
	}
}

func (this *SystemdJournalEntry) process() {
	// Replace generic timestamp
	this.Message = messageReplace["*"].ReplaceAllString(this.Message, "")

	re := messageReplace[ this.Syslog_identifier ]
	if nil == re {
		return
	}

	m := re.FindStringSubmatch(this.Message)
	if m == nil {
		return
	}

	// Store subpatterns in fields
	for idx, key := range re.SubexpNames() {
		if "Priority" == key {
			this.Priority = priorities[strings.ToLower(m[idx])]
		}
	}

	this.Message = re.ReplaceAllString(this.Message, "")
}

func (this *SystemdJournalEntry) sameSource(message *SystemdJournalEntry) bool {
	if this.Syslog_identifier != message.Syslog_identifier {
		return false
	}

	if this.Priority != message.Priority {
		return false
	}

	if this.Realtime_timestamp-message.Realtime_timestamp > SAMESOURCE_TIME_DIFFERENCE {
		return false
	}

	return true
}

func (this *SystemdJournalEntry) send() {
	message := this.toGelf()

	if err := writer.WriteMessage(message); err != nil {
		fmt.Fprintln(os.Stderr, err)
	}
}

func (this *SystemdJournalEntry) isJsonMessage() bool {
	return len(this.Message) > 64 && this.Message[0] == '{' && this.Message[1] == '"'
}

func (this *SystemdJournalEntry) extendWith(message *SystemdJournalEntry) {
	if this.FullMessage == "" {
		this.FullMessage = this.Message
	}

	this.FullMessage += "\n" + message.Message
}

var (
	pendingEntry *SystemdJournalEntry
	writer       *gelf.Writer
)

const (
	WRITE_INTERVAL             = 50 * time.Millisecond
	SAMESOURCE_TIME_DIFFERENCE = 100 * 1000
)

func main() {
	if len(os.Args) < 3 {
		fmt.Fprintln(os.Stderr, "Pass server:12201 as first argument and append journalctl parameters to use")
		os.Exit(1)
	}

	if w, err := gelf.NewWriter(os.Args[1]); err != nil {
		fmt.Fprintf(os.Stderr, "While connecting to Graylog server: %s\n", err)
		os.Exit(1)
	} else {
		writer = w
	}

	journalArgs := []string{"--all", "--output=json"}
	journalArgs = append(journalArgs, os.Args[2:]...)
	cmd := exec.Command("journalctl", journalArgs...)

	stderr, _ := cmd.StderrPipe()
	go io.Copy(os.Stderr, stderr)
	stdout, _ := cmd.StdoutPipe()
	s := bufio.NewScanner(stdout)

	go writePendingEntry()

	cmd.Start()

	for s.Scan() {
		line := s.Text()

		var entry = &SystemdJournalEntry{}
		if err := json.Unmarshal([]byte(line), &entry); err != nil {
			//fmt.Fprintf(os.Stderr, "Could not parse line, skipping: %s\n", line)
			continue
		}

		entry.process()

		if pendingEntry == nil {
			pendingEntry = entry
		} else if !pendingEntry.sameSource(entry) || pendingEntry.isJsonMessage() {
			pendingEntry.send()
			pendingEntry = entry
		} else {
			pendingEntry.extendWith(entry)

			// Keeps writePendingEntry waiting longer for us to append even more
			pendingEntry.Realtime_timestamp = entry.Realtime_timestamp
		}

		// Prevent saturation and throttling
		time.Sleep(1 * time.Millisecond)
	}

	if err := s.Err(); err != nil {
		fmt.Fprintf(os.Stderr, "Error from Scanner: %s\n", err)
		cmd.Process.Kill()
	}

	cmd.Wait()
	pendingEntry.send()
}

func writePendingEntry() {
	var entry *SystemdJournalEntry

	for {
		time.Sleep(WRITE_INTERVAL)

		if pendingEntry != nil && (time.Now().UnixNano()/1000-pendingEntry.Realtime_timestamp) > SAMESOURCE_TIME_DIFFERENCE {
			entry = pendingEntry
			pendingEntry = nil

			entry.send()
		}
	}
}
<|endoftext|>