<commit_before>package db

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"
)

func (db *SQLDB) InsertVolume(data Volume) error {
	tx, err := db.conn.Begin()
	if err != nil {
		return err
	}

	var resourceVersion []byte

	resourceVersion, err = json.Marshal(data.ResourceVersion)
	if err != nil {
		return err
	}

	interval := fmt.Sprintf("%d second", int(data.TTL.Seconds()))

	defer tx.Rollback()
	if data.TTL == 0 {
		_, err = tx.Exec(`
		INSERT INTO volumes(
			worker_name,
			expires_at,
			ttl,
			handle,
			resource_version,
			resource_hash
		) VALUES (
			$1,
		  NULL,
			$2,
			$3,
			$4,
			$5
		)
	`, data.WorkerName,
			data.TTL,
			data.Handle,
			resourceVersion,
			data.ResourceHash,
		)
	} else {
		_, err = tx.Exec(`
		INSERT INTO volumes(
			worker_name,
			expires_at,
			ttl,
			handle,
			resource_version,
			resource_hash
		) VALUES (
			$1,
			NOW() + $2::INTERVAL,
			$3,
			$4,
			$5,
			$6
		)
	`, data.WorkerName,
			interval,
			data.TTL,
			data.Handle,
			resourceVersion,
			data.ResourceHash,
		)
	}
	if err != nil {
		if strings.Contains(err.Error(), `duplicate key value violates unique constraint "volumes_worker_name_handle_key"`) {
			return nil
		}

		return err
	}

	return tx.Commit()
}

func (db *SQLDB) InsertCOWVolume(originalVolumeHandle string, cowVolumeHandle string, ttl time.Duration) error {
	originalVolume, err := db.getVolume(originalVolumeHandle)
	if err != nil {
		return err
	}

	tx, err := db.conn.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	interval := fmt.Sprintf("%d second", int(ttl.Seconds()))

	resourceVersion, err := json.Marshal(originalVolume.ResourceVersion)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		INSERT INTO volumes (
			worker_name,
			expires_at,
			ttl,
			handle,
			resource_version,
			resource_hash,
			original_volume_handle
		) VALUES (
			$1,
			NOW() + $2::INTERVAL,
			$3,
			$4,
			$5,
			$6,
			$7
		)
	`, originalVolume.WorkerName,
		interval,
		ttl,
		cowVolumeHandle,
		resourceVersion,
		originalVolume.ResourceHash,
		originalVolumeHandle,
	)

	return tx.Commit()
}

func (db *SQLDB) InsertOutputVolume(outputVolume Volume) error {
	tx, err := db.conn.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	interval := fmt.Sprintf("%d second", int(outputVolume.TTL.Seconds()))

	_, err = tx.Exec(`
		INSERT INTO volumes (
			worker_name,
			expires_at,
			ttl,
			handle
		) VALUES (
			$1,
			NOW() + $2::INTERVAL,
			$3,
			$4
		)
	`,
		outputVolume.WorkerName,
		interval,
		outputVolume.TTL,
		outputVolume.Handle,
	)

	return tx.Commit()
}

func (db *SQLDB) ReapVolume(handle string) error {
	_, err := db.conn.Exec(`
		DELETE FROM volumes
		WHERE handle = $1
	`, handle)
	return err
}

func (db *SQLDB) GetVolumes() ([]SavedVolume, error) {
	err := db.expireVolumes()
	if err != nil {
		return nil, err
	}

	rows, err := db.conn.Query(`
		SELECT
			worker_name,
			ttl,
			EXTRACT(epoch FROM expires_at - NOW()),
			handle,
			resource_version,
			resource_hash,
			id,
			original_volume_handle
		FROM volumes
	`)
	if err != nil {
		return nil, err
	}

	volumes, err := scanVolumes(rows)
	return volumes, err
}

func (db *SQLDB) GetVolumesForOneOffBuildImageResources() ([]SavedVolume, error) {
	err := db.expireVolumes()
	if err != nil {
		return nil, err
	}

	rows, err := db.conn.Query(`
		SELECT DISTINCT
			v.worker_name,
			v.ttl,
			EXTRACT(epoch FROM v.expires_at - NOW()),
			v.handle,
			v.resource_version,
			v.resource_hash,
			v.id,
			v.original_volume_handle
		FROM volumes v
			INNER JOIN image_resource_versions i
				ON i.version = v.resource_version
				AND i.resource_hash = v.resource_hash
			INNER JOIN builds b
				ON b.id = i.build_id
		WHERE b.job_id IS NULL
	`)
	if err != nil {
		return nil, err
	}

	volumes, err := scanVolumes(rows)
	return volumes, err
}

func (db *SQLDB) SetVolumeTTL(handle string, ttl time.Duration) error {
	if ttl == 0 {
		_, err := db.conn.Exec(`
			UPDATE volumes
			SET expires_at = null, ttl = 0
			WHERE handle = $1
		`, handle)

		return err
	}

	interval := fmt.Sprintf("%d second", int(ttl.Seconds()))

	_, err := db.conn.Exec(`
		UPDATE volumes
		SET expires_at = NOW() + $1::INTERVAL,
		ttl = $2
		WHERE handle = $3
	`, interval, ttl, handle)

	return err
}

func (db *SQLDB) GetVolumeTTL(handle string) (time.Duration, error) {
	var ttl time.Duration

	err := db.conn.QueryRow(`
		SELECT ttl
		FROM volumes
		WHERE handle = $1
	`, handle).Scan(&ttl)

	return ttl, err
}

func (db *SQLDB) getVolume(originalVolumeHandle string) (SavedVolume, error) {
	err := db.expireVolumes()
	if err != nil {
		return SavedVolume{}, err
	}

	rows, err := db.conn.Query(`
		SELECT
			worker_name,
			ttl,
			EXTRACT(epoch FROM expires_at - NOW()),
			handle,
			resource_version,
			resource_hash,
			id,
			original_volume_handle
		FROM volumes
		WHERE handle = $1
	`, originalVolumeHandle)
	if err != nil {
		return SavedVolume{}, err
	}

	volumes, err := scanVolumes(rows)
	if err != nil {
		return SavedVolume{}, err
	}

	switch len(volumes) {
	case 0:
		return SavedVolume{}, errors.New(fmt.Sprintf("unable to find volume handle %s", originalVolumeHandle))
	case 1:
		return volumes[0], nil
	default:
		return SavedVolume{}, errors.New(fmt.Sprintf("multiple volumes found for handle %s", originalVolumeHandle))
	}
}

func (db *SQLDB) expireVolumes() error {
	_, err := db.conn.Exec(`
		DELETE FROM volumes
		WHERE expires_at IS NOT NULL
		AND expires_at < NOW()
	`)
	return err
}

func scanVolumes(rows *sql.Rows) ([]SavedVolume, error) {
	defer rows.Close()

	volumes := []SavedVolume{}

	for rows.Next() {
		var volume SavedVolume
		var ttlSeconds *float64
		var versionJSON sql.NullString
		var resourceHash sql.NullString
		var originalVolumeHandle sql.NullString

		err := rows.Scan(&volume.WorkerName, &volume.TTL, &ttlSeconds, &volume.Handle, &versionJSON, &resourceHash, &volume.ID, &originalVolumeHandle)
		if err != nil {
			return nil, err
		}

		if ttlSeconds != nil {
			volume.ExpiresIn = time.Duration(*ttlSeconds) * time.Second
		}

		if versionJSON.Valid {
			err = json.Unmarshal([]byte(versionJSON.String), &volume.ResourceVersion)
			if err != nil {
				return nil, err
			}
		}

		if resourceHash.Valid {
			volume.ResourceHash = resourceHash.String
		}

		if originalVolumeHandle.Valid {
			volume.OriginalVolumeHandle = originalVolumeHandle.String
		}

		volumes = append(volumes, volume)
	}

	return volumes, nil
}
<commit_msg>Revert "clearer error messages when volume is not found"<commit_after>package db

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"
)

func (db *SQLDB) InsertVolume(data Volume) error {
	tx, err := db.conn.Begin()
	if err != nil {
		return err
	}

	var resourceVersion []byte

	resourceVersion, err = json.Marshal(data.ResourceVersion)
	if err != nil {
		return err
	}

	interval := fmt.Sprintf("%d second", int(data.TTL.Seconds()))

	defer tx.Rollback()
	if data.TTL == 0 {
		_, err = tx.Exec(`
		INSERT INTO volumes(
			worker_name,
			expires_at,
			ttl,
			handle,
			resource_version,
			resource_hash
		) VALUES (
			$1,
		  NULL,
			$2,
			$3,
			$4,
			$5
		)
	`, data.WorkerName,
			data.TTL,
			data.Handle,
			resourceVersion,
			data.ResourceHash,
		)
	} else {
		_, err = tx.Exec(`
		INSERT INTO volumes(
			worker_name,
			expires_at,
			ttl,
			handle,
			resource_version,
			resource_hash
		) VALUES (
			$1,
			NOW() + $2::INTERVAL,
			$3,
			$4,
			$5,
			$6
		)
	`, data.WorkerName,
			interval,
			data.TTL,
			data.Handle,
			resourceVersion,
			data.ResourceHash,
		)
	}
	if err != nil {
		if strings.Contains(err.Error(), `duplicate key value violates unique constraint "volumes_worker_name_handle_key"`) {
			return nil
		}

		return err
	}

	return tx.Commit()
}

func (db *SQLDB) InsertCOWVolume(originalVolumeHandle string, cowVolumeHandle string, ttl time.Duration) error {
	originalVolume, err := db.getVolume(originalVolumeHandle)
	if err != nil {
		return err
	}

	tx, err := db.conn.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	interval := fmt.Sprintf("%d second", int(ttl.Seconds()))

	resourceVersion, err := json.Marshal(originalVolume.ResourceVersion)
	if err != nil {
		return err
	}

	_, err = tx.Exec(`
		INSERT INTO volumes (
			worker_name,
			expires_at,
			ttl,
			handle,
			resource_version,
			resource_hash,
			original_volume_handle
		) VALUES (
			$1,
			NOW() + $2::INTERVAL,
			$3,
			$4,
			$5,
			$6,
			$7
		)
	`, originalVolume.WorkerName,
		interval,
		ttl,
		cowVolumeHandle,
		resourceVersion,
		originalVolume.ResourceHash,
		originalVolumeHandle,
	)

	return tx.Commit()
}

func (db *SQLDB) InsertOutputVolume(outputVolume Volume) error {
	tx, err := db.conn.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	interval := fmt.Sprintf("%d second", int(outputVolume.TTL.Seconds()))

	_, err = tx.Exec(`
		INSERT INTO volumes (
			worker_name,
			expires_at,
			ttl,
			handle
		) VALUES (
			$1,
			NOW() + $2::INTERVAL,
			$3,
			$4
		)
	`,
		outputVolume.WorkerName,
		interval,
		outputVolume.TTL,
		outputVolume.Handle,
	)

	return tx.Commit()
}

func (db *SQLDB) ReapVolume(handle string) error {
	_, err := db.conn.Exec(`
		DELETE FROM volumes
		WHERE handle = $1
	`, handle)
	return err
}

func (db *SQLDB) GetVolumes() ([]SavedVolume, error) {
	err := db.expireVolumes()
	if err != nil {
		return nil, err
	}

	rows, err := db.conn.Query(`
		SELECT
			worker_name,
			ttl,
			EXTRACT(epoch FROM expires_at - NOW()),
			handle,
			resource_version,
			resource_hash,
			id,
			original_volume_handle
		FROM volumes
	`)
	if err != nil {
		return nil, err
	}

	volumes, err := scanVolumes(rows)
	return volumes, err
}

func (db *SQLDB) GetVolumesForOneOffBuildImageResources() ([]SavedVolume, error) {
	err := db.expireVolumes()
	if err != nil {
		return nil, err
	}

	rows, err := db.conn.Query(`
		SELECT DISTINCT
			v.worker_name,
			v.ttl,
			EXTRACT(epoch FROM v.expires_at - NOW()),
			v.handle,
			v.resource_version,
			v.resource_hash,
			v.id,
			v.original_volume_handle
		FROM volumes v
			INNER JOIN image_resource_versions i
				ON i.version = v.resource_version
				AND i.resource_hash = v.resource_hash
			INNER JOIN builds b
				ON b.id = i.build_id
		WHERE b.job_id IS NULL
	`)
	if err != nil {
		return nil, err
	}

	volumes, err := scanVolumes(rows)
	return volumes, err
}

func (db *SQLDB) SetVolumeTTL(handle string, ttl time.Duration) error {
	if ttl == 0 {
		_, err := db.conn.Exec(`
			UPDATE volumes
			SET expires_at = null, ttl = 0
			WHERE handle = $1
		`, handle)

		return err
	}

	interval := fmt.Sprintf("%d second", int(ttl.Seconds()))

	_, err := db.conn.Exec(`
		UPDATE volumes
		SET expires_at = NOW() + $1::INTERVAL,
		ttl = $2
		WHERE handle = $3
	`, interval, ttl, handle)

	return err
}

func (db *SQLDB) GetVolumeTTL(handle string) (time.Duration, error) {
	var ttl time.Duration

	err := db.conn.QueryRow(`
		SELECT ttl
		FROM volumes
		WHERE handle = $1
	`, handle).Scan(&ttl)

	return ttl, err
}

func (db *SQLDB) getVolume(originalVolumeHandle string) (SavedVolume, error) {
	err := db.expireVolumes()
	if err != nil {
		return SavedVolume{}, err
	}

	rows, err := db.conn.Query(`
		SELECT
			worker_name,
			ttl,
			EXTRACT(epoch FROM expires_at - NOW()),
			handle,
			resource_version,
			resource_hash,
			id,
			original_volume_handle
		FROM volumes
		WHERE handle = $1
	`, originalVolumeHandle)
	if err != nil {
		return SavedVolume{}, err
	}

	volumes, err := scanVolumes(rows)
	if err != nil {
		return SavedVolume{}, err
	}

	if len(volumes) != 1 {
		return SavedVolume{}, errors.New("0 or 2+ volumes found for same handle")
	}

	return volumes[0], nil
}

func (db *SQLDB) expireVolumes() error {
	_, err := db.conn.Exec(`
		DELETE FROM volumes
		WHERE expires_at IS NOT NULL
		AND expires_at < NOW()
	`)
	return err
}

func scanVolumes(rows *sql.Rows) ([]SavedVolume, error) {
	defer rows.Close()

	volumes := []SavedVolume{}

	for rows.Next() {
		var volume SavedVolume
		var ttlSeconds *float64
		var versionJSON sql.NullString
		var resourceHash sql.NullString
		var originalVolumeHandle sql.NullString

		err := rows.Scan(&volume.WorkerName, &volume.TTL, &ttlSeconds, &volume.Handle, &versionJSON, &resourceHash, &volume.ID, &originalVolumeHandle)
		if err != nil {
			return nil, err
		}

		if ttlSeconds != nil {
			volume.ExpiresIn = time.Duration(*ttlSeconds) * time.Second
		}

		if versionJSON.Valid {
			err = json.Unmarshal([]byte(versionJSON.String), &volume.ResourceVersion)
			if err != nil {
				return nil, err
			}
		}

		if resourceHash.Valid {
			volume.ResourceHash = resourceHash.String
		}

		if originalVolumeHandle.Valid {
			volume.OriginalVolumeHandle = originalVolumeHandle.String
		}

		volumes = append(volumes, volume)
	}

	return volumes, nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package networker

import (
	"fmt"

	"github.com/juju/juju/network"
	"github.com/juju/juju/state/api/base"
	"github.com/juju/juju/state/api/params"
	"github.com/juju/juju/state/api/watcher"
)

const networkerFacade = "Networker"

// State provides access to an networker worker's view of the state.
type State struct {
	facade base.FacadeCaller
}

// NewState creates a new client-side Machiner facade.
func NewState(caller base.APICaller) *State {
	return &State{base.NewFacadeCaller(caller, networkerFacade)}
}

// MachineNetworkInfo returns information about networks to setup only for a single machine.
func (st *State) MachineNetworkInfo(machineTag string) ([]network.Info, error) {
	args := params.Entities{
		Entities: []params.Entity{{Tag: machineTag}},
	}
	var results params.MachineNetworkInfoResults
	err := st.facade.FacadeCall("MachineNetworkInfo", args, &results)
	if err != nil {
		// TODO: Not directly tested
		return nil, err
	}
	if len(results.Results) != 1 {
		// TODO: Not directly tested
		err = fmt.Errorf("expected one result, got %d", len(results.Results))
		return nil, err
	}
	result := results.Results[0]
	if result.Error != nil {
		return nil, result.Error
	}
	return results.Results[0].Info, nil
}

// WatchInterfaces returns a NotifyWatcher that notifies of changes to network
// interfaces on the machine.
func (st *State) WatchInterfaces(machineTag string) (watcher.NotifyWatcher, error) {
	args := params.Entities{
		Entities: []params.Entity{{Tag: machineTag}},
	}
	var results params.NotifyWatchResults
	err := st.call("WatchInterfaces", args, &results)
	if err != nil {
		// TODO: Not directly tested
		return nil, err
	}
	if len(results.Results) != 1 {
		// TODO: Not directly tested
		err = fmt.Errorf("expected one result, got %d", len(results.Results))
		return nil, err
	}
	result := results.Results[0]
	if result.Error != nil {
		return nil, result.Error
	}
	w := watcher.NewNotifyWatcher(st.caller, result)
	return w, nil
}
<commit_msg>fix networker now that Vlad has updated the code.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package networker

import (
	"fmt"

	"github.com/juju/juju/network"
	"github.com/juju/juju/state/api/base"
	"github.com/juju/juju/state/api/params"
	"github.com/juju/juju/state/api/watcher"
)

const networkerFacade = "Networker"

// State provides access to an networker worker's view of the state.
type State struct {
	facade base.FacadeCaller
}

// NewState creates a new client-side Machiner facade.
func NewState(caller base.APICaller) *State {
	return &State{base.NewFacadeCaller(caller, networkerFacade)}
}

// MachineNetworkInfo returns information about networks to setup only for a single machine.
func (st *State) MachineNetworkInfo(machineTag string) ([]network.Info, error) {
	args := params.Entities{
		Entities: []params.Entity{{Tag: machineTag}},
	}
	var results params.MachineNetworkInfoResults
	err := st.facade.FacadeCall("MachineNetworkInfo", args, &results)
	if err != nil {
		// TODO: Not directly tested
		return nil, err
	}
	if len(results.Results) != 1 {
		// TODO: Not directly tested
		err = fmt.Errorf("expected one result, got %d", len(results.Results))
		return nil, err
	}
	result := results.Results[0]
	if result.Error != nil {
		return nil, result.Error
	}
	return results.Results[0].Info, nil
}

// WatchInterfaces returns a NotifyWatcher that notifies of changes to network
// interfaces on the machine.
func (st *State) WatchInterfaces(machineTag string) (watcher.NotifyWatcher, error) {
	args := params.Entities{
		Entities: []params.Entity{{Tag: machineTag}},
	}
	var results params.NotifyWatchResults
	err := st.facade.FacadeCall("WatchInterfaces", args, &results)
	if err != nil {
		// TODO: Not directly tested
		return nil, err
	}
	if len(results.Results) != 1 {
		// TODO: Not directly tested
		err = fmt.Errorf("expected one result, got %d", len(results.Results))
		return nil, err
	}
	result := results.Results[0]
	if result.Error != nil {
		return nil, result.Error
	}
	w := watcher.NewNotifyWatcher(st.facade.RawAPICaller(), result)
	return w, nil
}
<|endoftext|>


<commit_before>package gannoy

import (
	"fmt"
	"math"
	"runtime"
	"sort"
	"sync"

	"github.com/gansidui/priority_queue"
)

type GannoyIndex struct {
	meta      meta
	tree      int
	dim       int
	distance  Distance
	random    Random
	nodes     Nodes
	K         int
	numWorker int
	buildChan chan buildArgs
}

func NewGannoyIndex(metaFile string, distance Distance, random Random) (GannoyIndex, error) {

	meta, err := loadMeta(metaFile)
	if err != nil {
		return GannoyIndex{}, err
	}
	tree := meta.tree
	dim := meta.dim
	K := meta.K

	ann := meta.treePath()

	gannoy := GannoyIndex{
		meta:      meta,
		tree:      tree,
		dim:       dim,
		distance:  distance,
		random:    random,
		K:         K,
		nodes:     newNodes(ann, tree, dim, K),
		numWorker: numWorker(tree),
		buildChan: make(chan buildArgs, 1),
	}
	go gannoy.builder()
	return gannoy, nil
}

func (g GannoyIndex) Tree() {
	for i, root := range g.meta.roots() {
		n, err := g.nodes.getNode(root)
		if err != nil {
			fmt.Printf("%v\n", err)
			break
		}
		g.printTree(i, n, root, 0)
	}
}

func (g GannoyIndex) MetaFile() string {
	return g.meta.file.Name()
}

func (g *GannoyIndex) AddItem(key int, w []float64) error {
	args := buildArgs{action: ADD, key: key, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) RemoveItem(key int) error {
	args := buildArgs{action: DELETE, key: key, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) UpdateItem(key int, w []float64) error {
	args := buildArgs{action: UPDATE, key: key, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) GetNnsByKey(key, n, searchK int) ([]int, error) {
	m, err := g.nodes.getNodeByKey(key)
	if err != nil || !m.isLeaf() {
		return []int{}, fmt.Errorf("Not found")
	}
	return g.getAllNns(m.v, n, searchK)
}

func (g *GannoyIndex) getAllNns(v []float64, n, searchK int) ([]int, error) {
	if searchK == -1 {
		searchK = n * g.tree
	}

	q := priority_queue.New()
	for _, root := range g.meta.roots() {
		q.Push(&Queue{priority: math.Inf(1), value: root})
	}

	nns := []int{}
	for len(nns) < searchK && q.Len() > 0 {
		top := q.Top().(*Queue)
		d := top.priority
		i := top.value

		nd, err := g.nodes.getNode(i)
		if err != nil {
			return []int{}, err
		}
		q.Pop()
		if nd.isLeaf() {
			nns = append(nns, i)
		} else if nd.nDescendants <= g.K {
			dst := nd.children
			nns = append(nns, dst...)
		} else {
			margin := g.distance.margin(nd, v)
			q.Push(&Queue{priority: math.Min(d, +margin), value: nd.children[1]})
			q.Push(&Queue{priority: math.Min(d, -margin), value: nd.children[0]})
		}
	}

	sort.Ints(nns)
	nnsDist := make([]sorter, len(nns))
	last := -1
	dup := 0
	for idx, j := range nns {
		if j == last {
			dup++
			continue
		}
		last = j
		node, err := g.nodes.getNode(j)
		if err != nil {
			return []int{}, err
		}
		nnsDist[idx-dup] = sorter{value: g.distance.distance(v, node.v), id: node.key}
	}
	nnsDist = nnsDist[:len(nns)-dup]

	m := len(nnsDist)
	p := m
	if n < m {
		p = n
	}

	HeapSort(nnsDist, DESC, p)

	result := make([]int, p)
	for i := 0; i < p; i++ {
		result[i] = nnsDist[m-1-i].id
	}

	return result, nil
}

func (g *GannoyIndex) addItem(key int, w []float64) error {
	if len(w) != g.dim {
		return fmt.Errorf("Dimension mismatch. expect %d, but %d.\n", g.dim, len(w))
	}
	n := g.nodes.newNode()
	n.key = key
	n.v = w
	n.parents = make([]int, g.tree)
	err := n.save()
	if err != nil {
		return err
	}
	// fmt.Printf("id %d\n", n.id)

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for index := range buildChan {
			g.build(index, g.meta.roots()[index], n)
			wg.Done()
		}
	}

	for i := 0; i < g.numWorker; i++ {
		go worker(n)
	}

	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)
	g.nodes.maps.add(n.id, key)

	return nil
}

func (g *GannoyIndex) build(index, root int, n Node) {
	if root == -1 {
		// 最初のノード
		n.parents[index] = -1
		n.save()
		g.meta.updateRoot(index, n.id)
		return
	}
	id := g.findBranchByVector(root, n.v)
	found, _ := g.nodes.getNode(id)
	// fmt.Printf("Found %d\n", item)

	org_parent := found.parents[index]
	if found.isBucket() && len(found.children) < g.K {
		// ノードに余裕があれば追加
		// fmt.Printf("pattern bucket\n")
		n.updateParents(index, id)
		found.nDescendants++
		found.children = append(found.children, n.id)
		found.save()
	} else {
		// ノードが上限またはリーフノードであれば新しいノードを追加
		willDelete := false
		var indices []int
		if found.isLeaf() {
			// fmt.Printf("pattern leaf node\n")
			indices = []int{id, n.id}
		} else {
			// fmt.Printf("pattern full backet\n")
			indices = append(found.children, n.id)
			willDelete = true
		}

		m := g.makeTree(index, org_parent, indices)
		if root == 0 || org_parent == 0 || org_parent == -1 {
			// rootノードの入れ替え
			g.meta.updateRoot(index, m)
		} else {
			parent, _ := g.nodes.getNode(org_parent)
			parent.nDescendants++
			children := make([]int, len(parent.children))
			for i, child := range parent.children {
				if child == id {
					// 新しいノードに変更
					children[i] = m
				} else {
					// 既存のノードのまま
					children[i] = child
				}
			}
			parent.children = children
			parent.save()

		}
		if willDelete {
			found.destroy()
			g.nodes.free.push(found.id)
		}
	}
}

func (g *GannoyIndex) removeItem(key int) error {
	n, err := g.nodes.getNodeByKey(key)
	if err != nil {
		return err
	}

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for root := range buildChan {
			g.remove(root, n)
			wg.Done()
		}
	}

	for i := 0; i < g.numWorker; i++ {
		go worker(n)
	}
	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)

	g.nodes.maps.remove(key)
	n.destroy()
	g.nodes.free.push(n.id)

	return nil
}

func (g *GannoyIndex) remove(root int, node Node) {
	if node.isRoot(root) {
		g.meta.updateRoot(root, -1)
		return
	}
	parent, _ := g.nodes.getNode(node.parents[root])
	if parent.isBucket() && len(parent.children) > 2 {
		// fmt.Printf("pattern bucket\n")
		target := -1
		for i, child := range parent.children {
			if child == node.id {
				target = i
			}
		}
		if target == -1 {
			return
		}
		children := append(parent.children[:target], parent.children[(target+1):]...)
		parent.nDescendants--
		parent.children = children
		parent.save()
	} else {
		// fmt.Printf("pattern leaf node\n")
		var other int
		for _, child := range parent.children {
			if child != node.id {
				other = child
			}
		}
		if parent.isRoot(root) {
			g.meta.updateRoot(root, other)
		} else {
			grandParent, _ := g.nodes.getNode(parent.parents[root])
			children := []int{}
			for _, child := range grandParent.children {
				if child == node.parents[root] {
					children = append(children, other)
				} else {
					children = append(children, child)
				}
			}
			grandParent.nDescendants--
			grandParent.children = children
			grandParent.save()
		}

		otherNode, _ := g.nodes.getNode(other)
		otherNode.updateParents(root, parent.parents[root])

		parent.destroy()
		g.nodes.free.push(parent.id)
	}
}

func (g GannoyIndex) findBranchByVector(id int, v []float64) int {
	node, _ := g.nodes.getNode(id)
	if node.isLeaf() || node.isBucket() {
		return id
	}
	side := g.distance.side(node, v, g.random)
	return g.findBranchByVector(node.children[side], v)
}

func (g *GannoyIndex) makeTree(root, parent int, ids []int) int {
	if len(ids) == 1 {
		n, _ := g.nodes.getNode(ids[0])
		if len(n.parents) == 0 {
			n.parents = make([]int, g.tree)
		}
		n.updateParents(root, parent)
		return ids[0]
	}

	if len(ids) <= g.K {
		m := g.nodes.newNode()
		m.parents = make([]int, g.tree)
		m.nDescendants = len(ids)
		m.parents[root] = parent
		m.children = ids
		m.save()
		for _, child := range ids {
			c, _ := g.nodes.getNode(child)
			if len(c.parents) == 0 {
				c.parents = make([]int, g.tree)
			}
			c.updateParents(root, m.id)
		}
		return m.id
	}

	children := make([]Node, len(ids))
	for i, id := range ids {
		children[i], _ = g.nodes.getNode(id)
	}

	childrenIds := [2][]int{[]int{}, []int{}}

	m := g.nodes.newNode()
	m.parents = make([]int, g.tree)
	m.nDescendants = len(ids)
	m.parents[root] = parent

	m = g.distance.createSplit(children, g.random, m)
	for _, id := range ids {
		n, _ := g.nodes.getNode(id)
		side := g.distance.side(m, n.v, g.random)
		childrenIds[side] = append(childrenIds[side], id)
	}

	for len(childrenIds[0]) == 0 || len(childrenIds[1]) == 0 {
		childrenIds[0] = []int{}
		childrenIds[1] = []int{}
		for z := 0; z < g.dim; z++ {
			m.v[z] = 0.0
		}
		for _, id := range ids {
			side := g.random.flip()
			childrenIds[side] = append(childrenIds[side], id)
		}
	}

	var flip int
	if len(childrenIds[0]) > len(childrenIds[1]) {
		flip = 1
	}

	m.save()
	for side := 0; side < 2; side++ {
		m.children[side^flip] = g.makeTree(root, m.id, childrenIds[side^flip])
	}
	m.save()

	return m.id
}

type buildArgs struct {
	action int
	key    int
	w      []float64
	result chan error
}

func (g *GannoyIndex) builder() {
	for args := range g.buildChan {
		switch args.action {
		case ADD:
			args.result <- g.addItem(args.key, args.w)
		case DELETE:
			args.result <- g.removeItem(args.key)
		case UPDATE:
			err := g.removeItem(args.key)
			if err != nil {
				args.result <- err
			} else {
				args.result <- g.addItem(args.key, args.w)
			}
		}
	}
}

func (g GannoyIndex) printTree(root int, node Node, id, tab int) {
	for i := 0; i < tab*2; i++ {
		fmt.Print(" ")
	}
	fmt.Printf("%d [%d] (%d) [nDescendants: %d, v: %v]\n", id, node.key, node.parents[root], node.nDescendants, node.v)
	if !node.isLeaf() {
		for _, child := range node.children {
			n, err := g.nodes.getNode(child)
			if err != nil {
				fmt.Printf("%v\n", err)
				break
			}
			g.printTree(root, n, child, tab+1)
		}
	}
}

func numWorker(tree int) int {
	procs := runtime.GOMAXPROCS(0) // current setting
	if tree < procs {
		return tree
	} else {
		return procs
	}
}
<commit_msg>key existence is now checked when called addItem.<commit_after>package gannoy

import (
	"fmt"
	"math"
	"runtime"
	"sort"
	"sync"

	"github.com/gansidui/priority_queue"
)

type GannoyIndex struct {
	meta      meta
	tree      int
	dim       int
	distance  Distance
	random    Random
	nodes     Nodes
	K         int
	numWorker int
	buildChan chan buildArgs
}

func NewGannoyIndex(metaFile string, distance Distance, random Random) (GannoyIndex, error) {

	meta, err := loadMeta(metaFile)
	if err != nil {
		return GannoyIndex{}, err
	}
	tree := meta.tree
	dim := meta.dim
	K := meta.K

	ann := meta.treePath()

	gannoy := GannoyIndex{
		meta:      meta,
		tree:      tree,
		dim:       dim,
		distance:  distance,
		random:    random,
		K:         K,
		nodes:     newNodes(ann, tree, dim, K),
		numWorker: numWorker(tree),
		buildChan: make(chan buildArgs, 1),
	}
	go gannoy.builder()
	return gannoy, nil
}

func (g GannoyIndex) Tree() {
	for i, root := range g.meta.roots() {
		n, err := g.nodes.getNode(root)
		if err != nil {
			fmt.Printf("%v\n", err)
			break
		}
		g.printTree(i, n, root, 0)
	}
}

func (g GannoyIndex) MetaFile() string {
	return g.meta.file.Name()
}

func (g *GannoyIndex) AddItem(key int, w []float64) error {
	args := buildArgs{action: ADD, key: key, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) RemoveItem(key int) error {
	args := buildArgs{action: DELETE, key: key, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) UpdateItem(key int, w []float64) error {
	args := buildArgs{action: UPDATE, key: key, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) GetNnsByKey(key, n, searchK int) ([]int, error) {
	m, err := g.nodes.getNodeByKey(key)
	if err != nil || !m.isLeaf() {
		return []int{}, fmt.Errorf("Not found")
	}
	return g.getAllNns(m.v, n, searchK)
}

func (g *GannoyIndex) getAllNns(v []float64, n, searchK int) ([]int, error) {
	if searchK == -1 {
		searchK = n * g.tree
	}

	q := priority_queue.New()
	for _, root := range g.meta.roots() {
		q.Push(&Queue{priority: math.Inf(1), value: root})
	}

	nns := []int{}
	for len(nns) < searchK && q.Len() > 0 {
		top := q.Top().(*Queue)
		d := top.priority
		i := top.value

		nd, err := g.nodes.getNode(i)
		if err != nil {
			return []int{}, err
		}
		q.Pop()
		if nd.isLeaf() {
			nns = append(nns, i)
		} else if nd.nDescendants <= g.K {
			dst := nd.children
			nns = append(nns, dst...)
		} else {
			margin := g.distance.margin(nd, v)
			q.Push(&Queue{priority: math.Min(d, +margin), value: nd.children[1]})
			q.Push(&Queue{priority: math.Min(d, -margin), value: nd.children[0]})
		}
	}

	sort.Ints(nns)
	nnsDist := make([]sorter, len(nns))
	last := -1
	dup := 0
	for idx, j := range nns {
		if j == last {
			dup++
			continue
		}
		last = j
		node, err := g.nodes.getNode(j)
		if err != nil {
			return []int{}, err
		}
		nnsDist[idx-dup] = sorter{value: g.distance.distance(v, node.v), id: node.key}
	}
	nnsDist = nnsDist[:len(nns)-dup]

	m := len(nnsDist)
	p := m
	if n < m {
		p = n
	}

	HeapSort(nnsDist, DESC, p)

	result := make([]int, p)
	for i := 0; i < p; i++ {
		result[i] = nnsDist[m-1-i].id
	}

	return result, nil
}

func (g *GannoyIndex) addItem(key int, w []float64) error {
	if len(w) != g.dim {
		return fmt.Errorf("Dimension mismatch. expect %d, but %d.\n", g.dim, len(w))
	}
	if g.nodes.maps.isExist(key) {
		return fmt.Errorf("Key [%d] is already exist.\n", key)
	}
	n := g.nodes.newNode()
	n.key = key
	n.v = w
	n.parents = make([]int, g.tree)
	err := n.save()
	if err != nil {
		return err
	}
	// fmt.Printf("id %d\n", n.id)

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for index := range buildChan {
			g.build(index, g.meta.roots()[index], n)
			wg.Done()
		}
	}

	for i := 0; i < g.numWorker; i++ {
		go worker(n)
	}

	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)
	g.nodes.maps.add(n.id, key)

	return nil
}

func (g *GannoyIndex) build(index, root int, n Node) {
	if root == -1 {
		// 最初のノード
		n.parents[index] = -1
		n.save()
		g.meta.updateRoot(index, n.id)
		return
	}
	id := g.findBranchByVector(root, n.v)
	found, _ := g.nodes.getNode(id)
	// fmt.Printf("Found %d\n", item)

	org_parent := found.parents[index]
	if found.isBucket() && len(found.children) < g.K {
		// ノードに余裕があれば追加
		// fmt.Printf("pattern bucket\n")
		n.updateParents(index, id)
		found.nDescendants++
		found.children = append(found.children, n.id)
		found.save()
	} else {
		// ノードが上限またはリーフノードであれば新しいノードを追加
		willDelete := false
		var indices []int
		if found.isLeaf() {
			// fmt.Printf("pattern leaf node\n")
			indices = []int{id, n.id}
		} else {
			// fmt.Printf("pattern full backet\n")
			indices = append(found.children, n.id)
			willDelete = true
		}

		m := g.makeTree(index, org_parent, indices)
		if root == 0 || org_parent == 0 || org_parent == -1 {
			// rootノードの入れ替え
			g.meta.updateRoot(index, m)
		} else {
			parent, _ := g.nodes.getNode(org_parent)
			parent.nDescendants++
			children := make([]int, len(parent.children))
			for i, child := range parent.children {
				if child == id {
					// 新しいノードに変更
					children[i] = m
				} else {
					// 既存のノードのまま
					children[i] = child
				}
			}
			parent.children = children
			parent.save()

		}
		if willDelete {
			found.destroy()
			g.nodes.free.push(found.id)
		}
	}
}

func (g *GannoyIndex) removeItem(key int) error {
	n, err := g.nodes.getNodeByKey(key)
	if err != nil {
		return err
	}

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for root := range buildChan {
			g.remove(root, n)
			wg.Done()
		}
	}

	for i := 0; i < g.numWorker; i++ {
		go worker(n)
	}
	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)

	g.nodes.maps.remove(key)
	n.destroy()
	g.nodes.free.push(n.id)

	return nil
}

func (g *GannoyIndex) remove(root int, node Node) {
	if node.isRoot(root) {
		g.meta.updateRoot(root, -1)
		return
	}
	parent, _ := g.nodes.getNode(node.parents[root])
	if parent.isBucket() && len(parent.children) > 2 {
		// fmt.Printf("pattern bucket\n")
		target := -1
		for i, child := range parent.children {
			if child == node.id {
				target = i
			}
		}
		if target == -1 {
			return
		}
		children := append(parent.children[:target], parent.children[(target+1):]...)
		parent.nDescendants--
		parent.children = children
		parent.save()
	} else {
		// fmt.Printf("pattern leaf node\n")
		var other int
		for _, child := range parent.children {
			if child != node.id {
				other = child
			}
		}
		if parent.isRoot(root) {
			g.meta.updateRoot(root, other)
		} else {
			grandParent, _ := g.nodes.getNode(parent.parents[root])
			children := []int{}
			for _, child := range grandParent.children {
				if child == node.parents[root] {
					children = append(children, other)
				} else {
					children = append(children, child)
				}
			}
			grandParent.nDescendants--
			grandParent.children = children
			grandParent.save()
		}

		otherNode, _ := g.nodes.getNode(other)
		otherNode.updateParents(root, parent.parents[root])

		parent.destroy()
		g.nodes.free.push(parent.id)
	}
}

func (g GannoyIndex) findBranchByVector(id int, v []float64) int {
	node, _ := g.nodes.getNode(id)
	if node.isLeaf() || node.isBucket() {
		return id
	}
	side := g.distance.side(node, v, g.random)
	return g.findBranchByVector(node.children[side], v)
}

func (g *GannoyIndex) makeTree(root, parent int, ids []int) int {
	if len(ids) == 1 {
		n, _ := g.nodes.getNode(ids[0])
		if len(n.parents) == 0 {
			n.parents = make([]int, g.tree)
		}
		n.updateParents(root, parent)
		return ids[0]
	}

	if len(ids) <= g.K {
		m := g.nodes.newNode()
		m.parents = make([]int, g.tree)
		m.nDescendants = len(ids)
		m.parents[root] = parent
		m.children = ids
		m.save()
		for _, child := range ids {
			c, _ := g.nodes.getNode(child)
			if len(c.parents) == 0 {
				c.parents = make([]int, g.tree)
			}
			c.updateParents(root, m.id)
		}
		return m.id
	}

	children := make([]Node, len(ids))
	for i, id := range ids {
		children[i], _ = g.nodes.getNode(id)
	}

	childrenIds := [2][]int{[]int{}, []int{}}

	m := g.nodes.newNode()
	m.parents = make([]int, g.tree)
	m.nDescendants = len(ids)
	m.parents[root] = parent

	m = g.distance.createSplit(children, g.random, m)
	for _, id := range ids {
		n, _ := g.nodes.getNode(id)
		side := g.distance.side(m, n.v, g.random)
		childrenIds[side] = append(childrenIds[side], id)
	}

	for len(childrenIds[0]) == 0 || len(childrenIds[1]) == 0 {
		childrenIds[0] = []int{}
		childrenIds[1] = []int{}
		for z := 0; z < g.dim; z++ {
			m.v[z] = 0.0
		}
		for _, id := range ids {
			side := g.random.flip()
			childrenIds[side] = append(childrenIds[side], id)
		}
	}

	var flip int
	if len(childrenIds[0]) > len(childrenIds[1]) {
		flip = 1
	}

	m.save()
	for side := 0; side < 2; side++ {
		m.children[side^flip] = g.makeTree(root, m.id, childrenIds[side^flip])
	}
	m.save()

	return m.id
}

type buildArgs struct {
	action int
	key    int
	w      []float64
	result chan error
}

func (g *GannoyIndex) builder() {
	for args := range g.buildChan {
		switch args.action {
		case ADD:
			args.result <- g.addItem(args.key, args.w)
		case DELETE:
			args.result <- g.removeItem(args.key)
		case UPDATE:
			err := g.removeItem(args.key)
			if err != nil {
				args.result <- err
			} else {
				args.result <- g.addItem(args.key, args.w)
			}
		}
	}
}

func (g GannoyIndex) printTree(root int, node Node, id, tab int) {
	for i := 0; i < tab*2; i++ {
		fmt.Print(" ")
	}
	fmt.Printf("%d [%d] (%d) [nDescendants: %d, v: %v]\n", id, node.key, node.parents[root], node.nDescendants, node.v)
	if !node.isLeaf() {
		for _, child := range node.children {
			n, err := g.nodes.getNode(child)
			if err != nil {
				fmt.Printf("%v\n", err)
				break
			}
			g.printTree(root, n, child, tab+1)
		}
	}
}

func numWorker(tree int) int {
	procs := runtime.GOMAXPROCS(0) // current setting
	if tree < procs {
		return tree
	} else {
		return procs
	}
}
<|endoftext|>


<commit_before>package main

import (
	cryptorand "crypto/rand"
	"crypto/sha512"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"runtime"
	"time"
)

const (
	BSIZE    = 65536
	BSIZEMOD = BSIZE % 7
	BSIZE7   = BSIZE - BSIZEMOD
	SRC      = 8
)

var (
	narg    = int(0)
	phrase  = ""
	sources [8]rand.Source
)

// randomSeed produces a int64 seed based on crypto/rand and time.
func randomSeed() int64 {
	var seed int64

	urandom := make([]byte, 8)
	cryptorand.Reader.Read(urandom)

	for key, value := range urandom {
		seed ^= (int64(value) ^ time.Now().UTC().UnixNano()) << (uint(key) * 8)
	}

	return seed
}

func randomBytes(src rand.Source) []byte {
	buf := make([]byte, BSIZE)

	for i := 0; i < BSIZE7; i += 7 {
		r := src.Int63()
		buf[i] = byte(r)
		buf[i+1] = byte(r << 8)
		buf[i+2] = byte(r << 16)
		buf[i+3] = byte(r << 24)
		buf[i+4] = byte(r << 32)
		buf[i+5] = byte(r << 40)
		buf[i+6] = byte(r << 48)
	}

	r := src.Int63()

	for i := BSIZE - BSIZE7; i < BSIZE; i++ {
		buf[i] = byte(r)
		r = r << 8
	}

	return buf
}

func init() {
	flag.StringVar(&phrase, "phrase", "", "the Garble phrase, by default random")
	flag.Parse()

	narg = flag.NArg()

	if narg <= 0 {
		flag.Usage()
		os.Exit(1)
	}

	if phrase == "" {
		phrase = fmt.Sprintf("%016x", uint64(randomSeed()))
	}
}

func garble(index int, f *os.File, c chan []byte) {
	fmt.Println("garble(", index, f, c, ")")
	<-c
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	channels := make([]chan []byte, narg)
	files := make([]*os.File, narg)

	for i, arg := range flag.Args() {
		f, err := os.OpenFile(arg, os.O_RDWR, 0666)
		if err != nil {
			panic(err)
		}
		defer f.Close()
		files[i] = f
		channels[i] = make(chan []byte, 8)
		go garble(i, files[i], channels[i])
	}

	hash := sha512.New()
	seed := []byte(phrase)
	b := make([]byte, hash.Size())
	for i := 0; i < 2; i++ {
		hash.Write(seed)
		hash.Sum(b[:0])
		fmt.Println(b)
	}
}
<commit_msg>stupid randomBytes<commit_after>package main

import (
	cryptorand "crypto/rand"
	"crypto/sha512"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"runtime"
	"time"
)

const (
	BSIZE    = 65536
	BSIZEMOD = BSIZE % 7
	BSIZE7   = BSIZE - BSIZEMOD
	SRC      = 8
)

var (
	narg    = int(0)
	phrase  = ""
	sources [8]rand.Source
)

// randomSeed produces a int64 seed based on crypto/rand and time.
func randomSeed() int64 {
	var seed int64

	urandom := make([]byte, 8)
	cryptorand.Reader.Read(urandom)

	for key, value := range urandom {
		seed ^= (int64(value) ^ time.Now().UTC().UnixNano()) << (uint(key) * 8)
	}

	return seed
}

func randomBytes(src rand.Source) []byte {
	buf := make([]byte, BSIZE)

	for i := 0; i < BSIZE7; i++ {
		r := src.Int63()
		buf[i] = byte(r)
		i++
		buf[i] = byte(r >> 8)
		i++
		buf[i] = byte(r >> 16)
		i++
		buf[i] = byte(r >> 24)
		i++
		buf[i] = byte(r >> 32)
		i++
		buf[i] = byte(r >> 40)
		i++
		buf[i] = byte(r >> 48)
	}

	r := src.Int63()

	for i := BSIZE - BSIZEMOD; i < BSIZE; i++ {
		buf[i] = byte(r)
		r >>= 8
	}

	return buf
}

func init() {
	flag.StringVar(&phrase, "phrase", "", "the Garble phrase, by default random")
	flag.Parse()

	narg = flag.NArg()

	if narg <= 0 {
		flag.Usage()
		os.Exit(1)
	}

	if phrase == "" {
		phrase = fmt.Sprintf("%016x", uint64(randomSeed()))
	}
}

func garble(index int, f *os.File, c chan []byte) {
	fmt.Println("garble(", index, f, c, ")")
	<-c
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	channels := make([]chan []byte, narg)
	files := make([]*os.File, narg)

	for i, arg := range flag.Args() {
		f, err := os.OpenFile(arg, os.O_RDWR, 0666)
		if err != nil {
			panic(err)
		}
		defer f.Close()
		files[i] = f
		channels[i] = make(chan []byte, 8)
		go garble(i, files[i], channels[i])
	}

	hash := sha512.New()
	seed := []byte(phrase)
	b := make([]byte, hash.Size())
	for i := 0; i < 2; i++ {
		hash.Write(seed)
		hash.Sum(b[:0])
		fmt.Println(b)
	}

	src := rand.NewSource(int64(b[0]) + int64(b[1])<<8)

	for i := 0; i < 2; i++ {
		os.Stdout.Write(randomBytes(src))
	}
}
<|endoftext|>


<commit_before>package arn

import "sort"

// Post represents a forum post.
type Post struct {
	ID       string   `json:"id"`
	Text     string   `json:"text"`
	AuthorID string   `json:"authorId"`
	ThreadID string   `json:"threadId"`
	Tags     []string `json:"tags"`
	Likes    []string `json:"likes"`
	Created  string   `json:"created"`
	Edited   string   `json:"edited"`

	author *User
	thread *Thread
}

// Author returns the post author.
func (post *Post) Author() *User {
	if post.author != nil {
		return post.author
	}

	post.author, _ = GetUser(post.AuthorID)
	return post.author
}

// Thread returns the thread this post was posted in.
func (post *Post) Thread() *Thread {
	if post.thread != nil {
		return post.thread
	}

	post.thread, _ = GetThread(post.ThreadID)
	return post.thread
}

// Link returns the relative URL of the post.
func (post *Post) Link() string {
	return "/posts/" + post.ID
}

// ToPostable converts a post into an object that implements the Postable interface.
func (post *Post) ToPostable() Postable {
	return &PostPostable{post}
}

// GetPost ...
func GetPost(id string) (*Post, error) {
	obj, err := DB.Get("Post", id)
	return obj.(*Post), err
}

// GetPosts ...
func GetPosts() ([]*Post, error) {
	var posts []*Post

	channel := make(chan *Post)
	err := DB.Scan("Post", channel)

	for post := range channel {
		posts = append(posts, post)
	}

	return posts, err
}

// SortPostsLatestFirst sorts the slice of posts.
func SortPostsLatestFirst(posts []*Post) {
	sort.Slice(posts, func(i, j int) bool {
		return posts[i].Created > posts[j].Created
	})
}

// SortPostsLatestLast sorts the slice of posts.
func SortPostsLatestLast(posts []*Post) {
	sort.Slice(posts, func(i, j int) bool {
		return posts[i].Created < posts[j].Created
	})
}

// GetPostsByUser ...
func GetPostsByUser(user *User) ([]*Post, error) {
	var posts []*Post

	scan := make(chan *Post)
	err := DB.Scan("Post", scan)

	for post := range scan {
		if post.AuthorID == user.ID {
			posts = append(posts, post)
		}
	}

	return posts, err
}

// FilterPosts filters all forum posts by a custom function.
func FilterPosts(filter func(*Post) bool) ([]*Post, error) {
	var filtered []*Post

	channel := make(chan *Post)
	err := DB.Scan("Post", channel)

	if err != nil {
		return filtered, err
	}

	for post := range channel {
		if filter(post) {
			filtered = append(filtered, post)
		}
	}

	return filtered, nil
}
<commit_msg>Updated Post.go<commit_after>package arn

import "sort"

// Post represents a forum post.
type Post struct {
	ID       string   `json:"id"`
	Text     string   `json:"text"`
	AuthorID string   `json:"authorId"`
	ThreadID string   `json:"threadId"`
	Tags     []string `json:"tags"`
	Likes    []string `json:"likes"`
	Created  string   `json:"created"`
	Edited   string   `json:"edited"`

	author *User
	thread *Thread
}

// Author returns the post author.
func (post *Post) Author() *User {
	if post.author != nil {
		return post.author
	}

	post.author, _ = GetUser(post.AuthorID)
	return post.author
}

// Thread returns the thread this post was posted in.
func (post *Post) Thread() *Thread {
	if post.thread != nil {
		return post.thread
	}

	post.thread, _ = GetThread(post.ThreadID)
	return post.thread
}

// Link returns the relative URL of the post.
func (post *Post) Link() string {
	return "/posts/" + post.ID
}

// ToPostable converts a post into an object that implements the Postable interface.
func (post *Post) ToPostable() Postable {
	return &PostPostable{post}
}

// GetPost ...
func GetPost(id string) (*Post, error) {
	obj, err := DB.Get("Post", id)
	return obj.(*Post), err
}

// GetPosts ...
func GetPosts() ([]*Post, error) {
	var posts []*Post

	channel := make(chan *Post)
	err := DB.Scan("Post", channel)

	for post := range channel {
		posts = append(posts, post)
	}

	return posts, err
}

// SortPostsLatestFirst sorts the slice of posts.
func SortPostsLatestFirst(posts []*Post) {
	sort.Slice(posts, func(i, j int) bool {
		return posts[i].Created > posts[j].Created
	})
}

// SortPostsLatestLast sorts the slice of posts.
func SortPostsLatestLast(posts []*Post) {
	sort.Slice(posts, func(i, j int) bool {
		return posts[i].Created < posts[j].Created
	})
}

// GetPostsByUser ...
func GetPostsByUser(user *User) ([]*Post, error) {
	var posts []*Post

	scan := make(chan *Post)
	err := DB.Scan("Post", scan)

	if err != nil {
		return nil, err
	}

	for post := range scan {
		if post.AuthorID == user.ID {
			posts = append(posts, post)
		}
	}

	return posts, nil
}

// FilterPosts filters all forum posts by a custom function.
func FilterPosts(filter func(*Post) bool) ([]*Post, error) {
	var filtered []*Post

	channel := make(chan *Post)
	err := DB.Scan("Post", channel)

	if err != nil {
		return filtered, err
	}

	for post := range channel {
		if filter(post) {
			filtered = append(filtered, post)
		}
	}

	return filtered, nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rsa"

	cryptorand "crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"errors"
	"github.com/Sirupsen/logrus"
	"github.com/hlandau/acme/acmeapi"
	"github.com/hlandau/acme/acmeapi/acmeutils"
	context "golang.org/x/net/context"
	"net"
	"strings"
	"sync"
	"time"
	"fmt"
	"encoding/pem"
)

const (
	SNI01_EXPIRE_TOKEN time.Duration = time.Minute * 10
)

type acmeStruct struct {
	serverAddress string
	privateKey    *rsa.PrivateKey
	client        *acmeapi.Client

	mutex                *sync.Mutex
	acmeauthDomainsMutex *sync.Mutex
	acmeAuthDomains      map[string]time.Time
}

func (this *acmeStruct) Init() {
	this.client = &acmeapi.Client{
		AccountKey:   this.privateKey,
		DirectoryURL: this.serverAddress,
	}

	this.mutex = &sync.Mutex{}

	this.acmeauthDomainsMutex = &sync.Mutex{}
	this.acmeAuthDomains = make(map[string]time.Time)
	this.CleanupTimer()
}

func (this *acmeStruct) RegisterEnsure(ctx context.Context) (err error) {
	reg := &acmeapi.Registration{}
	for i := 0; i < TRY_COUNT+1; i++ { // +1 count need for request latest agreement uri
		reg.AgreementURI = reg.LatestAgreementURI
		if reg.AgreementURI != "" {
			logrus.Info("Auto agree with terms:", reg.LatestAgreementURI)
		}
		err = this.client.UpsertRegistration(reg, ctx)
		if reg.AgreementURI != "" && err == nil {
			return
		}
	}
	return err
}

func (this *acmeStruct) Cleanup() {
	this.mutex.Lock()
	defer this.mutex.Unlock()

	now := time.Now()
	for token, expire := range this.acmeAuthDomains {
		if expire.Before(now) {
			delete(this.acmeAuthDomains, token)
		}
	}
}

func (this *acmeStruct) CleanupTimer() {
	this.Cleanup()
	time.AfterFunc(SNI01_EXPIRE_TOKEN, this.Cleanup)
}

func (this *acmeStruct) CreateCertificate(domain string) (cert *tls.Certificate, err error) {
	// Check suffix for avoid mutex sync in DeleteAcmeAuthDomain
	if strings.HasSuffix(domain, ".acme.invalid") {
		logrus.Debugf("Detect auth-domain mode for domain '%v'", domain)
		if this.DeleteAcmeAuthDomain(domain) {
			logrus.Debugf("Return self-signed certificate for domain '%v'", domain)
			return this.createCertificateSelfSigned(domain)
		} else {
			logrus.Debugf("Detect auth-domain is not present in list '%v'", domain)
			return nil, errors.New("Now allowed auth-domain")
		}
	}

	// check about we serve the domain
	ips, err := net.LookupIP(domain)
	if err != nil {
		logrus.Warnf("Can't lookup ip for domain '%v': %v", domain, err)
		return nil, errors.New("Can't lookup ip of the domain")
	}
	isLocalIP := false
checkLocalIP:
	for _, ip := range ips {
		for _, localIP := range localIPs {
			if ip.Equal(localIP) {
				isLocalIP = true
				break checkLocalIP
			}
		}
	}
	if !isLocalIP {
		logrus.Warnf("Domain have ip of other server. Domain '%v', Domain ips: %v, Server ips: %v", domain, ips, localIPs)
		return nil, errors.New("Domain have ip of other server.")
	}

	return this.createCertificateAcme(domain)
}

func (this *acmeStruct) createCertificateAcme(domain string) (cert *tls.Certificate, err error) {
	this.mutex.Lock()
	defer this.mutex.Unlock()

	var auth *acmeapi.Authorization

	ctx, cancelFunc := context.WithTimeout(context.Background(), LETSENCRYPT_CREATE_CERTIFICATE_TIMEOUT)
	defer cancelFunc()

	for i := 0; i < TRY_COUNT; i++ {
		auth, err = this.client.NewAuthorization(domain, ctx)
		if err == nil {
			break
		} else {
			logrus.Infof("Can't create new authorization for domain '%v': %v", domain, err)
			time.Sleep(RETRY_SLEEP)
		}
	}
	if err == nil {
		logrus.Infof("Create authorization for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't create new authorization for domain '%v': %v", domain, err)
		return nil, errors.New("Can't create new authorization for domain")
	}

	if logrus.GetLevel() >= logrus.DebugLevel {
		challengeTypes := make([]string, len(auth.Challenges))
		for i := range auth.Challenges {
			challengeTypes[i] = auth.Challenges[i].Type
		}
		logrus.Debugf("Challenge types for domain '%v': %v. Challenge combinations: %v", domain, challengeTypes, auth.Combinations)
	}

	canAuthorize := false
	var challenge *acmeapi.Challenge
	for _, cmb := range auth.Combinations {
		if len(cmb) == 1 && auth.Challenges[cmb[0]].Type == "tls-sni-01" {
			canAuthorize = true
			challenge = auth.Challenges[cmb[0]]
			break
		}
	}
	if !canAuthorize {
		logrus.Errorf("Can't find good challange combination for domain: '%v'", domain)
		return nil, errors.New("Can't find good challange combination")
	}

	acmeHostName, err := acmeutils.TLSSNIHostname(this.privateKey, challenge.Token)
	if err == nil {
		logrus.Debugf("Create acme-auth hostname for domain '%v': %v", domain, acmeHostName)
	} else {
		logrus.Errorf("Can't create acme domain for domain '%v' token '%v': %v", domain, challenge.Token, err)
		return nil, errors.New("Can't create acme domain")
	}
	this.PutAcmeAuthDomain(acmeHostName)

	challengeResponse, err := acmeutils.ChallengeResponseJSON(this.privateKey, challenge.Token, challenge.Type)
	if err == nil {
		logrus.Debugf("Create challenge response for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't create challenge response for domain '%v', token '%v', challenge type %v: %v",
			domain, challenge.Token, challenge.Type, err)
		return nil, errors.New("Can't create challenge response")
	}
	for i := 0; i < TRY_COUNT; i++ {
		err = this.client.RespondToChallenge(challenge, challengeResponse, this.privateKey, ctx)
		if err == nil {
			break
		} else {
			logrus.Info("Can't send response for challenge of domain '%v': %v", domain, err)
			time.Sleep(RETRY_SLEEP)
		}
	}
	if err == nil {
		logrus.Debugf("Send challenge response for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't send response for challenge of domain '%v': %v", domain, err)
		return nil, errors.New("Can't send response for challenge")
	}

	for i := 0; i < TRY_COUNT; i++ {
		err = this.client.LoadChallenge(challenge, ctx)
		if err != nil {
			logrus.Infof("Can't load challenge for domain '%v': %v", domain, err)
			time.Sleep(RETRY_SLEEP)
		}
	}
	if err == nil {
		logrus.Debugf("Load challenge for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't load challenge for domain '%v': %v", domain, err)
		return nil, errors.New("Can't load challenge")
	}

	// Generate CSR
	certKey, err := rsa.GenerateKey(cryptorand.Reader, PRIVATE_KEY_BITS)
	if err == nil {
		logrus.Debugf("Create private key for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't create rsa key for domain '%v': %v", err)
		return nil, errors.New("Can't create rsa key")
	}
	certRequest := &x509.CertificateRequest{
		SignatureAlgorithm: x509.SHA256WithRSA,
		PublicKeyAlgorithm: x509.RSA,
		PublicKey:          &certKey.PublicKey,
		Subject:            pkix.Name{CommonName: domain},
		DNSNames:           []string{domain},
	}
	csrDER, err := x509.CreateCertificateRequest(cryptorand.Reader, certRequest, certKey)
	if err == nil {
		logrus.Debugf("Create CSR for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't create csr for domain '%v': %v", domain, err)
		return nil, errors.New("Can't create csr")
	}

	var certResponse *acmeapi.Certificate
	for i := 0; i < TRY_COUNT; i++ {
		certResponse, err = this.client.RequestCertificate(csrDER, ctx)
		if err != nil {
			logrus.Infof("Can't get certificate for domain '%v': %v", domain, err)
			time.Sleep(RETRY_SLEEP)
		}
	}
	if err == nil {
		logrus.Infof("Get certificate for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't get certificate for domain '%v': %v", domain, err)
		return nil, errors.New("Can't request certificate")
	}
	cert = &tls.Certificate{}
	cert.Certificate = [][]byte{certResponse.Certificate}
	fmt.Printf("CERT:\n%s\n", certResponse.Certificate)
	cert.Leaf, err = x509.ParseCertificate(certResponse.Certificate)
	if err == nil {
		logrus.Debugf("Parse certificate for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't parse certificate for domain '%v':%v", domain, err)
		return nil, errors.New("Can't parse certificate")
	}
	fmt.Printf("CERT PARSED\n%#v\n", cert)

	certPEM := pem.EncodeToMemory(&pem.Block{Bytes: cert.Certificate, Type: "CERTIFICATE"})
	fmt.Printf("CERT PEM:\n%s\n", certPEM)

	return cert, nil
}

func (this *acmeStruct) createCertificateSelfSigned(domain string) (cert *tls.Certificate, err error) {
	derCert, privateKey, err := acmeutils.CreateTLSSNICertificate(domain)
	if err != nil {
		logrus.Errorf("Can't create tls-sni-01 self-signed certificate for '%v': %v", domain, err)
		return nil, err
	}

	cert = &tls.Certificate{}
	cert.Certificate = [][]byte{derCert}
	cert.PrivateKey = privateKey
	return cert, nil
}

func (this *acmeStruct) PutAcmeAuthDomain(domain string) {
	this.acmeauthDomainsMutex.Lock()
	defer this.acmeauthDomainsMutex.Unlock()

	logrus.Debug("Put acme auth domain:", domain)
	this.acmeAuthDomains[domain] = time.Now().Add(SNI01_EXPIRE_TOKEN)
}

func (this *acmeStruct) DeleteAcmeAuthDomain(domain string) bool {
	this.acmeauthDomainsMutex.Lock()
	defer this.acmeauthDomainsMutex.Unlock()

	logrus.Debug("Delete acme auth domain:", domain)
	_, ok := this.acmeAuthDomains[domain]
	if ok {
		delete(this.acmeAuthDomains, domain)
	}
	return ok
}
<commit_msg>cert response<commit_after>package main

import (
	"crypto/rsa"

	cryptorand "crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"errors"
	"github.com/Sirupsen/logrus"
	"github.com/hlandau/acme/acmeapi"
	"github.com/hlandau/acme/acmeapi/acmeutils"
	context "golang.org/x/net/context"
	"net"
	"strings"
	"sync"
	"time"
	"fmt"
	"encoding/pem"
)

const (
	SNI01_EXPIRE_TOKEN time.Duration = time.Minute * 10
)

type acmeStruct struct {
	serverAddress string
	privateKey    *rsa.PrivateKey
	client        *acmeapi.Client

	mutex                *sync.Mutex
	acmeauthDomainsMutex *sync.Mutex
	acmeAuthDomains      map[string]time.Time
}

func (this *acmeStruct) Init() {
	this.client = &acmeapi.Client{
		AccountKey:   this.privateKey,
		DirectoryURL: this.serverAddress,
	}

	this.mutex = &sync.Mutex{}

	this.acmeauthDomainsMutex = &sync.Mutex{}
	this.acmeAuthDomains = make(map[string]time.Time)
	this.CleanupTimer()
}

func (this *acmeStruct) RegisterEnsure(ctx context.Context) (err error) {
	reg := &acmeapi.Registration{}
	for i := 0; i < TRY_COUNT+1; i++ { // +1 count need for request latest agreement uri
		reg.AgreementURI = reg.LatestAgreementURI
		if reg.AgreementURI != "" {
			logrus.Info("Auto agree with terms:", reg.LatestAgreementURI)
		}
		err = this.client.UpsertRegistration(reg, ctx)
		if reg.AgreementURI != "" && err == nil {
			return
		}
	}
	return err
}

func (this *acmeStruct) Cleanup() {
	this.mutex.Lock()
	defer this.mutex.Unlock()

	now := time.Now()
	for token, expire := range this.acmeAuthDomains {
		if expire.Before(now) {
			delete(this.acmeAuthDomains, token)
		}
	}
}

func (this *acmeStruct) CleanupTimer() {
	this.Cleanup()
	time.AfterFunc(SNI01_EXPIRE_TOKEN, this.Cleanup)
}

func (this *acmeStruct) CreateCertificate(domain string) (cert *tls.Certificate, err error) {
	// Check suffix for avoid mutex sync in DeleteAcmeAuthDomain
	if strings.HasSuffix(domain, ".acme.invalid") {
		logrus.Debugf("Detect auth-domain mode for domain '%v'", domain)
		if this.DeleteAcmeAuthDomain(domain) {
			logrus.Debugf("Return self-signed certificate for domain '%v'", domain)
			return this.createCertificateSelfSigned(domain)
		} else {
			logrus.Debugf("Detect auth-domain is not present in list '%v'", domain)
			return nil, errors.New("Now allowed auth-domain")
		}
	}

	// check about we serve the domain
	ips, err := net.LookupIP(domain)
	if err != nil {
		logrus.Warnf("Can't lookup ip for domain '%v': %v", domain, err)
		return nil, errors.New("Can't lookup ip of the domain")
	}
	isLocalIP := false
checkLocalIP:
	for _, ip := range ips {
		for _, localIP := range localIPs {
			if ip.Equal(localIP) {
				isLocalIP = true
				break checkLocalIP
			}
		}
	}
	if !isLocalIP {
		logrus.Warnf("Domain have ip of other server. Domain '%v', Domain ips: %v, Server ips: %v", domain, ips, localIPs)
		return nil, errors.New("Domain have ip of other server.")
	}

	return this.createCertificateAcme(domain)
}

func (this *acmeStruct) createCertificateAcme(domain string) (cert *tls.Certificate, err error) {
	this.mutex.Lock()
	defer this.mutex.Unlock()

	var auth *acmeapi.Authorization

	ctx, cancelFunc := context.WithTimeout(context.Background(), LETSENCRYPT_CREATE_CERTIFICATE_TIMEOUT)
	defer cancelFunc()

	for i := 0; i < TRY_COUNT; i++ {
		auth, err = this.client.NewAuthorization(domain, ctx)
		if err == nil {
			break
		} else {
			logrus.Infof("Can't create new authorization for domain '%v': %v", domain, err)
			time.Sleep(RETRY_SLEEP)
		}
	}
	if err == nil {
		logrus.Infof("Create authorization for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't create new authorization for domain '%v': %v", domain, err)
		return nil, errors.New("Can't create new authorization for domain")
	}

	if logrus.GetLevel() >= logrus.DebugLevel {
		challengeTypes := make([]string, len(auth.Challenges))
		for i := range auth.Challenges {
			challengeTypes[i] = auth.Challenges[i].Type
		}
		logrus.Debugf("Challenge types for domain '%v': %v. Challenge combinations: %v", domain, challengeTypes, auth.Combinations)
	}

	canAuthorize := false
	var challenge *acmeapi.Challenge
	for _, cmb := range auth.Combinations {
		if len(cmb) == 1 && auth.Challenges[cmb[0]].Type == "tls-sni-01" {
			canAuthorize = true
			challenge = auth.Challenges[cmb[0]]
			break
		}
	}
	if !canAuthorize {
		logrus.Errorf("Can't find good challange combination for domain: '%v'", domain)
		return nil, errors.New("Can't find good challange combination")
	}

	acmeHostName, err := acmeutils.TLSSNIHostname(this.privateKey, challenge.Token)
	if err == nil {
		logrus.Debugf("Create acme-auth hostname for domain '%v': %v", domain, acmeHostName)
	} else {
		logrus.Errorf("Can't create acme domain for domain '%v' token '%v': %v", domain, challenge.Token, err)
		return nil, errors.New("Can't create acme domain")
	}
	this.PutAcmeAuthDomain(acmeHostName)

	challengeResponse, err := acmeutils.ChallengeResponseJSON(this.privateKey, challenge.Token, challenge.Type)
	if err == nil {
		logrus.Debugf("Create challenge response for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't create challenge response for domain '%v', token '%v', challenge type %v: %v",
			domain, challenge.Token, challenge.Type, err)
		return nil, errors.New("Can't create challenge response")
	}
	for i := 0; i < TRY_COUNT; i++ {
		err = this.client.RespondToChallenge(challenge, challengeResponse, this.privateKey, ctx)
		if err == nil {
			break
		} else {
			logrus.Info("Can't send response for challenge of domain '%v': %v", domain, err)
			time.Sleep(RETRY_SLEEP)
		}
	}
	if err == nil {
		logrus.Debugf("Send challenge response for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't send response for challenge of domain '%v': %v", domain, err)
		return nil, errors.New("Can't send response for challenge")
	}

	for i := 0; i < TRY_COUNT; i++ {
		err = this.client.LoadChallenge(challenge, ctx)
		if err != nil {
			logrus.Infof("Can't load challenge for domain '%v': %v", domain, err)
			time.Sleep(RETRY_SLEEP)
		}
	}
	if err == nil {
		logrus.Debugf("Load challenge for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't load challenge for domain '%v': %v", domain, err)
		return nil, errors.New("Can't load challenge")
	}

	// Generate CSR
	certKey, err := rsa.GenerateKey(cryptorand.Reader, PRIVATE_KEY_BITS)
	if err == nil {
		logrus.Debugf("Create private key for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't create rsa key for domain '%v': %v", err)
		return nil, errors.New("Can't create rsa key")
	}
	certRequest := &x509.CertificateRequest{
		SignatureAlgorithm: x509.SHA256WithRSA,
		PublicKeyAlgorithm: x509.RSA,
		PublicKey:          &certKey.PublicKey,
		Subject:            pkix.Name{CommonName: domain},
		DNSNames:           []string{domain},
	}
	csrDER, err := x509.CreateCertificateRequest(cryptorand.Reader, certRequest, certKey)
	if err == nil {
		logrus.Debugf("Create CSR for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't create csr for domain '%v': %v", domain, err)
		return nil, errors.New("Can't create csr")
	}

	var certResponse *acmeapi.Certificate
	for i := 0; i < TRY_COUNT; i++ {
		certResponse, err = this.client.RequestCertificate(csrDER, ctx)
		if err != nil {
			logrus.Infof("Can't get certificate for domain '%v': %v", domain, err)
			time.Sleep(RETRY_SLEEP)
		}
	}
	if err == nil {
		logrus.Infof("Get certificate for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't get certificate for domain '%v': %v", domain, err)
		return nil, errors.New("Can't request certificate")
	}
	cert = &tls.Certificate{}
	cert.Certificate = [][]byte{certResponse.Certificate}
	fmt.Printf("CERT:\n%s\n", certResponse.Certificate)
	cert.Leaf, err = x509.ParseCertificate(certResponse.Certificate)
	if err == nil {
		logrus.Debugf("Parse certificate for domain '%v'", domain)
	} else {
		logrus.Errorf("Can't parse certificate for domain '%v':%v", domain, err)
		return nil, errors.New("Can't parse certificate")
	}
	fmt.Printf("CERT PARSED\n%#v\n", cert)

	certPEM := pem.EncodeToMemory(&pem.Block{Bytes: certResponse.Certificate, Type: "CERTIFICATE"})
	fmt.Printf("CERT PEM:\n%s\n", certPEM)

	return cert, nil
}

func (this *acmeStruct) createCertificateSelfSigned(domain string) (cert *tls.Certificate, err error) {
	derCert, privateKey, err := acmeutils.CreateTLSSNICertificate(domain)
	if err != nil {
		logrus.Errorf("Can't create tls-sni-01 self-signed certificate for '%v': %v", domain, err)
		return nil, err
	}

	cert = &tls.Certificate{}
	cert.Certificate = [][]byte{derCert}
	cert.PrivateKey = privateKey
	return cert, nil
}

func (this *acmeStruct) PutAcmeAuthDomain(domain string) {
	this.acmeauthDomainsMutex.Lock()
	defer this.acmeauthDomainsMutex.Unlock()

	logrus.Debug("Put acme auth domain:", domain)
	this.acmeAuthDomains[domain] = time.Now().Add(SNI01_EXPIRE_TOKEN)
}

func (this *acmeStruct) DeleteAcmeAuthDomain(domain string) bool {
	this.acmeauthDomainsMutex.Lock()
	defer this.acmeauthDomainsMutex.Unlock()

	logrus.Debug("Delete acme auth domain:", domain)
	_, ok := this.acmeAuthDomains[domain]
	if ok {
		delete(this.acmeAuthDomains, domain)
	}
	return ok
}
<|endoftext|>


<commit_before>package otp

import (
	"bytes"
	"crypto/hmac"
	"encoding/base32"
	"encoding/binary"
	"hash"
	"strconv"
	"time"
)

type Hash func() hash.Hash

// Returns the current time interval.
func GetInterval(period int64) int64 {
	return time.Now().Unix() / period
}

// Returns the TOTP/HOTP code.
func GetCode(secret string, iv int64, h Hash, digits int) (string, error) {
	key, err := base32.StdEncoding.DecodeString(secret)
	if err != nil {
		return "", err
	}

	msg := bytes.Buffer{}
	_ = binary.Write(&msg, binary.BigEndian, iv)

	mac := hmac.New(h, key)
	mac.Write(msg.Bytes())
	digest := mac.Sum(nil)

	offset := digest[len(digest)-1] & 0xF
	trunc := digest[offset : offset+4]

	var code int32
	truncBytes := bytes.NewBuffer(trunc)
	_ = binary.Read(truncBytes, binary.BigEndian, &code)

	code = (code & 0x7FFFFFFF) % 1000000
	stringCode := strconv.Itoa(int(code))
	for len(stringCode) < digits {
		stringCode = "0" + stringCode
	}
	return stringCode, nil
}
<commit_msg>header<commit_after>/*
One-time password (OTP) implementation. Supports HMAC-based (HOTP) and time-based (TOTP) methods.
*/
package otp

import (
	"bytes"
	"crypto/hmac"
	"encoding/base32"
	"encoding/binary"
	"hash"
	"strconv"
	"time"
)

type Hash func() hash.Hash

// Returns the current time interval.
func GetInterval(period int64) int64 {
	return time.Now().Unix() / period
}

// Returns the TOTP/HOTP code.
func GetCode(secret string, iv int64, h Hash, digits int) (string, error) {
	key, err := base32.StdEncoding.DecodeString(secret)
	if err != nil {
		return "", err
	}

	msg := bytes.Buffer{}
	_ = binary.Write(&msg, binary.BigEndian, iv)

	mac := hmac.New(h, key)
	mac.Write(msg.Bytes())
	digest := mac.Sum(nil)

	offset := digest[len(digest)-1] & 0xF
	trunc := digest[offset : offset+4]

	var code int32
	truncBytes := bytes.NewBuffer(trunc)
	_ = binary.Read(truncBytes, binary.BigEndian, &code)

	code = (code & 0x7FFFFFFF) % 1000000
	stringCode := strconv.Itoa(int(code))
	for len(stringCode) < digits {
		stringCode = "0" + stringCode
	}
	return stringCode, nil
}
<|endoftext|>


<commit_before>package shimmie

import (
	"crypto/md5"
	"database/sql"
	"fmt"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	"golang.org/x/net/context"
)

// Hash returns the MD5 checksum of a string s as type string.
func Hash(s string) string {
	return fmt.Sprintf("%x", md5.Sum([]byte(s)))
}

// Auth is a handler wrapper that checks if a user is authenticated to Shimmie.
// It checks for two cookies "shm_user" and "shm_session". The first
// contains the username which is used to query the database and the get user's
// password hash. Then it attempts to recreate the "shm_session" cookie value
// by using the username, user IP and password hash. If the recreated value
// does not match the "shm_session" cookie value then it redirects to
// redirectPath. If redirectURL is empty then "/user_admin/login" is used
// instead which is the default login URL for Shimmie.
func (shim *Shimmie) Auth(ctx context.Context, fn func(context.Context, http.ResponseWriter, *http.Request), redirectURL string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		const defaultLoginURL = "/user_admin/login"
		if redirectURL == "" {
			redirectURL = defaultLoginURL
		}
		usernameCookie, err := r.Cookie("shm_user")
		if err != nil || usernameCookie.Value == "" {
			http.Redirect(w, r, redirectURL, http.StatusFound)
			return
		}
		sessionCookie, err := r.Cookie("shm_session")
		if err != nil {
			log.Print("shimmie: no session cookie")
			http.Redirect(w, r, redirectURL, http.StatusFound)
			return
		}
		username := usernameCookie.Value
		user, err := shim.Store.GetUser(username)
		if err != nil {
			if err == sql.ErrNoRows {
				log.Printf("shimmie: user %q does not exist", username)
				http.Redirect(w, r, redirectURL, http.StatusFound)
				return
			}
			msg := fmt.Sprintf("shimmie: could not authenticate: get user %q failed: %v", username, err.Error())
			log.Print(msg)
			http.Error(w, msg, http.StatusInternalServerError)
			return
		}
		passwordHash := user.Pass
		userIP := getOriginalIP(r)
		sessionCookieValue := CookieValue(passwordHash, userIP)
		if sessionCookieValue != sessionCookie.Value {
			http.Redirect(w, r, redirectURL, http.StatusFound)
			return
		}
		ctx = context.WithValue(ctx, "user", user)
		fn(ctx, w, r)
	}
}

func getOriginalIP(r *http.Request) string {
	x := r.Header.Get("X-Forwarded-For")
	if x != "" && strings.Contains(r.RemoteAddr, "127.0.0.1") {
		// format is comma separated
		return strings.Split(x, ",")[0]
	}
	// it also contains the port
	return strings.Split(r.RemoteAddr, ":")[0]
}

// CookieValue recreates the Shimmie session cookie value based on the user
// password hash and the user IP.
//
// Shimmie creates a cookie "shm_session" containing an md5 digest value of the
// user password hash concatenated with the user IP masked with the 255.255.0.0
// mask. That's essentially:
//
//   md5(password_hash + masked_ip)
//
func CookieValue(passwordHash, userIP string) string {
	addr := net.ParseIP(strings.Split(userIP, ":")[0])
	mask := net.IPv4Mask(255, 255, 0, 0)
	addr = addr.Mask(mask)
	sessionHash := md5.Sum([]byte(fmt.Sprintf("%s%s", passwordHash, addr)))
	return fmt.Sprintf("%x", sessionHash)
}

const loginMemory = 365

// SetCookie creates a cookie on path "/" with 1 year expiration and other
// flags set to false mimicking the cookies that Shimmie creates.
func SetCookie(w http.ResponseWriter, name, value string) {
	expires := time.Now().Add(time.Second * 60 * 60 * 24 * loginMemory)
	c := http.Cookie{
		Name:    name,
		Value:   value,
		Expires: expires,
		Path:    "/",
	}
	http.SetCookie(w, &c)
}
<commit_msg>Make GetOriginalIP func exported<commit_after>package shimmie

import (
	"crypto/md5"
	"database/sql"
	"fmt"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	"golang.org/x/net/context"
)

// Hash returns the MD5 checksum of a string s as type string.
func Hash(s string) string {
	return fmt.Sprintf("%x", md5.Sum([]byte(s)))
}

// Auth is a handler wrapper that checks if a user is authenticated to Shimmie.
// It checks for two cookies "shm_user" and "shm_session". The first
// contains the username which is used to query the database and the get user's
// password hash. Then it attempts to recreate the "shm_session" cookie value
// by using the username, user IP and password hash. If the recreated value
// does not match the "shm_session" cookie value then it redirects to
// redirectPath. If redirectURL is empty then "/user_admin/login" is used
// instead which is the default login URL for Shimmie.
func (shim *Shimmie) Auth(ctx context.Context, fn func(context.Context, http.ResponseWriter, *http.Request), redirectURL string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		const defaultLoginURL = "/user_admin/login"
		if redirectURL == "" {
			redirectURL = defaultLoginURL
		}
		usernameCookie, err := r.Cookie("shm_user")
		if err != nil || usernameCookie.Value == "" {
			http.Redirect(w, r, redirectURL, http.StatusFound)
			return
		}
		sessionCookie, err := r.Cookie("shm_session")
		if err != nil {
			log.Print("shimmie: no session cookie")
			http.Redirect(w, r, redirectURL, http.StatusFound)
			return
		}
		username := usernameCookie.Value
		user, err := shim.Store.GetUser(username)
		if err != nil {
			if err == sql.ErrNoRows {
				log.Printf("shimmie: user %q does not exist", username)
				http.Redirect(w, r, redirectURL, http.StatusFound)
				return
			}
			msg := fmt.Sprintf("shimmie: could not authenticate: get user %q failed: %v", username, err.Error())
			log.Print(msg)
			http.Error(w, msg, http.StatusInternalServerError)
			return
		}
		passwordHash := user.Pass
		userIP := GetOriginalIP(r)
		sessionCookieValue := CookieValue(passwordHash, userIP)
		if sessionCookieValue != sessionCookie.Value {
			http.Redirect(w, r, redirectURL, http.StatusFound)
			return
		}
		ctx = context.WithValue(ctx, "user", user)
		fn(ctx, w, r)
	}
}

// GetOriginalIP gets the original IP of the HTTP for the case of being behind
// a proxy. It searches for the X-Forwarded-For header.
func GetOriginalIP(r *http.Request) string {
	x := r.Header.Get("X-Forwarded-For")
	if x != "" && strings.Contains(r.RemoteAddr, "127.0.0.1") {
		// format is comma separated
		return strings.Split(x, ",")[0]
	}
	// it also contains the port
	return strings.Split(r.RemoteAddr, ":")[0]
}

// CookieValue recreates the Shimmie session cookie value based on the user
// password hash and the user IP.
//
// Shimmie creates a cookie "shm_session" containing an md5 digest value of the
// user password hash concatenated with the user IP masked with the 255.255.0.0
// mask. That's essentially:
//
//   md5(password_hash + masked_ip)
//
func CookieValue(passwordHash, userIP string) string {
	addr := net.ParseIP(strings.Split(userIP, ":")[0])
	mask := net.IPv4Mask(255, 255, 0, 0)
	addr = addr.Mask(mask)
	sessionHash := md5.Sum([]byte(fmt.Sprintf("%s%s", passwordHash, addr)))
	return fmt.Sprintf("%x", sessionHash)
}

const loginMemory = 365

// SetCookie creates a cookie on path "/" with 1 year expiration and other
// flags set to false mimicking the cookies that Shimmie creates.
func SetCookie(w http.ResponseWriter, name, value string) {
	expires := time.Now().Add(time.Second * 60 * 60 * 24 * loginMemory)
	c := http.Cookie{
		Name:    name,
		Value:   value,
		Expires: expires,
		Path:    "/",
	}
	http.SetCookie(w, &c)
}
<|endoftext|>


<commit_before>package env

import (
	"errors"
	"reflect"
	"strconv"
)

// ErrNotAStructPtr is returned if you pass something that is not a pointer to a
// Struct to Parse
var ErrNotAStructPtr = errors.New("Expected a pointer to a Struct")

// ErrUnsuportedType if the struct field type is not supported by env
var ErrUnsuportedType = errors.New("Type is not supported")

// Parse parses a struct containing `env` tags and loads its values from
// environment variables.
func Parse(val interface{}) error {
	ptrRef := reflect.ValueOf(val)
	if ptrRef.Kind() != reflect.Ptr {
		return ErrNotAStructPtr
	}
	ref := ptrRef.Elem()
	if ref.Kind() != reflect.Struct {
		return ErrNotAStructPtr
	}
	refType := ref.Type()
	for i := 0; i < refType.NumField(); i++ {
		field := refType.Field(i)
		value := get(field)
		if value == "" {
			continue
		}
		if err := set(ref.Field(i), value); err != nil {
			return err
		}
	}
	return nil
}

func get(field reflect.StructField) string {
	defaultValue := field.Tag.Get("default")
	return GetOr(field.Tag.Get("env"), defaultValue)
}

func set(field reflect.Value, value string) error {
	switch field.Kind() {
	case reflect.String:
		field.SetString(value)
	case reflect.Bool:
		bvalue, err := strconv.ParseBool(value)
		if err != nil {
			return err
		}
		field.SetBool(bvalue)
	case reflect.Int:
		intValue, err := strconv.ParseInt(value, 10, 32)
		if err != nil {
			return err
		}
		field.SetInt(intValue)
	}
	return nil
}
<commit_msg>refactoring<commit_after>package env

import (
	"errors"
	"reflect"
	"strconv"
)

// ErrNotAStructPtr is returned if you pass something that is not a pointer to a
// Struct to Parse
var ErrNotAStructPtr = errors.New("Expected a pointer to a Struct")

// ErrUnsuportedType if the struct field type is not supported by env
var ErrUnsuportedType = errors.New("Type is not supported")

// Parse parses a struct containing `env` tags and loads its values from
// environment variables.
func Parse(val interface{}) error {
	ptrRef := reflect.ValueOf(val)
	if ptrRef.Kind() != reflect.Ptr {
		return ErrNotAStructPtr
	}
	ref := ptrRef.Elem()
	if ref.Kind() != reflect.Struct {
		return ErrNotAStructPtr
	}
	refType := ref.Type()
	for i := 0; i < refType.NumField(); i++ {
		value := get(refType.Field(i))
		if value == "" {
			continue
		}
		if err := set(ref.Field(i), value); err != nil {
			return err
		}
	}
	return nil
}

func get(field reflect.StructField) string {
	defaultValue := field.Tag.Get("default")
	return GetOr(field.Tag.Get("env"), defaultValue)
}

func set(field reflect.Value, value string) error {
	switch field.Kind() {
	case reflect.String:
		field.SetString(value)
	case reflect.Bool:
		bvalue, err := strconv.ParseBool(value)
		if err != nil {
			return err
		}
		field.SetBool(bvalue)
	case reflect.Int:
		intValue, err := strconv.ParseInt(value, 10, 32)
		if err != nil {
			return err
		}
		field.SetInt(intValue)
	}
	return nil
}
<|endoftext|>


<commit_before>package DatabaseConnector

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"sync"
	"time"

	mdl "github.com/LoRaWanSoFa/LoRaWanSoFa/Components"
	_ "github.com/lib/pq"
)

type DatabaseConnector struct {
	Database                  *sql.DB
	CheckDevEUISTMT           *sql.Stmt
	GetNodeHeaderSTMT         *sql.Stmt
	InsertMessageSTMT         *sql.Stmt
	InsertPayloadSTMT         *sql.Stmt
	InsertDownlinkMessageSTMT *sql.Stmt
}

type WorkRequest struct {
	Query         string
	ResultChannel chan (WorkResult)
	Arguments     []interface{}
	F             func(w *WorkRequest)
}

type WorkResult struct {
	Result interface{}
	err    error
}

func newDBC(db *sql.DB) *DatabaseConnector {
	dbc := DatabaseConnector{}
	dbc.Database = db
	return &dbc
}

var instantiated *DatabaseConnector
var once sync.Once
var WorkQueue = make(chan WorkRequest, 100)

// Get the instantiated instance of the DatabaseConnector or create it.
func GetInstance() *DatabaseConnector {
	once.Do(func() {
		settings := mdl.GetConfiguration().Db
		dbConnectionInfo := fmt.Sprintf("user=%s password=%s dbname=%s host=%s port=%v sslmode=disable", settings.User, settings.Password, settings.Name, settings.Network, settings.Port)
		println(dbConnectionInfo)
		actualDb, err := sql.Open("postgres", dbConnectionInfo)
		instantiated = newDBC(actualDb)
		if err != nil {
			log.Fatal(err)
		}
		log.Print("not connected yet")
		StartDispatcher(settings.NumberOfWorkers)
	})
	return instantiated
}

func Connect() error {
	//IDEA move sql.Open to here so we can reconnect.
	db := GetInstance()
	err := db.Database.Ping()
	if err != nil {
		return err
	}
	time.Sleep(1000 * time.Millisecond) //Wait for connections to finnish setting up. 1 sec is propably too long, can be finetuned later
	db.CheckDevEUISTMT, err = db.Database.Prepare("select exists(select 1 from nodes where deveui=$1)")
	if err != nil {
		return err
	}
	db.GetNodeHeaderSTMT, err = db.Database.Prepare("select sensors.id, number_of_values, lenght_of_values, header_order, conversion_expression, data_type " +
		"from sensors " +
		"join public.sensortypes on sensors.sensortype_id = sensortypes.id " +
		"where deveui =$1 " +
		"order by header_order asc;")
	if err != nil {
		return err
	}
	db.InsertMessageSTMT, err = db.Database.Prepare("INSERT INTO public.messages(" +
		"deveui, created_at, down) " +
		"VALUES ($1, NOW(), false) " +
		"RETURNING id;")
	if err != nil {
		return err
	}
	db.InsertDownlinkMessageSTMT, err = db.Database.Prepare("INSERT INTO public.messages(" +
		"deveui, created_at, down) " +
		"VALUES ($1, $2, true) " +
		"RETURNING id;")
	if err != nil {
		return err
	}
	db.InsertPayloadSTMT, err = db.Database.Prepare("INSERT INTO public.message_payloads(" +
		"message_id, sensor_id, payload) " +
		"VALUES ($1, $2, $3);")
	return err
}

func Close() error {
	return GetInstance().Database.Close()
}

//Not used, but can be used as example
// Checks if the devEUI exists in the database.
// Uses a database worker to execute the query.
func CheckDevEUI(devEUI string) bool {
	//log.Println(devEUI)
	result := make(chan WorkResult)
	args := make([]interface{}, 1)
	args[0] = devEUI
	WorkQueue <- WorkRequest{Query: "", Arguments: args, ResultChannel: result, F: func(w *WorkRequest) {
		rows, err := GetInstance().CheckDevEUISTMT.Query(w.Arguments...)
		defer rows.Close()
		checkErr(err)
		if err != nil {
			w.ResultChannel <- WorkResult{Result: false, err: err}
			return
		}
		rows.Next()
		var exisistsboolean = false
		err = rows.Scan(&exisistsboolean)
		w.ResultChannel <- WorkResult{Result: exisistsboolean, err: err}
	}}
	defer close(result)
	var workResult = <-result
	checkErr(workResult.err)
	exists := workResult.Result.(bool)
	log.Printf("Work done for Q:%s A:%t\n", devEUI, exists)
	return exists
}

//Store a message
func AddMessage(devEUI string) (mdl.MessageUplinkI, error) {
	//create response channel
	result := make(chan WorkResult)
	defer close(result)
	//create arguments
	args := make([]interface{}, 1)
	args[0] = devEUI
	//create and add new WorkRequest
	WorkQueue <- WorkRequest{Query: "", Arguments: args, ResultChannel: result, F: func(w *WorkRequest) {
		var messageId int64
		rows, err := GetInstance().InsertMessageSTMT.Query(w.Arguments...)

		checkErr(err)
		if err != nil {
			log.Println("check errors1")
			w.ResultChannel <- WorkResult{Result: 0, err: err}
			return
		}
		defer rows.Close()

		rows.Next()
		err = rows.Scan(&messageId)
		if err != nil {
			w.ResultChannel <- WorkResult{Result: 0, err: err}
		}
		w.ResultChannel <- WorkResult{Result: messageId, err: nil}
	}}
	response := <-result
	if response.err != nil {
		log.Println("check errors3")
		return nil, response.err
	}
	message := mdl.NewMessageUplink(response.Result.(int64), devEUI)

	return message, response.err
}

//Get a message
func StoreMessagePayloads(message mdl.MessageUplinkI) error {
	if message == nil {
		return errors.New("nil given as message parameter")
	}
	if message.GetId() == 0 {
		return errors.New("Message has not been initalized/stored")
	}
	payloads := message.GetPayloads()
	if len(payloads) == 0 {
		return errors.New("Nothing to store!")
	}
	var parameters []interface{}
	var err error
	for _, payload := range payloads {
		parameters = make([]interface{}, 0)
		parameters = append(parameters, message.GetId())               //message id
		parameters = append(parameters, payload.GetSensor().Id)        //sensor id
		parameters = append(parameters, payload.GetPayload().(string)) //payload
		//log.Printf("parameters: %+v", parameters)
		err = insertPayload(parameters)
	}
	if err != nil {
		return nil
	}
	return nil
}

func insertPayload(parameters []interface{}) error {
	result := make(chan WorkResult)
	defer close(result)
	WorkQueue <- WorkRequest{Query: "", Arguments: parameters, ResultChannel: result, F: func(w *WorkRequest) {
		rows, err := GetInstance().InsertPayloadSTMT.Query(w.Arguments...)

		checkErr(err)
		if err != nil {
			log.Println("Query could not be executed!")
			w.ResultChannel <- WorkResult{Result: nil, err: err}
			return
		}
		defer rows.Close()

		rows.Next()
		if err != nil {
			w.ResultChannel <- WorkResult{Result: nil, err: err}
		}
		w.ResultChannel <- WorkResult{Result: true, err: nil}
	}}
	response := <-result
	if response.err != nil {
		log.Println("The worker could not execute the work properly")
		return response.err
	}
	return nil
}

//Stores a DownlinkMessage which has an id,payload and deveui set.
//if no time is set NOW() will be used
func StoreDownlinkMessage(message *mdl.MessageDownLink) error {
	if message.Id != 0 {
		return errors.New("Message already has an id, can not insert it")
	}
	if message.Payload == "" {
		return errors.New("Message has an empty payload")
	}
	if message.Deveui == "" {
		return errors.New("Message has no DevEUI set")
	}
	if message.Time.IsZero() {
		message.Time = time.Now()
	}
	err := addDownlinkMessage(message)
	if err != nil {
		return err
	}
	var parameters []interface{}
	parameters = append(parameters, message.Id)
	parameters = append(parameters, nil)
	parameters = append(parameters, message.Payload)
	err = insertPayload(parameters)
	return err
}

// Executes the query to insert the message
// Formats the time to UTC and rounds it to second percision.
// Sets a new id for the message
func addDownlinkMessage(message *mdl.MessageDownLink) error {
	//create response channel
	result := make(chan WorkResult)
	defer close(result)
	//create arguments
	args := make([]interface{}, 2)
	args[0] = message.Deveui
	args[1] = message.Time.UTC().Round(time.Second).Format(time.RFC3339)
	//create and add new WorkRequest
	WorkQueue <- WorkRequest{Query: "", Arguments: args, ResultChannel: result, F: func(w *WorkRequest) {
		var messageId int64
		rows, err := GetInstance().InsertDownlinkMessageSTMT.Query(w.Arguments...)

		checkErr(err)
		if err != nil {
			log.Println("Query could not be executed")
			w.ResultChannel <- WorkResult{Result: 0, err: err}
			return
		}
		defer rows.Close()

		rows.Next()
		err = rows.Scan(&messageId)
		if err != nil {
			w.ResultChannel <- WorkResult{Result: 0, err: err}
		}
		w.ResultChannel <- WorkResult{Result: messageId, err: nil}
	}}
	response := <-result
	if response.err != nil {
		log.Println("Worker could finnish its work properly")
		return response.err
	}
	message.Id = response.Result.(int64)

	return response.err
}

//Get sensors that belong to one node
func GetNodeSensors(devEUI string) []mdl.Sensor {
	result := make(chan WorkResult)
	defer close(result)
	args := make([]interface{}, 1)
	args[0] = devEUI
	WorkQueue <- WorkRequest{Query: "", Arguments: args, ResultChannel: result, F: func(w *WorkRequest) {
		//rows, err := GetInstance().Database.Query(w.Query, w.Arguments...)
		rows, err := GetInstance().GetNodeHeaderSTMT.Query(w.Arguments...)
		defer rows.Close()
		checkErr(err)
		if err != nil {
			w.ResultChannel <- WorkResult{Result: false, err: err}
			return
		}
		sensors := make([]mdl.Sensor, 0)
		var id int64
		var number_of_values, lenght_of_values, header_order, data_type int
		var conversion_expression string

		for rows.Next() {
			err = rows.Scan(&id, &number_of_values, &lenght_of_values, &header_order, &conversion_expression, &data_type)
			if err != nil {
				panic(err.Error())
			}
			s := mdl.NewHeaderSensor(
				id,
				number_of_values,
				lenght_of_values,
				header_order,
				data_type,
				conversion_expression)
			sensors = append(sensors, s)
		}
		w.ResultChannel <- WorkResult{Result: sensors, err: err}
	}}
	var workResult = <-result
	if workResult.err != nil {
		log.Printf("A problem occured when getting the sensorheaders: %+v", workResult.err)
		return make([]mdl.Sensor, 0)
	}
	sensors := workResult.Result.([]mdl.Sensor)
	return sensors
}

func isOffline() bool {
	return instantiated == nil
}

func checkErr(err error) {
	if err != nil {
		log.Println(err)
	}
}

func panicErr(err error) {
	if err != nil {
		log.Println(err)
		panic(err)
	}
}
<commit_msg>added update header function<commit_after>package DatabaseConnector

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"sync"
	"time"

	mdl "github.com/LoRaWanSoFa/LoRaWanSoFa/Components"
	_ "github.com/lib/pq"
)

type DatabaseConnector struct {
	Database                  *sql.DB
	CheckDevEUISTMT           *sql.Stmt
	GetNodeHeaderSTMT         *sql.Stmt
	InsertMessageSTMT         *sql.Stmt
	InsertPayloadSTMT         *sql.Stmt
	InsertDownlinkMessageSTMT *sql.Stmt
}

type WorkRequest struct {
	Query         string
	ResultChannel chan (WorkResult)
	Arguments     []interface{}
	F             func(w *WorkRequest)
}

type WorkResult struct {
	Result interface{}
	err    error
}

func newDBC(db *sql.DB) *DatabaseConnector {
	dbc := DatabaseConnector{}
	dbc.Database = db
	return &dbc
}

var instantiated *DatabaseConnector
var once sync.Once
var WorkQueue = make(chan WorkRequest, 100)

// Get the instantiated instance of the DatabaseConnector or create it.
func GetInstance() *DatabaseConnector {
	once.Do(func() {
		settings := mdl.GetConfiguration().Db
		dbConnectionInfo := fmt.Sprintf("user=%s password=%s dbname=%s host=%s port=%v sslmode=disable", settings.User, settings.Password, settings.Name, settings.Network, settings.Port)
		println(dbConnectionInfo)
		actualDb, err := sql.Open("postgres", dbConnectionInfo)
		instantiated = newDBC(actualDb)
		if err != nil {
			log.Fatal(err)
		}
		log.Print("not connected yet")
		StartDispatcher(settings.NumberOfWorkers)
	})
	return instantiated
}

func Connect() error {
	//IDEA move sql.Open to here so we can reconnect.
	db := GetInstance()
	err := db.Database.Ping()
	if err != nil {
		return err
	}
	time.Sleep(1000 * time.Millisecond) //Wait for connections to finnish setting up. 1 sec is propably too long, can be finetuned later
	db.CheckDevEUISTMT, err = db.Database.Prepare("select exists(select 1 from nodes where deveui=$1)")
	if err != nil {
		return err
	}
	db.GetNodeHeaderSTMT, err = db.Database.Prepare("select sensors.id, number_of_values, lenght_of_values, header_order, conversion_expression, data_type " +
		"from sensors " +
		"join public.sensortypes on sensors.sensortype_id = sensortypes.id " +
		"where deveui =$1 and header_order >= 0" +
		"order by header_order asc;")
	if err != nil {
		return err
	}
	db.InsertMessageSTMT, err = db.Database.Prepare("INSERT INTO public.messages(" +
		"deveui, created_at, down) " +
		"VALUES ($1, NOW(), false) " +
		"RETURNING id;")
	if err != nil {
		return err
	}
	db.InsertDownlinkMessageSTMT, err = db.Database.Prepare("INSERT INTO public.messages(" +
		"deveui, created_at, down) " +
		"VALUES ($1, $2, true) " +
		"RETURNING id;")
	if err != nil {
		return err
	}
	db.InsertPayloadSTMT, err = db.Database.Prepare("INSERT INTO public.message_payloads(" +
		"message_id, sensor_id, payload) " +
		"VALUES ($1, $2, $3);")
	return err
}

func Close() error {
	return GetInstance().Database.Close()
}

// Checks if the devEUI exists in the database.
// Uses a database worker to execute the query.
func CheckDevEUI(devEUI string) bool {
	//log.Println(devEUI)
	result := make(chan WorkResult)
	args := make([]interface{}, 1)
	args[0] = devEUI
	WorkQueue <- WorkRequest{Query: "", Arguments: args, ResultChannel: result, F: func(w *WorkRequest) {
		rows, err := GetInstance().CheckDevEUISTMT.Query(w.Arguments...)
		defer rows.Close()
		checkErr(err)
		if err != nil {
			w.ResultChannel <- WorkResult{Result: false, err: err}
			return
		}
		rows.Next()
		var exisistsboolean = false
		err = rows.Scan(&exisistsboolean)
		w.ResultChannel <- WorkResult{Result: exisistsboolean, err: err}
	}}
	defer close(result)
	var workResult = <-result
	checkErr(workResult.err)
	exists := workResult.Result.(bool)
	log.Printf("Work done for Q:%s A:%t\n", devEUI, exists)
	return exists
}

//Store a message
func AddMessage(devEUI string) (mdl.MessageUplinkI, error) {
	//create response channel
	result := make(chan WorkResult)
	defer close(result)
	//create arguments
	args := make([]interface{}, 1)
	args[0] = devEUI
	//create and add new WorkRequest
	WorkQueue <- WorkRequest{Query: "", Arguments: args, ResultChannel: result, F: func(w *WorkRequest) {
		var messageId int64
		rows, err := GetInstance().InsertMessageSTMT.Query(w.Arguments...)

		checkErr(err)
		if err != nil {
			log.Println("check errors1")
			w.ResultChannel <- WorkResult{Result: 0, err: err}
			return
		}
		defer rows.Close()

		rows.Next()
		err = rows.Scan(&messageId)
		if err != nil {
			w.ResultChannel <- WorkResult{Result: 0, err: err}
		}
		w.ResultChannel <- WorkResult{Result: messageId, err: nil}
	}}
	response := <-result
	if response.err != nil {
		log.Println("check errors3")
		return nil, response.err
	}
	message := mdl.NewMessageUplink(response.Result.(int64), devEUI)

	return message, response.err
}

//Get a message
func StoreMessagePayloads(message mdl.MessageUplinkI) error {
	if message == nil {
		return errors.New("nil given as message parameter")
	}
	if message.GetId() == 0 {
		return errors.New("Message has not been initalized/stored")
	}
	payloads := message.GetPayloads()
	if len(payloads) == 0 {
		return errors.New("Nothing to store!")
	}
	var parameters []interface{}
	var err error
	for _, payload := range payloads {
		parameters = make([]interface{}, 0)
		parameters = append(parameters, message.GetId())               //message id
		parameters = append(parameters, payload.GetSensor().Id)        //sensor id
		parameters = append(parameters, payload.GetPayload().(string)) //payload
		//log.Printf("parameters: %+v", parameters)
		err = insertPayload(parameters)
	}
	if err != nil {
		return nil
	}
	return nil
}

func insertPayload(parameters []interface{}) error {
	result := make(chan WorkResult)
	defer close(result)
	WorkQueue <- WorkRequest{Query: "", Arguments: parameters, ResultChannel: result, F: func(w *WorkRequest) {
		rows, err := GetInstance().InsertPayloadSTMT.Query(w.Arguments...)

		checkErr(err)
		if err != nil {
			log.Println("Query could not be executed!")
			w.ResultChannel <- WorkResult{Result: nil, err: err}
			return
		}
		defer rows.Close()

		rows.Next()
		if err != nil {
			w.ResultChannel <- WorkResult{Result: nil, err: err}
		}
		w.ResultChannel <- WorkResult{Result: true, err: nil}
	}}
	response := <-result
	if response.err != nil {
		log.Println("The worker could not execute the work properly")
		return response.err
	}
	return nil
}

//Stores a DownlinkMessage which has an id,payload and deveui set.
//if no time is set NOW() will be used
func StoreDownlinkMessage(message *mdl.MessageDownLink) error {
	if message.Id != 0 {
		return errors.New("Message already has an id, can not insert it")
	}
	if message.Payload == "" {
		return errors.New("Message has an empty payload")
	}
	if message.Deveui == "" {
		return errors.New("Message has no DevEUI set")
	}
	if message.Time.IsZero() {
		message.Time = time.Now()
	}
	err := addDownlinkMessage(message)
	if err != nil {
		return err
	}
	var parameters []interface{}
	parameters = append(parameters, message.Id)
	parameters = append(parameters, nil)
	parameters = append(parameters, message.Payload)
	err = insertPayload(parameters)
	return err
}

// Executes the query to insert the message
// Formats the time to UTC and rounds it to second percision.
// Sets a new id for the message
func addDownlinkMessage(message *mdl.MessageDownLink) error {
	//create response channel
	result := make(chan WorkResult)
	defer close(result)
	//create arguments
	args := make([]interface{}, 2)
	args[0] = message.Deveui
	args[1] = message.Time.UTC().Round(time.Second).Format(time.RFC3339)
	//create and add new WorkRequest
	WorkQueue <- WorkRequest{Query: "", Arguments: args, ResultChannel: result, F: func(w *WorkRequest) {
		var messageId int64
		rows, err := GetInstance().InsertDownlinkMessageSTMT.Query(w.Arguments...)

		checkErr(err)
		if err != nil {
			log.Println("Query could not be executed")
			w.ResultChannel <- WorkResult{Result: 0, err: err}
			return
		}
		defer rows.Close()

		rows.Next()
		err = rows.Scan(&messageId)
		if err != nil {
			w.ResultChannel <- WorkResult{Result: 0, err: err}
		}
		w.ResultChannel <- WorkResult{Result: messageId, err: nil}
	}}
	response := <-result
	if response.err != nil {
		log.Println("Worker could finnish its work properly")
		return response.err
	}
	message.Id = response.Result.(int64)

	return response.err
}

//Get sensors that belong to one node
func GetNodeSensors(devEUI string) []mdl.Sensor {
	result := make(chan WorkResult)
	defer close(result)
	args := make([]interface{}, 1)
	args[0] = devEUI
	WorkQueue <- WorkRequest{Query: "", Arguments: args, ResultChannel: result, F: func(w *WorkRequest) {
		//rows, err := GetInstance().Database.Query(w.Query, w.Arguments...)
		rows, err := GetInstance().GetNodeHeaderSTMT.Query(w.Arguments...)
		defer rows.Close()
		checkErr(err)
		if err != nil {
			w.ResultChannel <- WorkResult{Result: false, err: err}
			return
		}
		sensors := make([]mdl.Sensor, 0)
		var id int64
		var number_of_values, lenght_of_values, header_order, data_type int
		var conversion_expression string

		for rows.Next() {
			err = rows.Scan(&id, &number_of_values, &lenght_of_values, &header_order, &conversion_expression, &data_type)
			if err != nil {
				panic(err.Error())
			}
			s := mdl.NewHeaderSensor(
				id,
				number_of_values,
				lenght_of_values,
				header_order,
				data_type,
				conversion_expression)
			sensors = append(sensors, s)
		}
		w.ResultChannel <- WorkResult{Result: sensors, err: err}
	}}
	var workResult = <-result
	if workResult.err != nil {
		log.Printf("A problem occured when getting the sensorheaders: %+v", workResult.err)
		return make([]mdl.Sensor, 0)
	}
	sensors := workResult.Result.([]mdl.Sensor)
	return sensors
}

func UpdateHeader(devEUI string, newheader []mdl.Sensor) error {
	//check first
	if len(newheader) == 0 {
		return errors.New("No Sensors given")
	}
	if devEUI == "" {
		return errors.New("Deveui must not be empty")
	}
	if !CheckDevEUI(devEUI) {
		return errors.New("Deveui does not exist")
	}
	oldheader := GetNodeSensors(devEUI)
	newSensortyps := make(map[string]mdl.Sensor)
	if len(oldheader) == 0 {
		//old header does not exist, insert new ones

		for _, newS := range newheader {
			newSensortyps[fmt.Sprintf("%s %d", newS.Conversion_expression, newS.DataType)] = newS
		}
		for _, sensor := range newSensortyps {
			//Query: insert new sensortypes if not exist
			log.Printf("insert into sensor_typs where !%+v and !%+v on dubplicate ignore", sensor.Conversion_expression, sensor.DataType)
		}
		//Query: insert new sensors with Deveui
		return nil
	}
	sensortyps := make(map[string]bool)
	for _, s := range oldheader {
		sensortyps[fmt.Sprintf("%s %d", s.Conversion_expression, s.DataType)] = true
	}
	log.Printf("sensortyps map: %+v", sensortyps)
	for _, newS := range newheader {
		//check which newS is not in sensortyps; add those to newSensortyps
		newS.Conversion_expression = ""
	}
	//Query: softdelete where Deveui
	//Query: un-delete where Deveui and sensr iotype/ioaddress/sensorstype
	//Query: insert new sensortypes
	//Query: insert new sensors with Deveui

	return nil
}

func isOffline() bool {
	return instantiated == nil
}

func checkErr(err error) {
	if err != nil {
		log.Println(err)
	}
}

func panicErr(err error) {
	if err != nil {
		log.Println(err)
		panic(err)
	}
}
<|endoftext|>