<commit_before>// Copyright 2019 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package verifier

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"

	"github.com/golang/glog"
	"github.com/golang/protobuf/proto"
	"github.com/google/trillian/types"
	"github.com/kr/pretty"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"github.com/google/keytransparency/core/crypto/commitments"
	"github.com/google/keytransparency/core/crypto/vrf"
	"github.com/google/keytransparency/core/crypto/vrf/p256"
	"github.com/google/keytransparency/core/mutator/entry"

	pb "github.com/google/keytransparency/core/api/v1/keytransparency_go_proto"
	tclient "github.com/google/trillian/client"
	_ "github.com/google/trillian/merkle/coniks"  // Register hasher
	_ "github.com/google/trillian/merkle/rfc6962" // Register hasher
)

var (
	// ErrNilProof occurs when the provided GetUserResponse contains a nil proof.
	ErrNilProof = errors.New("nil proof")
)

// LogTracker tracks a series of consistent log roots.
type LogTracker interface {
	// LastVerifiedLogRoot retrieves the tree size of the latest log root.
	LastVerifiedLogRoot() *pb.LogRootRequest
	// VerifyLogRoot verifies root and updates the trusted root if it is newer.
	VerifyLogRoot(state *pb.LogRootRequest, newRoot *pb.LogRoot) (*types.LogRootV1, error)
}

// LogTrackerFactory allows the caller of NewFromDirectory to supply different
// tracker implementations.  eg. no-op, single-synchronous-head, and multi-head tracking.
type LogTrackerFactory func(*tclient.LogVerifier) LogTracker

// Verifier is a client helper library for verifying requests and responses.
type Verifier struct {
	vrf     vrf.PublicKey
	mv      *tclient.MapVerifier
	lv      *tclient.LogVerifier
	lt      LogTracker
	verbose *log.Logger
}

// New creates a new instance of the client verifier.
func New(vrf vrf.PublicKey,
	mv *tclient.MapVerifier,
	lv *tclient.LogVerifier,
	lt LogTracker) *Verifier {
	return &Verifier{
		vrf:     vrf,
		mv:      mv,
		lv:      lv,
		lt:      lt,
		verbose: log.New(ioutil.Discard, "", 0),
	}
}

// NewFromDirectory creates a new instance of the client verifier from a config.
func NewFromDirectory(config *pb.Directory, f LogTrackerFactory) (*Verifier, error) {
	logVerifier, err := tclient.NewLogVerifierFromTree(config.GetLog())
	if err != nil {
		return nil, err
	}

	mapVerifier, err := tclient.NewMapVerifierFromTree(config.GetMap())
	if err != nil {
		return nil, err
	}

	// VRF key
	vrfPubKey, err := p256.NewVRFVerifierFromRawKey(config.GetVrf().GetDer())
	if err != nil {
		return nil, fmt.Errorf("error parsing vrf public key: %v", err)
	}

	tracker := f(logVerifier)

	return New(vrfPubKey, mapVerifier, logVerifier, tracker), nil
}

// Index computes the index from a VRF proof.
func (v *Verifier) Index(vrfProof []byte, directoryID, userID string) ([]byte, error) {
	index, err := v.vrf.ProofToHash([]byte(userID), vrfProof)
	if err != nil {
		return nil, fmt.Errorf("vrf.ProofToHash(): %v", err)
	}
	return index[:], nil
}

// VerifyMapLeaf verifies pb.MapLeaf:
//  - Verify commitment.
//  - Verify VRF and index.
//  - Verify map inclusion proof.
func (v *Verifier) VerifyMapLeaf(directoryID, userID string,
	in *pb.MapLeaf, mapRoot *types.MapRootV1) error {
	if mapRoot == nil {
		return status.Errorf(codes.Internal, "nil MapRoot")
	}
	glog.V(5).Infof("VerifyMapLeaf(%v/%v): %# v", directoryID, userID, pretty.Formatter(in))

	// Unpack the merkle tree leaf value.
	leafValue := in.GetMapInclusion().GetLeaf().GetLeafValue()
	signed, err := entry.FromLeafValue(leafValue)
	if err != nil {
		return err
	}
	var e pb.Entry
	if err := proto.Unmarshal(signed.GetEntry(), &e); err != nil {
		return err
	}

	// If this is not a proof of absence, verify the connection between
	// profileData and the commitment in the merkle tree leaf.
	if in.GetCommitted() != nil {
		commitment := e.GetCommitment()
		data := in.GetCommitted().GetData()
		nonce := in.GetCommitted().GetKey()
		if err := commitments.Verify(userID, commitment, data, nonce); err != nil {
			v.verbose.Printf("✗ Commitment verification failed.")
			return fmt.Errorf("commitments.Verify(%v, %x, %x, %v): %v", userID, commitment, data, nonce, err)
		}
	}
	v.verbose.Printf("✓ Commitment verified.")

	index, err := v.Index(in.GetVrfProof(), directoryID, userID)
	if err != nil {
		v.verbose.Printf("✗ VRF verification failed.")
		return err
	}

	if leafValue != nil && !bytes.Equal(index, e.Index) {
		v.verbose.Printf("✗ VRF verification failed.")
		return fmt.Errorf("entry has wrong index: %x, want %x", e.Index, index)
	}
	v.verbose.Printf("✓ VRF verified.")

	leafProof := in.GetMapInclusion()
	if leafProof == nil {
		return ErrNilProof
	}
	leafProof.Leaf.Index = index

	if err := v.mv.VerifyMapLeafInclusionHash(mapRoot.RootHash, leafProof); err != nil {
		v.verbose.Printf("✗ Sparse tree proof verification failed.")
		return fmt.Errorf("map inclusion proof failed: %v", err)
	}
	v.verbose.Printf("✓ map inclusion proof verified.")
	return nil
}

// LastVerifiedLogRoot returns a LogRootRequest for making an RPC
func (v *Verifier) LastVerifiedLogRoot() *pb.LogRootRequest {
	return v.lt.LastVerifiedLogRoot()
}

// VerifyLogRoot verifies that revision.LogRoot is consistent with the last trusted SignedLogRoot.
func (v *Verifier) VerifyLogRoot(req *pb.LogRootRequest, slr *pb.LogRoot) (*types.LogRootV1, error) {
	// TODO(gdbelvin): Gossip root.
	return v.lt.VerifyLogRoot(req, slr)
}

// VerifyMapRevision verifies that the map revision is correctly signed and included in the append only log.
func (v *Verifier) VerifyMapRevision(lr *types.LogRootV1, smr *pb.MapRoot) (*types.MapRootV1, error) {
	mapRoot, err := v.mv.VerifySignedMapRoot(smr.GetMapRoot())
	if err != nil {
		v.verbose.Printf("✗ Signed Map Head signature verification failed.")
		return nil, err
	}
	v.verbose.Printf("✓ Signed Map Head signature verified.")

	// Verify inclusion proof.
	b := smr.GetMapRoot().GetMapRoot()
	leafIndex := int64(mapRoot.Revision)
	if err := v.lv.VerifyInclusionAtIndex(lr, b, leafIndex, smr.GetLogInclusion()); err != nil {
		return nil, fmt.Errorf("logVerifier: VerifyInclusionAtIndex(%x, %v, _): %v", b, leafIndex, err)
	}
	v.verbose.Printf("✓ Log inclusion proof verified.")
	return mapRoot, nil
}
<commit_msg>Remove dependency on kr/pretty (#1467)<commit_after>// Copyright 2019 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package verifier

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"

	"github.com/golang/glog"
	"github.com/golang/protobuf/proto"
	"github.com/google/trillian/types"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"github.com/google/keytransparency/core/crypto/commitments"
	"github.com/google/keytransparency/core/crypto/vrf"
	"github.com/google/keytransparency/core/crypto/vrf/p256"
	"github.com/google/keytransparency/core/mutator/entry"

	pb "github.com/google/keytransparency/core/api/v1/keytransparency_go_proto"
	tclient "github.com/google/trillian/client"
	_ "github.com/google/trillian/merkle/coniks"  // Register hasher
	_ "github.com/google/trillian/merkle/rfc6962" // Register hasher
)

var (
	// ErrNilProof occurs when the provided GetUserResponse contains a nil proof.
	ErrNilProof = errors.New("nil proof")
)

// LogTracker tracks a series of consistent log roots.
type LogTracker interface {
	// LastVerifiedLogRoot retrieves the tree size of the latest log root.
	LastVerifiedLogRoot() *pb.LogRootRequest
	// VerifyLogRoot verifies root and updates the trusted root if it is newer.
	VerifyLogRoot(state *pb.LogRootRequest, newRoot *pb.LogRoot) (*types.LogRootV1, error)
}

// LogTrackerFactory allows the caller of NewFromDirectory to supply different
// tracker implementations.  eg. no-op, single-synchronous-head, and multi-head tracking.
type LogTrackerFactory func(*tclient.LogVerifier) LogTracker

// Verifier is a client helper library for verifying requests and responses.
type Verifier struct {
	vrf     vrf.PublicKey
	mv      *tclient.MapVerifier
	lv      *tclient.LogVerifier
	lt      LogTracker
	verbose *log.Logger
}

// New creates a new instance of the client verifier.
func New(vrf vrf.PublicKey,
	mv *tclient.MapVerifier,
	lv *tclient.LogVerifier,
	lt LogTracker) *Verifier {
	return &Verifier{
		vrf:     vrf,
		mv:      mv,
		lv:      lv,
		lt:      lt,
		verbose: log.New(ioutil.Discard, "", 0),
	}
}

// NewFromDirectory creates a new instance of the client verifier from a config.
func NewFromDirectory(config *pb.Directory, f LogTrackerFactory) (*Verifier, error) {
	logVerifier, err := tclient.NewLogVerifierFromTree(config.GetLog())
	if err != nil {
		return nil, err
	}

	mapVerifier, err := tclient.NewMapVerifierFromTree(config.GetMap())
	if err != nil {
		return nil, err
	}

	// VRF key
	vrfPubKey, err := p256.NewVRFVerifierFromRawKey(config.GetVrf().GetDer())
	if err != nil {
		return nil, fmt.Errorf("error parsing vrf public key: %v", err)
	}

	tracker := f(logVerifier)

	return New(vrfPubKey, mapVerifier, logVerifier, tracker), nil
}

// Index computes the index from a VRF proof.
func (v *Verifier) Index(vrfProof []byte, directoryID, userID string) ([]byte, error) {
	index, err := v.vrf.ProofToHash([]byte(userID), vrfProof)
	if err != nil {
		return nil, fmt.Errorf("vrf.ProofToHash(): %v", err)
	}
	return index[:], nil
}

// VerifyMapLeaf verifies pb.MapLeaf:
//  - Verify commitment.
//  - Verify VRF and index.
//  - Verify map inclusion proof.
func (v *Verifier) VerifyMapLeaf(directoryID, userID string,
	in *pb.MapLeaf, mapRoot *types.MapRootV1) error {
	if mapRoot == nil {
		return status.Errorf(codes.Internal, "nil MapRoot")
	}
	glog.V(5).Infof("VerifyMapLeaf(%v/%v): %# v", directoryID, userID, in)

	// Unpack the merkle tree leaf value.
	leafValue := in.GetMapInclusion().GetLeaf().GetLeafValue()
	signed, err := entry.FromLeafValue(leafValue)
	if err != nil {
		return err
	}
	var e pb.Entry
	if err := proto.Unmarshal(signed.GetEntry(), &e); err != nil {
		return err
	}

	// If this is not a proof of absence, verify the connection between
	// profileData and the commitment in the merkle tree leaf.
	if in.GetCommitted() != nil {
		commitment := e.GetCommitment()
		data := in.GetCommitted().GetData()
		nonce := in.GetCommitted().GetKey()
		if err := commitments.Verify(userID, commitment, data, nonce); err != nil {
			v.verbose.Printf("✗ Commitment verification failed.")
			return fmt.Errorf("commitments.Verify(%v, %x, %x, %v): %v", userID, commitment, data, nonce, err)
		}
	}
	v.verbose.Printf("✓ Commitment verified.")

	index, err := v.Index(in.GetVrfProof(), directoryID, userID)
	if err != nil {
		v.verbose.Printf("✗ VRF verification failed.")
		return err
	}

	if leafValue != nil && !bytes.Equal(index, e.Index) {
		v.verbose.Printf("✗ VRF verification failed.")
		return fmt.Errorf("entry has wrong index: %x, want %x", e.Index, index)
	}
	v.verbose.Printf("✓ VRF verified.")

	leafProof := in.GetMapInclusion()
	if leafProof == nil {
		return ErrNilProof
	}
	leafProof.Leaf.Index = index

	if err := v.mv.VerifyMapLeafInclusionHash(mapRoot.RootHash, leafProof); err != nil {
		v.verbose.Printf("✗ Sparse tree proof verification failed.")
		return fmt.Errorf("map inclusion proof failed: %v", err)
	}
	v.verbose.Printf("✓ map inclusion proof verified.")
	return nil
}

// LastVerifiedLogRoot returns a LogRootRequest for making an RPC
func (v *Verifier) LastVerifiedLogRoot() *pb.LogRootRequest {
	return v.lt.LastVerifiedLogRoot()
}

// VerifyLogRoot verifies that revision.LogRoot is consistent with the last trusted SignedLogRoot.
func (v *Verifier) VerifyLogRoot(req *pb.LogRootRequest, slr *pb.LogRoot) (*types.LogRootV1, error) {
	// TODO(gdbelvin): Gossip root.
	return v.lt.VerifyLogRoot(req, slr)
}

// VerifyMapRevision verifies that the map revision is correctly signed and included in the append only log.
func (v *Verifier) VerifyMapRevision(lr *types.LogRootV1, smr *pb.MapRoot) (*types.MapRootV1, error) {
	mapRoot, err := v.mv.VerifySignedMapRoot(smr.GetMapRoot())
	if err != nil {
		v.verbose.Printf("✗ Signed Map Head signature verification failed.")
		return nil, err
	}
	v.verbose.Printf("✓ Signed Map Head signature verified.")

	// Verify inclusion proof.
	b := smr.GetMapRoot().GetMapRoot()
	leafIndex := int64(mapRoot.Revision)
	if err := v.lv.VerifyInclusionAtIndex(lr, b, leafIndex, smr.GetLogInclusion()); err != nil {
		return nil, fmt.Errorf("logVerifier: VerifyInclusionAtIndex(%x, %v, _): %v", b, leafIndex, err)
	}
	v.verbose.Printf("✓ Log inclusion proof verified.")
	return mapRoot, nil
}
<|endoftext|>


<commit_before>package vmware

import (
	"fmt"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"path/filepath"
)

// This step creates the virtual disks for the VM.
//
// Uses:
//   config *config
//   driver Driver
//   ui     packer.Ui
//
// Produces:
//   <nothing>
type stepCreateDisk struct{}

func (stepCreateDisk) Run(state map[string]interface{}) multistep.StepAction {
	config := state["config"].(*config)
	driver := state["driver"].(Driver)
	ui := state["ui"].(packer.Ui)

	ui.Say("Creating virtual machine disk")
	full_disk_path := filepath.Join(config.OutputDir, config.DiskName+".vmdk")
	state["full_disk_path"] = full_disk_path
	if err := driver.CreateDisk(full_disk_path, fmt.Sprintf("%dM", config.DiskSize)); err != nil {
		err := fmt.Errorf("Error creating disk: %s", err)
		state["error"] = err
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	return multistep.ActionContinue
}

func (stepCreateDisk) Cleanup(map[string]interface{}) {}
<commit_msg>builder/vmware: Style nitpick<commit_after>package vmware

import (
	"fmt"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"path/filepath"
)

// This step creates the virtual disks for the VM.
//
// Uses:
//   config *config
//   driver Driver
//   ui     packer.Ui
//
// Produces:
//   full_disk_path (string) - The full path to the created disk.
type stepCreateDisk struct{}

func (stepCreateDisk) Run(state map[string]interface{}) multistep.StepAction {
	config := state["config"].(*config)
	driver := state["driver"].(Driver)
	ui := state["ui"].(packer.Ui)

	ui.Say("Creating virtual machine disk")
	full_disk_path := filepath.Join(config.OutputDir, config.DiskName+".vmdk")
	if err := driver.CreateDisk(full_disk_path, fmt.Sprintf("%dM", config.DiskSize)); err != nil {
		err := fmt.Errorf("Error creating disk: %s", err)
		state["error"] = err
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	state["full_disk_path"] = full_disk_path

	return multistep.ActionContinue
}

func (stepCreateDisk) Cleanup(map[string]interface{}) {}
<|endoftext|>


<commit_before>/*

Package tictactoe is an exceedingly simple game based on boardgame. It serves
as an example, and also helps verify that the design and implementation of
boardgame are useful for real games.

*/
package tictactoe

import (
	"github.com/jkomoros/boardgame"
)

const gameName = "tictactoe"

const DIM = 3

func ticTacToeGame() *boardgame.Game {

	chest := boardgame.NewComponentChest(gameName)

	tokens := &boardgame.Deck{}

	//How many tokens of each of x's and o's do we need so that no matter who
	//goes first we always have enough?
	numTokens := 5

	//TODO: use deck.AddComponentMulti when that exists
	for i := 0; i < numTokens; i++ {
		tokens.AddComponent(&boardgame.Component{
			Values: &playerToken{
				Value: X,
			},
		})
	}

	for i := 0; i < numTokens; i++ {
		tokens.AddComponent(&boardgame.Component{
			Values: &playerToken{
				Value: O,
			},
		})
	}

	chest.AddDeck("tokens", tokens)

	starterPayload := &statePayload{
		game: &gameState{
			Slots: boardgame.NewSizedStack(tokens, DIM*DIM),
		},
	}

	game := &boardgame.Game{
		Name:  gameName,
		State: boardgame.NewStarterState(starterPayload),
	}

	game.SetChest(chest)

	return game

}
<commit_msg>Initialize user states in the created game. Part of #14.<commit_after>/*

Package tictactoe is an exceedingly simple game based on boardgame. It serves
as an example, and also helps verify that the design and implementation of
boardgame are useful for real games.

*/
package tictactoe

import (
	"github.com/jkomoros/boardgame"
)

const gameName = "tictactoe"

const DIM = 3

func ticTacToeGame() *boardgame.Game {

	chest := boardgame.NewComponentChest(gameName)

	tokens := &boardgame.Deck{}

	//How many tokens of each of x's and o's do we need so that no matter who
	//goes first we always have enough?
	numTokens := 5

	//TODO: use deck.AddComponentMulti when that exists
	for i := 0; i < numTokens; i++ {
		tokens.AddComponent(&boardgame.Component{
			Values: &playerToken{
				Value: X,
			},
		})
	}

	for i := 0; i < numTokens; i++ {
		tokens.AddComponent(&boardgame.Component{
			Values: &playerToken{
				Value: O,
			},
		})
	}

	chest.AddDeck("tokens", tokens)

	starterPayload := &statePayload{
		game: &gameState{
			Slots: boardgame.NewSizedStack(tokens, DIM*DIM),
		},
		users: []*userState{
			&userState{
				UnusedTokens: boardgame.NewGrowableStack(tokens, 0),
			},
			&userState{
				UnusedTokens: boardgame.NewGrowableStack(tokens, 0),
			},
		},
	}

	for i, user := range starterPayload.users {
		user.playerIndex = i
	}

	game := &boardgame.Game{
		Name:  gameName,
		State: boardgame.NewStarterState(starterPayload),
	}

	game.SetChest(chest)

	return game

}
<|endoftext|>


<commit_before>package main

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/apache/incubator-trafficcontrol/lib/go-log"
	"github.com/apache/incubator-trafficcontrol/lib/go-tc"
	"github.com/basho/riak-go-client"
	"github.com/jmoiron/sqlx"
	"github.com/lestrrat/go-jwx/jwk"
	"io/ioutil"
	"net/http"
	"strings"
)

const RiakPort = 8087
const CdnURIKeysBucket = "cdn_uri_sig_keys" // riak namespace for cdn uri signing keys.

type URISignerKeyset struct {
	RenewalKid *string               `json:"renewal_kid"`
	Keys       []jwk.EssentialHeader `json:"keys"`
}

// deletes an object from riak storage
func deleteObject(key string, bucket string, cluster *riak.Cluster) error {
	// build store command and execute.
	cmd, err := riak.NewDeleteValueCommandBuilder().
		WithBucket(bucket).
		WithKey(key).
		Build()
	if err != nil {
		return err
	}
	if err := cluster.Execute(cmd); err != nil {
		return err
	}

	return nil
}

// fetch an object from riak storage
func fetchObjectValues(key string, bucket string, cluster *riak.Cluster) ([]*riak.Object, error) {
	// build the fetch command
	cmd, err := riak.NewFetchValueCommandBuilder().
		WithBucket(bucket).
		WithKey(key).
		Build()
	if err != nil {
		return nil, err
	}

	if err = cluster.Execute(cmd); err != nil {
		return nil, err
	}
	fvc := cmd.(*riak.FetchValueCommand)

	// no object found with given key
	if fvc.Response == nil || fvc.Response.IsNotFound {
		return nil, nil
	}
	return fvc.Response.Values, nil
}

// saves an object to riak storage
func saveObject(obj *riak.Object, bucket string, cluster *riak.Cluster) error {
	// build store command and execute.
	cmd, err := riak.NewStoreValueCommandBuilder().
		WithBucket(bucket).
		WithContent(obj).
		Build()
	if err != nil {
		return err
	}
	if err := cluster.Execute(cmd); err != nil {
		return err
	}

	return nil
}

// returns a riak cluster of online riak nodes.
func getRiakCluster(db *sqlx.DB, cfg Config) (*riak.Cluster, error) {
	riakServerQuery := `
		SELECT s.host_name, s.domain_name FROM server s
		INNER JOIN type t on s.type = t.id
		INNER JOIN status st on s.status = st.id
		WHERE t.name = 'RIAK' AND st.name = 'ONLINE'
		`

	if cfg.RiakAuthOptions == nil {
		return nil, errors.New("ERROR: no riak auth information from riak.conf, cannot authenticate to any riak servers.")
	}

	var nodes []*riak.Node
	rows, err := db.Query(riakServerQuery)
	if err != nil {
		return nil, err
	}

	for rows.Next() {
		var s tc.Server
		var n *riak.Node
		if err := rows.Scan(&s.HostName, &s.DomainName); err != nil {
			return nil, err
		}
		addr := fmt.Sprintf("%s.%s:%d", s.HostName, s.DomainName, RiakPort)
		nodeOpts := &riak.NodeOptions{
			RemoteAddress: addr,
			AuthOptions:   cfg.RiakAuthOptions,
		}
		nodeOpts.AuthOptions.TlsConfig.ServerName = fmt.Sprintf("%s.%s", s.HostName, s.DomainName)
		n, err := riak.NewNode(nodeOpts)
		if err != nil {
			return nil, err
		}
		nodes = append(nodes, n)
	}

	if len(nodes) == 0 {
		return nil, errors.New("ERROR: no available riak servers.")
	}

	opts := &riak.ClusterOptions{
		Nodes: nodes,
	}
	cluster, err := riak.NewCluster(opts)

	return cluster, err
}

// endpoint handler for fetching uri signing keys from riak
func getURIsignkeysHandler(db *sqlx.DB, cfg Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		handleErr := tc.GetHandleErrorFunc(w, r)

		ctx := r.Context()
		pathParams, err := getPathParams(ctx)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		xmlId := pathParams["xmlId"]

		// create and start a cluster
		cluster, err := getRiakCluster(db, cfg)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		if err = cluster.Start(); err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		defer func() {
			if err := cluster.Stop(); err != nil {
				log.Errorf("%v\n", err)
			}
		}()

		ro, err := fetchObjectValues(xmlId, CdnURIKeysBucket, cluster)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		var respBytes []byte

		if ro == nil {
			var empty URISignerKeyset
			respBytes, err = json.Marshal(empty)
			if err != nil {
				log.Errorf("failed to marshal an empty response: %s\n", err)
				w.WriteHeader(http.StatusInternalServerError)
				fmt.Fprintf(w, http.StatusText(http.StatusInternalServerError))
				return
			}
		} else {
			respBytes = ro[0].Value
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, "%s", respBytes)
	}
}

// Http POST handler used to store urisigning keys to a delivery service.
func assignDeliveryServiceURIKeysHandler(db *sqlx.DB, cfg Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		handleErr := tc.GetHandleErrorFunc(w, r)

		defer r.Body.Close()

		ctx := r.Context()
		pathParams, err := getPathParams(ctx)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		xmlId := pathParams["xmlId"]
		data, err := ioutil.ReadAll(r.Body)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		// validate that the received data is a valid jwk keyset
		var keySet map[string]URISignerKeyset
		if err := json.Unmarshal(data, &keySet); err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusBadRequest)
			return
		}
		if err := validateURIKeyset(keySet); err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusBadRequest)
			return
		}

		// create and start a cluster
		cluster, err := getRiakCluster(db, cfg)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		if err = cluster.Start(); err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		defer func() {
			if err := cluster.Stop(); err != nil {
				log.Errorf("%v\n", err)
			}
		}()

		ro, err := fetchObjectValues(xmlId, CdnURIKeysBucket, cluster)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		// object exists.
		if ro != nil && ro[0].Value != nil {
			handleErr(fmt.Errorf("a keyset already exists for this delivery service."), http.StatusBadRequest)
			return
		}

		// create a storage object and store the data
		obj := &riak.Object{
			ContentType:     "text/json",
			Charset:         "utf-8",
			ContentEncoding: "utf-8",
			Key:             xmlId,
			Value:           []byte(data),
		}

		err = saveObject(obj, CdnURIKeysBucket, cluster)
		if err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, "%s", data)
	}
}

// Http DELETE handler used to remove urisigning keys assigned to a delivery service.
func removeDeliveryServiceURIKeysHandler(db *sqlx.DB, cfg Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		handleErr := tc.GetHandleErrorFunc(w, r)

		ctx := r.Context()
		pathParams, err := getPathParams(ctx)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		xmlId := pathParams["xmlId"]

		// create and start a cluster
		cluster, err := getRiakCluster(db, cfg)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		if err = cluster.Start(); err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		defer func() {
			if err := cluster.Stop(); err != nil {
				log.Errorf("%v\n", err)
			}
		}()

		ro, err := fetchObjectValues(xmlId, CdnURIKeysBucket, cluster)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		// fetch the object and delete it if it exists.
		var alert tc.Alerts

		if ro == nil || ro[0].Value == nil {
			alert = tc.CreateAlerts(tc.InfoLevel, "not deleted, no object found to delete.")
		} else if err := deleteObject(xmlId, CdnURIKeysBucket, cluster); err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		} else { // object successfully deleted
			alert = tc.CreateAlerts(tc.SuccessLevel, "object deleted")
		}

		// send response
		respBytes, err := json.Marshal(alert)
		if err != nil {
			log.Errorf("failed to marshal an alert response: %s\n", err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, http.StatusText(http.StatusInternalServerError))
			return
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, "%s", respBytes)
	}
}

// Http POST handler used to store urisigning keys to a delivery service.
func updateDeliveryServiceURIKeysHandler(db *sqlx.DB, cfg Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		handleErr := tc.GetHandleErrorFunc(w, r)

		defer r.Body.Close()

		ctx := r.Context()
		pathParams, err := getPathParams(ctx)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		xmlId := pathParams["xmlId"]
		data, err := ioutil.ReadAll(r.Body)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		// validate that the received data is a valid jwk keyset
		var keySet map[string]URISignerKeyset
		if err := json.Unmarshal(data, &keySet); err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusBadRequest)
			return
		}
		if err := validateURIKeyset(keySet); err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusBadRequest)
			return
		}

		// create and start a cluster
		cluster, err := getRiakCluster(db, cfg)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		if err = cluster.Start(); err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		defer func() {
			if err := cluster.Stop(); err != nil {
				log.Errorf("%v\n", err)
			}
		}()

		// create a storage object and store the data
		obj := &riak.Object{
			ContentType:     "text/json",
			Charset:         "utf-8",
			ContentEncoding: "utf-8",
			Key:             xmlId,
			Value:           []byte(data),
		}

		err = saveObject(obj, CdnURIKeysBucket, cluster)
		if err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, "%s", data)
	}
}

// validates URISigingKeyset json.
func validateURIKeyset(msg map[string]URISignerKeyset) error {
	var renewalKidFound int = 0
	var renewalKidMatched = false

	for key, value := range msg {
		issuer := key
		renewalKid := value.RenewalKid
		if issuer == "" {
			return errors.New("JSON Keyset has no issuer")
		}

		if renewalKid != nil {
			renewalKidFound++
		}

		for _, skey := range value.Keys {
			if skey.Algorithm == "" {
				return errors.New("A Key has no algorithm, alg, specified.\n")
			}
			if skey.KeyID == "" {
				return errors.New("A Key has no key id, kid, specified.\n")
			}
			if renewalKid != nil && strings.Compare(*renewalKid, skey.KeyID) == 0 {
				renewalKidMatched = true
			}
		}
	}

	// should only have one renewal_kid
	switch renewalKidFound {
	case 0:
		return errors.New("No renewal_kid was found in any keyset\n")
	case 1: // okay, this is what we want
		break
	default:
		return errors.New("More than one renewal_kid was found in the keysets\n")
	}

	// the renewal_kid should match the kid of one key
	if !renewalKidMatched {
		return errors.New("No key was found with a kid that matches the renewal kid\n")
	}

	return nil
}
<commit_msg>changed variable name from CdnUriKeysBucket to CDNURIKeysBucket as CDN is an acronym.<commit_after>package main

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/apache/incubator-trafficcontrol/lib/go-log"
	"github.com/apache/incubator-trafficcontrol/lib/go-tc"
	"github.com/basho/riak-go-client"
	"github.com/jmoiron/sqlx"
	"github.com/lestrrat/go-jwx/jwk"
	"io/ioutil"
	"net/http"
	"strings"
)

const RiakPort = 8087
const CDNURIKeysBucket = "cdn_uri_sig_keys" // riak namespace for cdn uri signing keys.

type URISignerKeyset struct {
	RenewalKid *string               `json:"renewal_kid"`
	Keys       []jwk.EssentialHeader `json:"keys"`
}

// deletes an object from riak storage
func deleteObject(key string, bucket string, cluster *riak.Cluster) error {
	// build store command and execute.
	cmd, err := riak.NewDeleteValueCommandBuilder().
		WithBucket(bucket).
		WithKey(key).
		Build()
	if err != nil {
		return err
	}
	if err := cluster.Execute(cmd); err != nil {
		return err
	}

	return nil
}

// fetch an object from riak storage
func fetchObjectValues(key string, bucket string, cluster *riak.Cluster) ([]*riak.Object, error) {
	// build the fetch command
	cmd, err := riak.NewFetchValueCommandBuilder().
		WithBucket(bucket).
		WithKey(key).
		Build()
	if err != nil {
		return nil, err
	}

	if err = cluster.Execute(cmd); err != nil {
		return nil, err
	}
	fvc := cmd.(*riak.FetchValueCommand)

	// no object found with given key
	if fvc.Response == nil || fvc.Response.IsNotFound {
		return nil, nil
	}
	return fvc.Response.Values, nil
}

// saves an object to riak storage
func saveObject(obj *riak.Object, bucket string, cluster *riak.Cluster) error {
	// build store command and execute.
	cmd, err := riak.NewStoreValueCommandBuilder().
		WithBucket(bucket).
		WithContent(obj).
		Build()
	if err != nil {
		return err
	}
	if err := cluster.Execute(cmd); err != nil {
		return err
	}

	return nil
}

// returns a riak cluster of online riak nodes.
func getRiakCluster(db *sqlx.DB, cfg Config) (*riak.Cluster, error) {
	riakServerQuery := `
		SELECT s.host_name, s.domain_name FROM server s
		INNER JOIN type t on s.type = t.id
		INNER JOIN status st on s.status = st.id
		WHERE t.name = 'RIAK' AND st.name = 'ONLINE'
		`

	if cfg.RiakAuthOptions == nil {
		return nil, errors.New("ERROR: no riak auth information from riak.conf, cannot authenticate to any riak servers.")
	}

	var nodes []*riak.Node
	rows, err := db.Query(riakServerQuery)
	if err != nil {
		return nil, err
	}

	for rows.Next() {
		var s tc.Server
		var n *riak.Node
		if err := rows.Scan(&s.HostName, &s.DomainName); err != nil {
			return nil, err
		}
		addr := fmt.Sprintf("%s.%s:%d", s.HostName, s.DomainName, RiakPort)
		nodeOpts := &riak.NodeOptions{
			RemoteAddress: addr,
			AuthOptions:   cfg.RiakAuthOptions,
		}
		nodeOpts.AuthOptions.TlsConfig.ServerName = fmt.Sprintf("%s.%s", s.HostName, s.DomainName)
		n, err := riak.NewNode(nodeOpts)
		if err != nil {
			return nil, err
		}
		nodes = append(nodes, n)
	}

	if len(nodes) == 0 {
		return nil, errors.New("ERROR: no available riak servers.")
	}

	opts := &riak.ClusterOptions{
		Nodes: nodes,
	}
	cluster, err := riak.NewCluster(opts)

	return cluster, err
}

// endpoint handler for fetching uri signing keys from riak
func getURIsignkeysHandler(db *sqlx.DB, cfg Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		handleErr := tc.GetHandleErrorFunc(w, r)

		ctx := r.Context()
		pathParams, err := getPathParams(ctx)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		xmlId := pathParams["xmlId"]

		// create and start a cluster
		cluster, err := getRiakCluster(db, cfg)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		if err = cluster.Start(); err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		defer func() {
			if err := cluster.Stop(); err != nil {
				log.Errorf("%v\n", err)
			}
		}()

		ro, err := fetchObjectValues(xmlId, CDNURIKeysBucket, cluster)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		var respBytes []byte

		if ro == nil {
			var empty URISignerKeyset
			respBytes, err = json.Marshal(empty)
			if err != nil {
				log.Errorf("failed to marshal an empty response: %s\n", err)
				w.WriteHeader(http.StatusInternalServerError)
				fmt.Fprintf(w, http.StatusText(http.StatusInternalServerError))
				return
			}
		} else {
			respBytes = ro[0].Value
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, "%s", respBytes)
	}
}

// Http POST handler used to store urisigning keys to a delivery service.
func assignDeliveryServiceURIKeysHandler(db *sqlx.DB, cfg Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		handleErr := tc.GetHandleErrorFunc(w, r)

		defer r.Body.Close()

		ctx := r.Context()
		pathParams, err := getPathParams(ctx)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		xmlId := pathParams["xmlId"]
		data, err := ioutil.ReadAll(r.Body)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		// validate that the received data is a valid jwk keyset
		var keySet map[string]URISignerKeyset
		if err := json.Unmarshal(data, &keySet); err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusBadRequest)
			return
		}
		if err := validateURIKeyset(keySet); err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusBadRequest)
			return
		}

		// create and start a cluster
		cluster, err := getRiakCluster(db, cfg)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		if err = cluster.Start(); err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		defer func() {
			if err := cluster.Stop(); err != nil {
				log.Errorf("%v\n", err)
			}
		}()

		ro, err := fetchObjectValues(xmlId, CDNURIKeysBucket, cluster)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		// object exists.
		if ro != nil && ro[0].Value != nil {
			handleErr(fmt.Errorf("a keyset already exists for this delivery service."), http.StatusBadRequest)
			return
		}

		// create a storage object and store the data
		obj := &riak.Object{
			ContentType:     "text/json",
			Charset:         "utf-8",
			ContentEncoding: "utf-8",
			Key:             xmlId,
			Value:           []byte(data),
		}

		err = saveObject(obj, CDNURIKeysBucket, cluster)
		if err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, "%s", data)
	}
}

// Http DELETE handler used to remove urisigning keys assigned to a delivery service.
func removeDeliveryServiceURIKeysHandler(db *sqlx.DB, cfg Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		handleErr := tc.GetHandleErrorFunc(w, r)

		ctx := r.Context()
		pathParams, err := getPathParams(ctx)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		xmlId := pathParams["xmlId"]

		// create and start a cluster
		cluster, err := getRiakCluster(db, cfg)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		if err = cluster.Start(); err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		defer func() {
			if err := cluster.Stop(); err != nil {
				log.Errorf("%v\n", err)
			}
		}()

		ro, err := fetchObjectValues(xmlId, CDNURIKeysBucket, cluster)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		// fetch the object and delete it if it exists.
		var alert tc.Alerts

		if ro == nil || ro[0].Value == nil {
			alert = tc.CreateAlerts(tc.InfoLevel, "not deleted, no object found to delete.")
		} else if err := deleteObject(xmlId, CDNURIKeysBucket, cluster); err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		} else { // object successfully deleted
			alert = tc.CreateAlerts(tc.SuccessLevel, "object deleted")
		}

		// send response
		respBytes, err := json.Marshal(alert)
		if err != nil {
			log.Errorf("failed to marshal an alert response: %s\n", err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, http.StatusText(http.StatusInternalServerError))
			return
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, "%s", respBytes)
	}
}

// Http POST handler used to store urisigning keys to a delivery service.
func updateDeliveryServiceURIKeysHandler(db *sqlx.DB, cfg Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		handleErr := tc.GetHandleErrorFunc(w, r)

		defer r.Body.Close()

		ctx := r.Context()
		pathParams, err := getPathParams(ctx)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		xmlId := pathParams["xmlId"]
		data, err := ioutil.ReadAll(r.Body)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}

		// validate that the received data is a valid jwk keyset
		var keySet map[string]URISignerKeyset
		if err := json.Unmarshal(data, &keySet); err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusBadRequest)
			return
		}
		if err := validateURIKeyset(keySet); err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusBadRequest)
			return
		}

		// create and start a cluster
		cluster, err := getRiakCluster(db, cfg)
		if err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		if err = cluster.Start(); err != nil {
			handleErr(err, http.StatusInternalServerError)
			return
		}
		defer func() {
			if err := cluster.Stop(); err != nil {
				log.Errorf("%v\n", err)
			}
		}()

		// create a storage object and store the data
		obj := &riak.Object{
			ContentType:     "text/json",
			Charset:         "utf-8",
			ContentEncoding: "utf-8",
			Key:             xmlId,
			Value:           []byte(data),
		}

		err = saveObject(obj, CDNURIKeysBucket, cluster)
		if err != nil {
			log.Errorf("%v\n", err)
			handleErr(err, http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, "%s", data)
	}
}

// validates URISigingKeyset json.
func validateURIKeyset(msg map[string]URISignerKeyset) error {
	var renewalKidFound int = 0
	var renewalKidMatched = false

	for key, value := range msg {
		issuer := key
		renewalKid := value.RenewalKid
		if issuer == "" {
			return errors.New("JSON Keyset has no issuer")
		}

		if renewalKid != nil {
			renewalKidFound++
		}

		for _, skey := range value.Keys {
			if skey.Algorithm == "" {
				return errors.New("A Key has no algorithm, alg, specified.\n")
			}
			if skey.KeyID == "" {
				return errors.New("A Key has no key id, kid, specified.\n")
			}
			if renewalKid != nil && strings.Compare(*renewalKid, skey.KeyID) == 0 {
				renewalKidMatched = true
			}
		}
	}

	// should only have one renewal_kid
	switch renewalKidFound {
	case 0:
		return errors.New("No renewal_kid was found in any keyset\n")
	case 1: // okay, this is what we want
		break
	default:
		return errors.New("More than one renewal_kid was found in the keysets\n")
	}

	// the renewal_kid should match the kid of one key
	if !renewalKidMatched {
		return errors.New("No key was found with a kid that matches the renewal kid\n")
	}

	return nil
}
<|endoftext|>


<commit_before>package sms // import "github.com/mndrix/sms-over-xmpp"

import (
	"fmt"
	"strings"

	"github.com/sheenobu/go-xco"
)

func Main(config Config) {
	opts := xco.Options{
		Name:         config.ComponentName(),
		SharedSecret: config.SharedSecret(),
		Address:      fmt.Sprintf("%s:%d", config.XmppHost(), config.XmppPort()),
	}
	c, err := xco.NewComponent(opts)
	if err != nil {
		panic(err)
	}

	// Uppercase Echo Component
	c.MessageHandler = xco.BodyResponseHandler(func(msg *xco.Message) (string, error) {
		return strings.ToUpper(msg.Body), nil
	})

	err = c.Run()
	if err != nil {
		log.Printf("ERROR: Run: %s", err)
	}
}
<commit_msg>Log for XMPP events<commit_after>package sms // import "github.com/mndrix/sms-over-xmpp"

import (
	"encoding/xml"
	"fmt"
	"log"
	"strings"

	"github.com/pkg/errors"
	"github.com/sheenobu/go-xco"
)

func Main(config Config) {
	opts := xco.Options{
		Name:         config.ComponentName(),
		SharedSecret: config.SharedSecret(),
		Address:      fmt.Sprintf("%s:%d", config.XmppHost(), config.XmppPort()),
	}
	c, err := xco.NewComponent(opts)
	if err != nil {
		panic(err)
	}

	// Uppercase Echo Component
	c.MessageHandler = func(c *xco.Component, m *xco.Message) error {
		log.Printf("Message: %+v", m)
		if m.Body == "" {
			log.Printf("  ignoring message with empty body")
			return nil
		}
		resp := &xco.Message{
			Header: xco.Header{
				From: m.To,
				To:   m.From,
				ID:   m.ID,
			},
			Subject: m.Subject,
			Thread:  m.Thread,
			Type:    m.Type,
			Body:    strings.ToUpper(m.Body),
			XMLName: m.XMLName,
		}
		log.Printf("Responding: %+v", resp)
		return errors.Wrap(c.Send(resp), "sending response")
	}
	c.PresenceHandler = func(c *xco.Component, p *xco.Presence) error {
		log.Printf("Presence: %+v", p)
		return nil
	}
	c.IqHandler = func(c *xco.Component, iq *xco.Iq) error {
		log.Printf("Iq: %+v", iq)
		return nil
	}
	c.UnknownHandler = func(c *xco.Component, x *xml.StartElement) error {
		log.Printf("Unknown: %+v", x)
		return nil
	}

	err = c.Run()
	if err != nil {
		log.Printf("ERROR: Run: %s", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/brotherlogic/godiscogs"
	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"

	pb "github.com/brotherlogic/discogssyncer/server"
	pbd "github.com/brotherlogic/godiscogs"
)

// GetRelease Gets the release and metadata for the release
func (syncer *Syncer) GetRelease(id int, folder int) (*pbd.Release, *pb.ReleaseMetadata) {
	releaseData, err := ioutil.ReadFile(syncer.saveLocation + "/" + strconv.Itoa(folder) + "/" + strconv.Itoa(id) + ".release")
	if err != nil {
		return nil, nil
	}

	metadataData, _ := ioutil.ReadFile(syncer.saveLocation + "/static-metadata/" + strconv.Itoa(id) + ".metadata")
	release := &pbd.Release{}
	metadata := &pb.ReleaseMetadata{}

	proto.Unmarshal(releaseData, release)
	proto.Unmarshal(metadataData, metadata)

	return release, metadata
}

// GetMonthlySpend gets the monthly spend
func (syncer *Syncer) GetMonthlySpend(ctx context.Context, req *pb.SpendRequest) (*pb.SpendResponse, error) {
	spend := 0
	for _, rel := range syncer.relMap {
		_, metadata := syncer.GetRelease(int(rel.Id), int(rel.FolderId))
		datev := time.Unix(metadata.DateAdded, 0)
		if datev.Year() == int(req.Year) && int32(datev.Month()) == req.Month {
			spend += int(metadata.Cost)
		}
	}

	return &pb.SpendResponse{TotalSpend: int32(spend)}, nil
}

func (syncer *Syncer) saveMetadata(rel *godiscogs.Release) {
	metadataRoot := syncer.saveLocation + "/static-metadata/"
	metadataPath := metadataRoot + strconv.Itoa(int(rel.Id)) + ".metadata"
	if _, err := os.Stat(metadataRoot); os.IsNotExist(err) {
		os.MkdirAll(metadataRoot, 0777)
	}

	metadata := &pb.ReleaseMetadata{}
	if _, err := os.Stat(metadataPath); os.IsNotExist(err) {
		metadata.DateAdded = time.Now().Unix()
		metadata.DateRefreshed = time.Now().Unix()
	} else {
		data, _ := ioutil.ReadFile(metadataPath)
		proto.Unmarshal(data, metadata)
		metadata.DateRefreshed = time.Now().Unix()
	}
	data, _ := proto.Marshal(metadata)
	ioutil.WriteFile(metadataPath, data, 0644)
}

func (syncer *Syncer) saveRelease(rel *godiscogs.Release, folder int) {
	//Check that the save folder exists
	savePath := syncer.saveLocation + "/" + strconv.Itoa(folder) + "/"
	if _, err := os.Stat(savePath); os.IsNotExist(err) {
		os.MkdirAll(savePath, 0777)
	}

	data, _ := proto.Marshal(rel)
	ioutil.WriteFile(savePath+strconv.Itoa(int(rel.Id))+".release", data, 0644)
	syncer.saveMetadata(rel)
}

type saver interface {
	GetCollection() []godiscogs.Release
	GetFolders() []godiscogs.Folder
	GetRelease(id int) (godiscogs.Release, error)
	MoveToFolder(folderID int, releaseID int, instanceID int, newFolderID int)
	AddToFolder(folderID int, releaseID int)
	SetRating(folderID int, releaseID int, instanceID int, rating int)
	GetWantlist() ([]pbd.Release, error)
	RemoveFromWantlist(releaseID int)
	AddToWantlist(releaseID int)
}

// EditWant edits a want in the wantlist
func (syncer *Syncer) EditWant(ctx context.Context, wantIn *pb.Want) (*pb.Want, error) {
	for _, want := range syncer.wants.Want {
		if want.ReleaseId == wantIn.ReleaseId {
			want.Valued = wantIn.Valued
		}
	}
	syncer.saveWantList()

	return wantIn, nil
}

// SaveCollection writes out the full collection to files.
func (syncer *Syncer) SaveCollection(retr saver) {
	releases := retr.GetCollection()
	for _, release := range releases {
		fullRelease, _ := retr.GetRelease(int(release.Id))
		fullRelease.InstanceId = release.InstanceId
		fullRelease.FolderId = release.FolderId
		fullRelease.Rating = release.Rating
		syncer.saveRelease(&fullRelease, int(release.FolderId))
	}
	folders := retr.GetFolders()
	folderList := pb.FolderList{}
	for i := range folders {
		folder := folders[i]
		folderList.Folders = append(folderList.Folders, &folder)
	}
	syncer.SaveFolders(&folderList)
}

// SyncWantlist syncs the wantlist with the server
func (syncer *Syncer) SyncWantlist() {
	wants, _ := syncer.retr.GetWantlist()

	for _, want := range wants {
		seen := false
		var val *pb.Want
		for _, swant := range syncer.wants.Want {
			if swant.ReleaseId == want.Id {
				seen = true
				val = swant
			}
		}

		if seen {
			val.Wanted = true
		} else {
			syncer.wants.Want = append(syncer.wants.Want, &pb.Want{ReleaseId: want.Id, Valued: false, Wanted: true})
		}
	}

	// Cache the want list releases
	for _, want := range syncer.wants.Want {
		release, _ := syncer.retr.GetRelease(int(want.ReleaseId))
		syncer.saveRelease(&release, -5)
	}

	syncer.saveWantList()
}

func (syncer *Syncer) getFolders() *pb.FolderList {
	data, _ := ioutil.ReadFile(syncer.saveLocation + "/metadata/folders")
	folderData := &pb.FolderList{}
	proto.Unmarshal(data, folderData)
	return folderData
}

// GetSingleRelease gets a single release
func (syncer *Syncer) GetSingleRelease(ctx context.Context, in *pbd.Release) (*pbd.Release, error) {
	if val, ok := syncer.relMap[in.Id]; ok {
		return val, nil
	}
	return nil, errors.New("Unable to find release")
}

// CollapseWantlist collapses the wantlist
func (syncer *Syncer) CollapseWantlist(ctx context.Context, in *pb.Empty) (*pb.Wantlist, error) {
	for _, want := range syncer.wants.Want {
		if !want.Valued {
			log.Printf("AVOIDING %v", want)
			syncer.retr.RemoveFromWantlist(int(want.ReleaseId))
			want.Wanted = false
		}
	}

	return &syncer.wants, nil
}

// RebuildWantlist rebuilds the wantlist
func (syncer *Syncer) RebuildWantlist(ctx context.Context, in *pb.Empty) (*pb.Wantlist, error) {
	for _, want := range syncer.wants.Want {
		syncer.retr.AddToWantlist(int(want.ReleaseId))
		want.Wanted = true
	}

	return &syncer.wants, nil
}

// AddToFolder adds a release to the specified folder
func (syncer *Syncer) AddToFolder(ctx context.Context, in *pb.ReleaseMove) (*pb.Empty, error) {
	syncer.retr.AddToFolder(int(in.NewFolderId), int(in.Release.Id))
	fullRelease, _ := syncer.retr.GetRelease(int(in.Release.Id))
	fullRelease.FolderId = int32(in.NewFolderId)
	syncer.saveRelease(&fullRelease, int(in.NewFolderId))
	syncer.relMap[in.Release.Id] = &fullRelease
	return &pb.Empty{}, nil
}

// UpdateRating updates the rating of a release
func (syncer *Syncer) UpdateRating(ctx context.Context, in *pbd.Release) (*pb.Empty, error) {
	syncer.retr.SetRating(int(in.FolderId), int(in.Id), int(in.InstanceId), int(in.Rating))
	fullRelease, _ := syncer.GetRelease(int(in.Id), int(in.FolderId))
	fullRelease.Rating = int32(in.Rating)
	syncer.relMap[in.Id] = fullRelease
	syncer.saveRelease(fullRelease, int(fullRelease.FolderId))
	return &pb.Empty{}, nil
}

// UpdateMetadata updates the metadata of a given record
func (syncer *Syncer) UpdateMetadata(ctx context.Context, in *pb.MetadataUpdate) (*pb.ReleaseMetadata, error) {
	release, metadata := syncer.GetRelease(int(in.Release.Id), int(in.Release.FolderId))
	if release == nil {
		return nil, errors.New("Unable to locate release")
	}
	proto.Merge(metadata, in.Update)

	metadataRoot := syncer.saveLocation + "/static-metadata/"
	metadataPath := metadataRoot + strconv.Itoa(int(in.Release.Id)) + ".metadata"
	data, _ := proto.Marshal(metadata)
	ioutil.WriteFile(metadataPath, data, 0644)

	return metadata, nil
}

// GetWantlist gets the wantlist
func (syncer *Syncer) GetWantlist(ctx context.Context, in *pb.Empty) (*pb.Wantlist, error) {
	return &syncer.wants, nil
}

// GetMetadata gets the metadata for a given release
func (syncer *Syncer) GetMetadata(ctx context.Context, in *pbd.Release) (*pb.ReleaseMetadata, error) {
	_, metadata := syncer.GetRelease(int(in.Id), int(in.FolderId))
	return metadata, nil
}

// GetReleasesInFolder serves up the releases in a given folder
func (syncer *Syncer) GetReleasesInFolder(ctx context.Context, in *pb.FolderList) (*pb.ReleaseList, error) {

	releases := pb.ReleaseList{}
	for _, folderSpec := range in.Folders {
		folders := syncer.getFolders()
		for _, folder := range folders.Folders {
			if folder.Name == folderSpec.Name || folder.Id == folderSpec.Id {
				innerReleases := syncer.getReleases(int(folder.Id))
				releases.Releases = append(releases.Releases, innerReleases.Releases...)
			}
		}
	}

	return &releases, nil
}

func (syncer *Syncer) getReleases(folderID int) *pb.ReleaseList {
	releases := pb.ReleaseList{}
	files, _ := ioutil.ReadDir(syncer.saveLocation + "/" + strconv.Itoa(folderID) + "/")
	for _, file := range files {
		if strings.HasSuffix(file.Name(), ".release") {
			data, _ := ioutil.ReadFile(syncer.saveLocation + "/" + strconv.Itoa(folderID) + "/" + file.Name())
			release := &pbd.Release{}
			proto.Unmarshal(data, release)
			releases.Releases = append(releases.Releases, release)
		}
	}
	return &releases
}

func (syncer *Syncer) saveWantList() {
	data, _ := proto.Marshal(&syncer.wants)
	savePath := syncer.saveLocation + "/metadata/"
	if _, err := os.Stat(savePath); os.IsNotExist(err) {
		os.MkdirAll(savePath, 0777)
	}
	ioutil.WriteFile(savePath+"wantlist", data, 0644)
}

// SaveFolders saves out the list of folders
func (syncer *Syncer) SaveFolders(list *pb.FolderList) {
	savePath := syncer.saveLocation + "/metadata/"
	if _, err := os.Stat(savePath); os.IsNotExist(err) {
		os.MkdirAll(savePath, 0777)
	}

	data, _ := proto.Marshal(list)
	ioutil.WriteFile(savePath+"folders", data, 0644)
}

// GetCollection serves up the whole of the collection
func (syncer *Syncer) GetCollection(ctx context.Context, in *pb.Empty) (*pb.ReleaseList, error) {
	releases := pb.ReleaseList{}
	bfiles, _ := ioutil.ReadDir(syncer.saveLocation)
	for _, bfile := range bfiles {
		if bfile.IsDir() && bfile.Name() != "-5" {
			folderID, _ := strconv.Atoi(bfile.Name())
			for _, release := range syncer.getReleases(folderID).Releases {
				releases.Releases = append(releases.Releases, release)
			}
		}
	}
	return &releases, nil
}

// DeleteWant removes a want from the system
func (syncer *Syncer) DeleteWant(ctx context.Context, in *pb.Want) (*pb.Empty, error) {
	//Remove the want file and remove from
	index := -1
	for i, val := range syncer.wants.Want {
		if val.ReleaseId == in.ReleaseId {
			index = i
		}
	}

	if index >= 0 {
		syncer.wants.Want = append(syncer.wants.Want[:index], syncer.wants.Want[index+1:]...)
	}

	return &pb.Empty{}, nil
}
<commit_msg>Added remove from want on remove<commit_after>package main

import (
	"errors"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/brotherlogic/godiscogs"
	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"

	pb "github.com/brotherlogic/discogssyncer/server"
	pbd "github.com/brotherlogic/godiscogs"
)

// GetRelease Gets the release and metadata for the release
func (syncer *Syncer) GetRelease(id int, folder int) (*pbd.Release, *pb.ReleaseMetadata) {
	releaseData, err := ioutil.ReadFile(syncer.saveLocation + "/" + strconv.Itoa(folder) + "/" + strconv.Itoa(id) + ".release")
	if err != nil {
		return nil, nil
	}

	metadataData, _ := ioutil.ReadFile(syncer.saveLocation + "/static-metadata/" + strconv.Itoa(id) + ".metadata")
	release := &pbd.Release{}
	metadata := &pb.ReleaseMetadata{}

	proto.Unmarshal(releaseData, release)
	proto.Unmarshal(metadataData, metadata)

	return release, metadata
}

// GetMonthlySpend gets the monthly spend
func (syncer *Syncer) GetMonthlySpend(ctx context.Context, req *pb.SpendRequest) (*pb.SpendResponse, error) {
	spend := 0
	for _, rel := range syncer.relMap {
		_, metadata := syncer.GetRelease(int(rel.Id), int(rel.FolderId))
		datev := time.Unix(metadata.DateAdded, 0)
		if datev.Year() == int(req.Year) && int32(datev.Month()) == req.Month {
			spend += int(metadata.Cost)
		}
	}

	return &pb.SpendResponse{TotalSpend: int32(spend)}, nil
}

func (syncer *Syncer) saveMetadata(rel *godiscogs.Release) {
	metadataRoot := syncer.saveLocation + "/static-metadata/"
	metadataPath := metadataRoot + strconv.Itoa(int(rel.Id)) + ".metadata"
	if _, err := os.Stat(metadataRoot); os.IsNotExist(err) {
		os.MkdirAll(metadataRoot, 0777)
	}

	metadata := &pb.ReleaseMetadata{}
	if _, err := os.Stat(metadataPath); os.IsNotExist(err) {
		metadata.DateAdded = time.Now().Unix()
		metadata.DateRefreshed = time.Now().Unix()
	} else {
		data, _ := ioutil.ReadFile(metadataPath)
		proto.Unmarshal(data, metadata)
		metadata.DateRefreshed = time.Now().Unix()
	}
	data, _ := proto.Marshal(metadata)
	ioutil.WriteFile(metadataPath, data, 0644)
}

func (syncer *Syncer) saveRelease(rel *godiscogs.Release, folder int) {
	//Check that the save folder exists
	savePath := syncer.saveLocation + "/" + strconv.Itoa(folder) + "/"
	if _, err := os.Stat(savePath); os.IsNotExist(err) {
		os.MkdirAll(savePath, 0777)
	}

	data, _ := proto.Marshal(rel)
	ioutil.WriteFile(savePath+strconv.Itoa(int(rel.Id))+".release", data, 0644)
	syncer.saveMetadata(rel)
}

type saver interface {
	GetCollection() []godiscogs.Release
	GetFolders() []godiscogs.Folder
	GetRelease(id int) (godiscogs.Release, error)
	MoveToFolder(folderID int, releaseID int, instanceID int, newFolderID int)
	AddToFolder(folderID int, releaseID int)
	SetRating(folderID int, releaseID int, instanceID int, rating int)
	GetWantlist() ([]pbd.Release, error)
	RemoveFromWantlist(releaseID int)
	AddToWantlist(releaseID int)
}

// EditWant edits a want in the wantlist
func (syncer *Syncer) EditWant(ctx context.Context, wantIn *pb.Want) (*pb.Want, error) {
	for _, want := range syncer.wants.Want {
		if want.ReleaseId == wantIn.ReleaseId {
			want.Valued = wantIn.Valued
		}
	}
	syncer.saveWantList()

	return wantIn, nil
}

// SaveCollection writes out the full collection to files.
func (syncer *Syncer) SaveCollection(retr saver) {
	releases := retr.GetCollection()
	for _, release := range releases {
		fullRelease, _ := retr.GetRelease(int(release.Id))
		fullRelease.InstanceId = release.InstanceId
		fullRelease.FolderId = release.FolderId
		fullRelease.Rating = release.Rating
		syncer.saveRelease(&fullRelease, int(release.FolderId))
	}
	folders := retr.GetFolders()
	folderList := pb.FolderList{}
	for i := range folders {
		folder := folders[i]
		folderList.Folders = append(folderList.Folders, &folder)
	}
	syncer.SaveFolders(&folderList)
}

// SyncWantlist syncs the wantlist with the server
func (syncer *Syncer) SyncWantlist() {
	wants, _ := syncer.retr.GetWantlist()

	for _, want := range wants {
		seen := false
		var val *pb.Want
		for _, swant := range syncer.wants.Want {
			if swant.ReleaseId == want.Id {
				seen = true
				val = swant
			}
		}

		if seen {
			val.Wanted = true
		} else {
			syncer.wants.Want = append(syncer.wants.Want, &pb.Want{ReleaseId: want.Id, Valued: false, Wanted: true})
		}
	}

	// Cache the want list releases
	for _, want := range syncer.wants.Want {
		release, _ := syncer.retr.GetRelease(int(want.ReleaseId))
		syncer.saveRelease(&release, -5)
	}

	syncer.saveWantList()
}

func (syncer *Syncer) getFolders() *pb.FolderList {
	data, _ := ioutil.ReadFile(syncer.saveLocation + "/metadata/folders")
	folderData := &pb.FolderList{}
	proto.Unmarshal(data, folderData)
	return folderData
}

// GetSingleRelease gets a single release
func (syncer *Syncer) GetSingleRelease(ctx context.Context, in *pbd.Release) (*pbd.Release, error) {
	if val, ok := syncer.relMap[in.Id]; ok {
		return val, nil
	}
	return nil, errors.New("Unable to find release")
}

// CollapseWantlist collapses the wantlist
func (syncer *Syncer) CollapseWantlist(ctx context.Context, in *pb.Empty) (*pb.Wantlist, error) {
	for _, want := range syncer.wants.Want {
		if !want.Valued {
			log.Printf("AVOIDING %v", want)
			syncer.retr.RemoveFromWantlist(int(want.ReleaseId))
			want.Wanted = false
		}
	}

	return &syncer.wants, nil
}

// RebuildWantlist rebuilds the wantlist
func (syncer *Syncer) RebuildWantlist(ctx context.Context, in *pb.Empty) (*pb.Wantlist, error) {
	for _, want := range syncer.wants.Want {
		syncer.retr.AddToWantlist(int(want.ReleaseId))
		want.Wanted = true
	}

	return &syncer.wants, nil
}

// AddToFolder adds a release to the specified folder
func (syncer *Syncer) AddToFolder(ctx context.Context, in *pb.ReleaseMove) (*pb.Empty, error) {
	syncer.retr.AddToFolder(int(in.NewFolderId), int(in.Release.Id))
	fullRelease, _ := syncer.retr.GetRelease(int(in.Release.Id))
	fullRelease.FolderId = int32(in.NewFolderId)
	syncer.saveRelease(&fullRelease, int(in.NewFolderId))
	syncer.relMap[in.Release.Id] = &fullRelease
	return &pb.Empty{}, nil
}

// UpdateRating updates the rating of a release
func (syncer *Syncer) UpdateRating(ctx context.Context, in *pbd.Release) (*pb.Empty, error) {
	syncer.retr.SetRating(int(in.FolderId), int(in.Id), int(in.InstanceId), int(in.Rating))
	fullRelease, _ := syncer.GetRelease(int(in.Id), int(in.FolderId))
	fullRelease.Rating = int32(in.Rating)
	syncer.relMap[in.Id] = fullRelease
	syncer.saveRelease(fullRelease, int(fullRelease.FolderId))
	return &pb.Empty{}, nil
}

// UpdateMetadata updates the metadata of a given record
func (syncer *Syncer) UpdateMetadata(ctx context.Context, in *pb.MetadataUpdate) (*pb.ReleaseMetadata, error) {
	release, metadata := syncer.GetRelease(int(in.Release.Id), int(in.Release.FolderId))
	if release == nil {
		return nil, errors.New("Unable to locate release")
	}
	proto.Merge(metadata, in.Update)

	metadataRoot := syncer.saveLocation + "/static-metadata/"
	metadataPath := metadataRoot + strconv.Itoa(int(in.Release.Id)) + ".metadata"
	data, _ := proto.Marshal(metadata)
	ioutil.WriteFile(metadataPath, data, 0644)

	return metadata, nil
}

// GetWantlist gets the wantlist
func (syncer *Syncer) GetWantlist(ctx context.Context, in *pb.Empty) (*pb.Wantlist, error) {
	return &syncer.wants, nil
}

// GetMetadata gets the metadata for a given release
func (syncer *Syncer) GetMetadata(ctx context.Context, in *pbd.Release) (*pb.ReleaseMetadata, error) {
	_, metadata := syncer.GetRelease(int(in.Id), int(in.FolderId))
	return metadata, nil
}

// GetReleasesInFolder serves up the releases in a given folder
func (syncer *Syncer) GetReleasesInFolder(ctx context.Context, in *pb.FolderList) (*pb.ReleaseList, error) {

	releases := pb.ReleaseList{}
	for _, folderSpec := range in.Folders {
		folders := syncer.getFolders()
		for _, folder := range folders.Folders {
			if folder.Name == folderSpec.Name || folder.Id == folderSpec.Id {
				innerReleases := syncer.getReleases(int(folder.Id))
				releases.Releases = append(releases.Releases, innerReleases.Releases...)
			}
		}
	}

	return &releases, nil
}

func (syncer *Syncer) getReleases(folderID int) *pb.ReleaseList {
	releases := pb.ReleaseList{}
	files, _ := ioutil.ReadDir(syncer.saveLocation + "/" + strconv.Itoa(folderID) + "/")
	for _, file := range files {
		if strings.HasSuffix(file.Name(), ".release") {
			data, _ := ioutil.ReadFile(syncer.saveLocation + "/" + strconv.Itoa(folderID) + "/" + file.Name())
			release := &pbd.Release{}
			proto.Unmarshal(data, release)
			releases.Releases = append(releases.Releases, release)
		}
	}
	return &releases
}

func (syncer *Syncer) saveWantList() {
	data, _ := proto.Marshal(&syncer.wants)
	savePath := syncer.saveLocation + "/metadata/"
	if _, err := os.Stat(savePath); os.IsNotExist(err) {
		os.MkdirAll(savePath, 0777)
	}
	ioutil.WriteFile(savePath+"wantlist", data, 0644)
}

// SaveFolders saves out the list of folders
func (syncer *Syncer) SaveFolders(list *pb.FolderList) {
	savePath := syncer.saveLocation + "/metadata/"
	if _, err := os.Stat(savePath); os.IsNotExist(err) {
		os.MkdirAll(savePath, 0777)
	}

	data, _ := proto.Marshal(list)
	ioutil.WriteFile(savePath+"folders", data, 0644)
}

// GetCollection serves up the whole of the collection
func (syncer *Syncer) GetCollection(ctx context.Context, in *pb.Empty) (*pb.ReleaseList, error) {
	releases := pb.ReleaseList{}
	bfiles, _ := ioutil.ReadDir(syncer.saveLocation)
	for _, bfile := range bfiles {
		if bfile.IsDir() && bfile.Name() != "-5" {
			folderID, _ := strconv.Atoi(bfile.Name())
			for _, release := range syncer.getReleases(folderID).Releases {
				releases.Releases = append(releases.Releases, release)
			}
		}
	}
	return &releases, nil
}

// DeleteWant removes a want from the system
func (syncer *Syncer) DeleteWant(ctx context.Context, in *pb.Want) (*pb.Empty, error) {
	//Remove the want file and remove from
	index := -1
	for i, val := range syncer.wants.Want {
		if val.ReleaseId == in.ReleaseId {
			index = i
		}
	}

	if index >= 0 {
		syncer.wants.Want = append(syncer.wants.Want[:index], syncer.wants.Want[index+1:]...)
	}

	syncer.retr.RemoveFromWantlist(int(in.ReleaseId))

	return &pb.Empty{}, nil
}
<|endoftext|>


<commit_before>package gcssync

import (
	"code.google.com/p/google-api-go-client/storage/v1"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"
)

const (
	fileWalkerBufferSize = 20
	fileWalkerEstimate   = 500
)

type fileInfo struct {
	info os.FileInfo
	path string
}

func getLocalFiles(dirname string) []fileInfo {
	rawFiles := make(chan fileInfo, fileWalkerBufferSize)
	var done sync.WaitGroup

	done.Add(1)
	go func() {
		filepath.Walk(dirname, func(path string, f os.FileInfo, err error) error {
			if err == nil && !f.IsDir() {
				rel, _ := filepath.Rel(dirname, path)
				rawFiles <- fileInfo{
					info: f,
					path: rel,
				}
			}
			return nil
		})
		close(rawFiles)
		done.Done()
	}()

	result := make([]fileInfo, 0, fileWalkerEstimate)

	// this is unnecessary, but left there in case filters for files will be implemented
	done.Add(1)
	go func() {
		for {
			f, more := <-rawFiles
			if more {
				result = append(result, f)
			} else {
				done.Done()
				return
			}
		}
	}()
	done.Wait()
	return result
}

func (c *Client) SyncFolder(from, to string) {
	fmt.Println(from, to)
	var localFiles []fileInfo
	var listsFetched sync.WaitGroup

	listsFetched.Add(1)
	go func() {
		localFiles = getLocalFiles(from)
		listsFetched.Done()
	}()

	var objects *storage.Objects
	var err error

	listsFetched.Add(1)
	go func() {
		objects, err = c.ListFiles()
		if err != nil {
			fmt.Println(err)
		}
		listsFetched.Done()
	}()

	listsFetched.Wait()

	fmt.Printf("Found %d local files\n", len(localFiles))
	fmt.Printf("Found %d remote files\n", len(objects.Items))

	remoteCache := make(map[string]time.Time, len(objects.Items))

	for _, object := range objects.Items {
		time, err := time.Parse(time.RFC3339, object.Updated)
		if err != nil {
			continue
		}
		remoteCache[object.Name] = time
	}
}
<commit_msg>feat(sync): basic sync works<commit_after>package gcssync

import (
	"code.google.com/p/google-api-go-client/storage/v1"
	"fmt"
	"github.com/dustin/go-humanize"
	"os"
	"path/filepath"
	"sync"
	"time"
)

const (
	fileWalkerBufferSize = 20
	fileWalkerEstimate   = 500
	uploadGoroutinesNum  = 10
)

type fileInfo struct {
	info os.FileInfo
	path string
}

func getLocalFiles(dirname string) []fileInfo {
	rawFiles := make(chan fileInfo, fileWalkerBufferSize)
	var done sync.WaitGroup

	done.Add(1)
	go func() {
		filepath.Walk(dirname, func(path string, f os.FileInfo, err error) error {
			if err == nil && !f.IsDir() {
				rel, _ := filepath.Rel(dirname, path)
				rawFiles <- fileInfo{
					info: f,
					path: rel,
				}
			}
			return nil
		})
		close(rawFiles)
		done.Done()
	}()

	result := make([]fileInfo, 0, fileWalkerEstimate)

	// this is unnecessary, but left there in case filters for files will be implemented
	done.Add(1)
	go func() {
		for {
			f, more := <-rawFiles
			if more {
				result = append(result, f)
			} else {
				done.Done()
				return
			}
		}
	}()
	done.Wait()
	return result
}

func (c *Client) SyncFolder(from, to string) {
	fmt.Println(from, to)
	var localFiles []fileInfo
	var listsFetched sync.WaitGroup

	listsFetched.Add(1)
	go func() {
		localFiles = getLocalFiles(from)
		listsFetched.Done()
	}()

	var objects *storage.Objects
	var err error

	listsFetched.Add(1)
	go func() {
		objects, err = c.ListFiles()
		if err != nil {
			fmt.Println(err)
		}
		listsFetched.Done()
	}()

	listsFetched.Wait()

	fmt.Printf("Found %d local files\n", len(localFiles))
	fmt.Printf("Found %d remote files\n", len(objects.Items))

	remoteCache := make(map[string]time.Time, len(objects.Items))

	for _, object := range objects.Items {
		time, err := time.Parse(time.RFC3339, object.Updated)
		if err != nil {
			continue
		}
		remoteCache[object.Name] = time
	}

	already := 0
	created := 0
	update := 0

	toDo := make([]fileInfo, 0, len(localFiles))

	for _, file := range localFiles {
		remoteTime, exists := remoteCache[file.path]
		if exists {
			already++
			if file.info.ModTime().After(remoteTime) {
				toDo = append(toDo, file)
				update++
			}
		} else {
			created++
			toDo = append(toDo, file)
		}
	}

	var uploadDone sync.WaitGroup
	uploadFile := make(chan fileInfo, uploadGoroutinesNum)
	uploadDone.Add(1)
	go func() {
		for _, file := range toDo {
			uploadFile <- file
		}
		close(uploadFile)
		uploadDone.Done()
	}()

	for i := 0; i < uploadGoroutinesNum; i++ {
		uploadDone.Add(1)
		go func() {
			for {
				file, more := <-uploadFile
				if more {
					success, object, err := c.UploadFile(filepath.Join(from, file.path), file.path)
					if err != nil || success == false {
						fmt.Println(err)
					} else {
						fmt.Printf("Uploaded %s %s\n", object.Name, humanize.Bytes(object.Size))
					}
				} else {
					uploadDone.Done()
					return
				}
			}
		}()
	}

	uploadDone.Wait()

	fmt.Println(len(toDo))
	fmt.Println(already, created, update)
}
<|endoftext|>


<commit_before>package main

import "errors"
import "github.com/brotherlogic/godiscogs"
import "github.com/golang/protobuf/proto"
import "io/ioutil"
import "log"
import "os"
import "strconv"
import "strings"
import "time"

import "golang.org/x/net/context"
import pb "github.com/brotherlogic/discogssyncer/server"
import pbd "github.com/brotherlogic/godiscogs"

// GetRelease Gets the release and metadata
func (syncer *Syncer) GetRelease(id int, folder int) (*pbd.Release, *pb.ReleaseMetadata) {
	releaseData, _ := ioutil.ReadFile(syncer.saveLocation + "/" + strconv.Itoa(folder) + "/" + strconv.Itoa(id) + ".release")
	metadataData, _ := ioutil.ReadFile(syncer.saveLocation + "/static-metadata/" + strconv.Itoa(id) + ".metadata")
	release := &pbd.Release{}
	metadata := &pb.ReleaseMetadata{}

	proto.Unmarshal(releaseData, release)
	proto.Unmarshal(metadataData, metadata)

	log.Printf("Reading %v from %v", metadataData, syncer.saveLocation+"/static-metadata/"+strconv.Itoa(id)+".metadata")

	return release, metadata
}

func (syncer *Syncer) saveMetadata(rel *godiscogs.Release) {
	metadataRoot := syncer.saveLocation + "/static-metadata/"
	metadataPath := metadataRoot + strconv.Itoa(int(rel.Id)) + ".metadata"
	if _, err := os.Stat(metadataRoot); os.IsNotExist(err) {
		os.MkdirAll(metadataRoot, 0777)
	}

	metadata := &pb.ReleaseMetadata{}
	if _, err := os.Stat(metadataPath); os.IsNotExist(err) {
		metadata.DateAdded = time.Now().Unix()
		metadata.DateRefreshed = time.Now().Unix()
	} else {
		data, _ := ioutil.ReadFile(metadataPath)
		proto.Unmarshal(data, metadata)
		metadata.DateRefreshed = time.Now().Unix()
	}
	log.Printf("Writing out %v", metadata)
	data, _ := proto.Marshal(metadata)
	ioutil.WriteFile(metadataPath, data, 0644)
}

func (syncer *Syncer) saveRelease(rel *godiscogs.Release, folder int) {
	//Check that the save folder exists
	savePath := syncer.saveLocation + "/" + strconv.Itoa(folder) + "/"
	if _, err := os.Stat(savePath); os.IsNotExist(err) {
		os.MkdirAll(savePath, 0777)
	}

	log.Printf("Saving %v", rel)
	data, _ := proto.Marshal(rel)
	ioutil.WriteFile(savePath+strconv.Itoa(int(rel.Id))+".release", data, 0644)
	syncer.saveMetadata(rel)
}

type saver interface {
	GetCollection() []godiscogs.Release
	GetFolders() []godiscogs.Folder
	GetRelease(id int) (godiscogs.Release, error)
	MoveToFolder(folderID int, releaseID int, instanceID int, newFolderID int)
	AddToFolder(folderID int, releaseID int)
}

// SaveCollection writes out the full collection to files.
func (syncer *Syncer) SaveCollection(retr saver) {
	releases := retr.GetCollection()
	for _, release := range releases {
		fullRelease, _ := retr.GetRelease(int(release.Id))
		fullRelease.InstanceId = release.InstanceId
		fullRelease.FolderId = release.FolderId
		syncer.saveRelease(&fullRelease, int(release.FolderId))
	}
	folders := retr.GetFolders()
	folderList := pb.FolderList{}
	for i := range folders {
		folder := folders[i]
		folderList.Folders = append(folderList.Folders, &folder)
	}
	syncer.SaveFolders(&folderList)
}

func (syncer *Syncer) getFolders() *pb.FolderList {
	data, _ := ioutil.ReadFile(syncer.saveLocation + "/metadata/folders")
	folderData := &pb.FolderList{}
	proto.Unmarshal(data, folderData)
	return folderData
}

// MoveToFolder moves a release to the specified folder
func (syncer *Syncer) MoveToFolder(ctx context.Context, in *pb.ReleaseMove) (*pb.Empty, error) {
	syncer.retr.MoveToFolder(int(in.Release.FolderId), int(in.Release.Id), int(in.Release.InstanceId), int(in.NewFolderId))
	return &pb.Empty{}, nil
}

// AddToFolder adds a release to the specified folder
func (syncer *Syncer) AddToFolder(ctx context.Context, in *pb.ReleaseMove) (*pb.Empty, error) {
	syncer.retr.AddToFolder(int(in.NewFolderId), int(in.Release.Id))
	return &pb.Empty{}, nil
}

// GetReleasesInFolder serves up the releases in a given folder
func (syncer *Syncer) GetReleasesInFolder(ctx context.Context, in *pb.FolderList) (*pb.ReleaseList, error) {

	releases := pb.ReleaseList{}
	for _, folderSpec := range in.Folders {
		folders := syncer.getFolders()
		for _, folder := range folders.Folders {
			if folder.Name == folderSpec.Name {
				innerReleases := syncer.getReleases(int(folder.Id))
				releases.Releases = append(releases.Releases, innerReleases.Releases...)
			}
		}
	}

	if len(releases.Releases) > 0 {
		return &releases, nil
	}

	return &pb.ReleaseList{}, errors.New("Folder does not exist in collection")
}

func (syncer *Syncer) getReleases(folderID int) *pb.ReleaseList {
	releases := pb.ReleaseList{}
	files, _ := ioutil.ReadDir(syncer.saveLocation + "/" + strconv.Itoa(folderID) + "/")
	for _, file := range files {
		if strings.HasSuffix(file.Name(), ".release") {
			data, _ := ioutil.ReadFile(syncer.saveLocation + "/" + strconv.Itoa(folderID) + "/" + file.Name())
			release := &pbd.Release{}
			proto.Unmarshal(data, release)
			releases.Releases = append(releases.Releases, release)
		}
	}
	return &releases
}

// SaveFolders saves out the list of folders
func (syncer *Syncer) SaveFolders(list *pb.FolderList) {
	savePath := syncer.saveLocation + "/metadata/"
	if _, err := os.Stat(savePath); os.IsNotExist(err) {
		os.MkdirAll(savePath, 0777)
	}

	data, _ := proto.Marshal(list)
	ioutil.WriteFile(savePath+"folders", data, 0644)
}

// GetCollection serves up the whole of the collection
func (syncer *Syncer) GetCollection(ctx context.Context, in *pb.Empty) (*pb.ReleaseList, error) {
	releases := pb.ReleaseList{}
	bfiles, _ := ioutil.ReadDir(syncer.saveLocation)
	for _, bfile := range bfiles {
		if bfile.IsDir() {
			log.Printf("Searching %v", bfile)
			folderID, _ := strconv.Atoi(bfile.Name())
			for _, release := range syncer.getReleases(folderID).Releases {
				releases.Releases = append(releases.Releases, release)
			}
		}
	}
	return &releases, nil
}
<commit_msg>Added end point stub<commit_after>package main

import "errors"
import "github.com/brotherlogic/godiscogs"
import "github.com/golang/protobuf/proto"
import "io/ioutil"
import "log"
import "os"
import "strconv"
import "strings"
import "time"

import "golang.org/x/net/context"
import pb "github.com/brotherlogic/discogssyncer/server"
import pbd "github.com/brotherlogic/godiscogs"

// GetRelease Gets the release and metadata
func (syncer *Syncer) GetRelease(id int, folder int) (*pbd.Release, *pb.ReleaseMetadata) {
	releaseData, _ := ioutil.ReadFile(syncer.saveLocation + "/" + strconv.Itoa(folder) + "/" + strconv.Itoa(id) + ".release")
	metadataData, _ := ioutil.ReadFile(syncer.saveLocation + "/static-metadata/" + strconv.Itoa(id) + ".metadata")
	release := &pbd.Release{}
	metadata := &pb.ReleaseMetadata{}

	proto.Unmarshal(releaseData, release)
	proto.Unmarshal(metadataData, metadata)

	log.Printf("Reading %v from %v", metadataData, syncer.saveLocation+"/static-metadata/"+strconv.Itoa(id)+".metadata")

	return release, metadata
}

func (syncer *Syncer) saveMetadata(rel *godiscogs.Release) {
	metadataRoot := syncer.saveLocation + "/static-metadata/"
	metadataPath := metadataRoot + strconv.Itoa(int(rel.Id)) + ".metadata"
	if _, err := os.Stat(metadataRoot); os.IsNotExist(err) {
		os.MkdirAll(metadataRoot, 0777)
	}

	metadata := &pb.ReleaseMetadata{}
	if _, err := os.Stat(metadataPath); os.IsNotExist(err) {
		metadata.DateAdded = time.Now().Unix()
		metadata.DateRefreshed = time.Now().Unix()
	} else {
		data, _ := ioutil.ReadFile(metadataPath)
		proto.Unmarshal(data, metadata)
		metadata.DateRefreshed = time.Now().Unix()
	}
	log.Printf("Writing out %v", metadata)
	data, _ := proto.Marshal(metadata)
	ioutil.WriteFile(metadataPath, data, 0644)
}

func (syncer *Syncer) saveRelease(rel *godiscogs.Release, folder int) {
	//Check that the save folder exists
	savePath := syncer.saveLocation + "/" + strconv.Itoa(folder) + "/"
	if _, err := os.Stat(savePath); os.IsNotExist(err) {
		os.MkdirAll(savePath, 0777)
	}

	log.Printf("Saving %v", rel)
	data, _ := proto.Marshal(rel)
	ioutil.WriteFile(savePath+strconv.Itoa(int(rel.Id))+".release", data, 0644)
	syncer.saveMetadata(rel)
}

type saver interface {
	GetCollection() []godiscogs.Release
	GetFolders() []godiscogs.Folder
	GetRelease(id int) (godiscogs.Release, error)
	MoveToFolder(folderID int, releaseID int, instanceID int, newFolderID int)
	AddToFolder(folderID int, releaseID int)
}

// SaveCollection writes out the full collection to files.
func (syncer *Syncer) SaveCollection(retr saver) {
	releases := retr.GetCollection()
	for _, release := range releases {
		fullRelease, _ := retr.GetRelease(int(release.Id))
		fullRelease.InstanceId = release.InstanceId
		fullRelease.FolderId = release.FolderId
		syncer.saveRelease(&fullRelease, int(release.FolderId))
	}
	folders := retr.GetFolders()
	folderList := pb.FolderList{}
	for i := range folders {
		folder := folders[i]
		folderList.Folders = append(folderList.Folders, &folder)
	}
	syncer.SaveFolders(&folderList)
}

func (syncer *Syncer) getFolders() *pb.FolderList {
	data, _ := ioutil.ReadFile(syncer.saveLocation + "/metadata/folders")
	folderData := &pb.FolderList{}
	proto.Unmarshal(data, folderData)
	return folderData
}

// MoveToFolder moves a release to the specified folder
func (syncer *Syncer) MoveToFolder(ctx context.Context, in *pb.ReleaseMove) (*pb.Empty, error) {
	syncer.retr.MoveToFolder(int(in.Release.FolderId), int(in.Release.Id), int(in.Release.InstanceId), int(in.NewFolderId))
	return &pb.Empty{}, nil
}

// AddToFolder adds a release to the specified folder
func (syncer *Syncer) AddToFolder(ctx context.Context, in *pb.ReleaseMove) (*pb.Empty, error) {
	syncer.retr.AddToFolder(int(in.NewFolderId), int(in.Release.Id))
	return &pb.Empty{}, nil
}

// UpdateMetadata updates the metadata of a given record
func (syncer *Syncer) UpdateMetadata(ctx context.Context, in *pb.MetadataUpdate) (*pb.ReleaseMetadata, error) {
	return &pb.ReleaseMetadata{}, nil
}

// GetReleasesInFolder serves up the releases in a given folder
func (syncer *Syncer) GetReleasesInFolder(ctx context.Context, in *pb.FolderList) (*pb.ReleaseList, error) {

	releases := pb.ReleaseList{}
	for _, folderSpec := range in.Folders {
		folders := syncer.getFolders()
		for _, folder := range folders.Folders {
			if folder.Name == folderSpec.Name {
				innerReleases := syncer.getReleases(int(folder.Id))
				releases.Releases = append(releases.Releases, innerReleases.Releases...)
			}
		}
	}

	if len(releases.Releases) > 0 {
		return &releases, nil
	}

	return &pb.ReleaseList{}, errors.New("Folder does not exist in collection")
}

func (syncer *Syncer) getReleases(folderID int) *pb.ReleaseList {
	releases := pb.ReleaseList{}
	files, _ := ioutil.ReadDir(syncer.saveLocation + "/" + strconv.Itoa(folderID) + "/")
	for _, file := range files {
		if strings.HasSuffix(file.Name(), ".release") {
			data, _ := ioutil.ReadFile(syncer.saveLocation + "/" + strconv.Itoa(folderID) + "/" + file.Name())
			release := &pbd.Release{}
			proto.Unmarshal(data, release)
			releases.Releases = append(releases.Releases, release)
		}
	}
	return &releases
}

// SaveFolders saves out the list of folders
func (syncer *Syncer) SaveFolders(list *pb.FolderList) {
	savePath := syncer.saveLocation + "/metadata/"
	if _, err := os.Stat(savePath); os.IsNotExist(err) {
		os.MkdirAll(savePath, 0777)
	}

	data, _ := proto.Marshal(list)
	ioutil.WriteFile(savePath+"folders", data, 0644)
}

// GetCollection serves up the whole of the collection
func (syncer *Syncer) GetCollection(ctx context.Context, in *pb.Empty) (*pb.ReleaseList, error) {
	releases := pb.ReleaseList{}
	bfiles, _ := ioutil.ReadDir(syncer.saveLocation)
	for _, bfile := range bfiles {
		if bfile.IsDir() {
			log.Printf("Searching %v", bfile)
			folderID, _ := strconv.Atoi(bfile.Name())
			for _, release := range syncer.getReleases(folderID).Releases {
				releases.Releases = append(releases.Releases, release)
			}
		}
	}
	return &releases, nil
}
<|endoftext|>


<commit_before>package conncrypt

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"hash"
	"io"
	"net"

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

//Confg defaults
const DefaultIterations = 2048
const DefaultKeySize = 32 //256bits
var DefaultHashFunc = sha256.New
var DefaultSalt = []byte("conncrypt")

//Config stores the PBKDF2 key generation parameters
type Config struct {
	Password   string
	Salt       []byte
	Iterations int
	KeySize    int
	HashFunc   func() hash.Hash
}

//New creates an AES encrypted net.Conn by generating
//a key using PBKDF2 with the provided configuration
func New(conn net.Conn, c *Config) net.Conn {
	//set defaults
	if len(c.Salt) == 0 {
		c.Salt = DefaultSalt
	}
	if c.Iterations == 0 {
		c.Iterations = DefaultIterations
	}
	if c.KeySize != 16 && c.KeySize != 24 && c.KeySize != 32 {
		c.KeySize = DefaultKeySize
	}
	if c.HashFunc == nil {
		c.HashFunc = DefaultHashFunc
	}
	//generate key
	key := pbkdf2.Key([]byte(c.Password), c.Salt, c.Iterations, c.KeySize, c.HashFunc)
	//key will be always be the correct size so this will never error
	conn, _ = NewFromKey(conn, key)
	return conn
}

//NewFromKey creates an AES encrypted net.Conn using the provided key
func NewFromKey(conn net.Conn, key []byte) (net.Conn, error) {
	block, err := aes.NewCipher([]byte(key))
	if err != nil {
		return nil, err
	}
	//hash(key) -> read IV
	riv := DefaultHashFunc().Sum(key)
	rstream := cipher.NewOFB(block, riv[:aes.BlockSize])
	reader := &cipher.StreamReader{S: rstream, R: conn}
	//hash(read IV) -> write IV
	wiv := DefaultHashFunc().Sum(riv)
	wstream := cipher.NewOFB(block, wiv[:aes.BlockSize])
	writer := &cipher.StreamWriter{S: wstream, W: conn}

	return &cryptoConn{
		Conn: conn,
		r:    reader,
		w:    writer,
	}, nil
}

type cryptoConn struct {
	net.Conn
	r io.Reader
	w io.Writer
}

//replace read and write methods
func (c *cryptoConn) Read(p []byte) (int, error) {
	return c.r.Read(p)
}
func (c *cryptoConn) Write(p []byte) (int, error) {
	return c.w.Write(p)
}
<commit_msg>strengthen salt<commit_after>package conncrypt

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"hash"
	"io"
	"net"

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

//Confg defaults
const DefaultIterations = 2048
const DefaultKeySize = 32 //256bits
var DefaultHashFunc = sha256.New
var DefaultSalt = []byte(`
(;QUHj.BQ?RXzYSO]ifkXp/G!kFmWyXyEV6Nt!d|@bo+N$L9+<d$|g6e26T}
Ao<:>SOd,6acYKY_ec+(x"R";\'4&fTAVu92GVA-wxBptOTM^2,iP5%)wnhW
hwk=]Snsgymt!3gbP2pe=J//}1a?lp9ej=&TB!C_V(cT2?z8wyoL_-13fd[]
`) //salt must be predefined in order to derive the same key

//Config stores the PBKDF2 key generation parameters
type Config struct {
	Password   string
	Salt       []byte
	Iterations int
	KeySize    int
	HashFunc   func() hash.Hash
}

//New creates an AES encrypted net.Conn by generating
//a key using PBKDF2 with the provided configuration
func New(conn net.Conn, c *Config) net.Conn {
	//set defaults
	if len(c.Salt) == 0 {
		c.Salt = DefaultSalt
	}
	if c.Iterations == 0 {
		c.Iterations = DefaultIterations
	}
	if c.KeySize != 16 && c.KeySize != 24 && c.KeySize != 32 {
		c.KeySize = DefaultKeySize
	}
	if c.HashFunc == nil {
		c.HashFunc = DefaultHashFunc
	}

	//generate key
	key := pbkdf2.Key([]byte(c.Password), c.Salt, c.Iterations, c.KeySize, c.HashFunc)

	// could use scrypt, but it's a bit slow...
	// dk, err := scrypt.Key([]byte(c.Password), c.Salt, 16384, 8, 1, 32)

	//key will be always be the correct size so this will never error
	conn, _ = NewFromKey(conn, key)
	return conn
}

//NewFromKey creates an AES encrypted net.Conn using the provided key
func NewFromKey(conn net.Conn, key []byte) (net.Conn, error) {
	block, err := aes.NewCipher([]byte(key))
	if err != nil {
		return nil, err
	}
	//hash(key) -> read IV
	riv := DefaultHashFunc().Sum(key)
	rstream := cipher.NewOFB(block, riv[:aes.BlockSize])
	reader := &cipher.StreamReader{S: rstream, R: conn}
	//hash(read IV) -> write IV
	wiv := DefaultHashFunc().Sum(riv)
	wstream := cipher.NewOFB(block, wiv[:aes.BlockSize])
	writer := &cipher.StreamWriter{S: wstream, W: conn}

	return &cryptoConn{
		Conn: conn,
		r:    reader,
		w:    writer,
	}, nil
}

type cryptoConn struct {
	net.Conn
	r io.Reader
	w io.Writer
}

//replace read and write methods
func (c *cryptoConn) Read(p []byte) (int, error) {
	return c.r.Read(p)
}
func (c *cryptoConn) Write(p []byte) (int, error) {
	return c.w.Write(p)
}
<|endoftext|>


<commit_before>package tail

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

	fsnotify "github.com/fsnotify/fsnotify"
)

const (
	openRetryInterval = time.Second
	tailOldFileDelay  = 15 * time.Second
	linesCapacity     = 1024
	errorsCapacity    = 16
)

// Line is a line of the target file.
type Line struct {
	Text string
	Time time.Time
}

// Tail tails a file.
type Tail struct {
	Lines  <-chan *Line
	Errors <-chan error

	lines    chan<- *Line
	errors   chan<- error
	filename string
	wg       sync.WaitGroup
	ctx      context.Context
	cancel   context.CancelFunc
}

type tail struct {
	parent *Tail

	file    *os.File
	reader  *bufio.Reader
	watcher *fsnotify.Watcher
	buf     string
	ctx     context.Context
	cancel  context.CancelFunc
}

// NewTailFile starts tailing a file
func NewTailFile(filename string) (*Tail, error) {
	filename, err := filepath.Abs(filename)
	if err != nil {
		return nil, fmt.Errorf("tail: failed to get the absolute path: %w", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	lines := make(chan *Line, linesCapacity)
	errs := make(chan error, errorsCapacity)
	parent := &Tail{
		Lines:    lines,
		Errors:   errs,
		lines:    lines,
		errors:   errs,
		filename: filename,
		ctx:      ctx,
		cancel:   cancel,
	}
	parent.wg.Add(1)
	go parent.runFile(os.SEEK_END)
	go parent.wait()
	return parent, nil
}

// NewTailReader starts tailing io.Reader
func NewTailReader(reader io.Reader) (*Tail, error) {
	ctx, cancel := context.WithCancel(context.Background())
	lines := make(chan *Line, linesCapacity)
	errs := make(chan error, errorsCapacity)
	parent := &Tail{
		Lines:  lines,
		Errors: errs,
		lines:  lines,
		errors: errs,
		ctx:    ctx,
		cancel: cancel,
	}
	t := &tail{
		parent: parent,
		reader: bufio.NewReader(reader),
		ctx:    ctx,
		cancel: cancel,
	}
	parent.wg.Add(1)
	go t.runReader()
	go parent.wait()
	return parent, nil
}

// Close stops tailing the file.
func (t *Tail) Close() error {
	t.cancel()
	t.wg.Wait()
	return nil
}

func (t *Tail) wait() {
	t.wg.Wait()
	close(t.errors)
	close(t.lines)
}

// open opens the target file.
// If it does not exist, wait for creating new file.
func (t *Tail) open(seek int) (*tail, error) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, fmt.Errorf("tail: failed to initialize fsnotify: %w", err)
	}
	for {
		file, err := os.Open(t.filename)
		if err == nil {
			// success, seek and watch the file.
			if _, err := file.Seek(0, seek); err != nil {
				file.Close()
				watcher.Close()
				return nil, fmt.Errorf("tail: failed to seek: %w", err)
			}
			if err := watcher.Add(t.filename); err != nil {
				file.Close()
				watcher.Close()
				return nil, fmt.Errorf("tail: failed to watch fsnotify event: %w", err)
			}
			ctx, cancel := context.WithCancel(t.ctx)
			return &tail{
				parent:  t,
				file:    file,
				reader:  bufio.NewReader(file),
				watcher: watcher,
				ctx:     ctx,
				cancel:  cancel,
			}, nil
		}

		// fail. retry...
		seek = os.SEEK_SET
		timer := time.NewTimer(openRetryInterval)
		select {
		case <-t.ctx.Done():
			timer.Stop()
			return nil, t.ctx.Err()
		case <-timer.C:
		}
	}
}

// runFile tails target files
func (t *Tail) runFile(seek int) {
	defer t.wg.Done()
	child, err := t.open(seek)
	if err != nil {
		if t.ctx.Err() != nil {
			// stopping tailing now. suppress the error.
			return
		}
		t.errors <- err
		return
	}
	t.wg.Add(1)
	go child.runFile()
}

// runFile tails a file
func (t *tail) runFile() {
	defer t.parent.wg.Done()
	defer t.watcher.Close()
	defer t.cancel()

	cherr := make(chan error, 1)
	ch := make(chan struct{}, 1)
	defer close(ch)

	t.parent.wg.Add(1)
	go func() {
		defer t.parent.wg.Done()
		for {
			if err := t.restrict(); err != nil {
				select {
				case cherr <- err:
				case <-t.ctx.Done():
				}
				return
			}
			err := t.tail()
			if err == nil {
				continue
			}

			// wait for writing new lines.
			select {
			case cherr <- err:
			case <-t.ctx.Done():
				return
			}
			select {
			case _, ok := <-ch:
				if !ok {
					return
				}
			case <-t.ctx.Done():
				return
			}
		}
	}()

	var renamed bool
	var waiting bool // waiting for writing new lines?
	for {
		select {
		case event := <-t.watcher.Events:
			if (event.Op & fsnotify.Remove) != 0 {
				// the target file is removed, stop tailing.
				return
			}
			if (event.Op & fsnotify.Rename) != 0 {
				// log rotation is detected.
				if !renamed {
					// start to watch creating new file.
					t.parent.wg.Add(1)
					go t.parent.runFile(os.SEEK_SET)

					// wait a little, and stop tailing old file.
					go func() {
						timer := time.NewTimer(tailOldFileDelay)
						defer timer.Stop()
						select {
						case <-timer.C:
							t.cancel()
						case <-t.ctx.Done():
						}
					}()
				}
				name, err := getFileName(t.file)
				if err != nil {
					t.parent.errors <- err
					return
				}
				if err := t.watcher.Add(name); err != nil {
					t.parent.errors <- fmt.Errorf("tail: failed to watch fsnotify event: %w", err)
					return
				}
				renamed = true
			}

			// notify new lines are wrote.
			if waiting {
				ch <- struct{}{}
				waiting = false
			}
		case err := <-cherr:
			if errors.Is(err, io.EOF) {
				waiting = true
			} else {
				t.parent.errors <- err
				return
			}
		case err := <-t.watcher.Errors:
			t.parent.errors <- err
			return
		case <-t.ctx.Done():
			return
		}
	}
}

// runReader tails io.Reader
func (t *tail) runReader() {
	defer t.parent.wg.Done()
	defer t.cancel()
	err := t.tail()
	if errors.Is(err, io.EOF) || errors.Is(err, io.ErrClosedPipe) {
		return
	}
	if err != nil {
		t.parent.errors <- err
		return
	}
}

// restrict detects a file that is truncated
func (t *tail) restrict() error {
	stat, err := t.file.Stat()
	if err != nil {
		return fmt.Errorf("tail: failed to stat the file: %w", err)
	}
	pos, err := t.file.Seek(0, os.SEEK_CUR)
	if err != nil {
		return fmt.Errorf("tail: failed to seek: %w", err)
	}
	if stat.Size() < pos {
		// file is truncated. seek to head of file.
		_, err := t.file.Seek(0, os.SEEK_SET)
		if err != nil {
			return fmt.Errorf("tail: failed to seek: %w", err)
		}
	}
	return nil
}

// tail reads lines until EOF
func (t *tail) tail() error {
	for {
		line, err := t.reader.ReadString('\n')
		if err != nil {
			t.buf += line
			return fmt.Errorf("tail: failed to read the file: %w", err)
		}
		t.parent.lines <- &Line{t.buf + line, time.Now()}
		t.buf = ""
	}
}
<commit_msg>optimize memory allocation<commit_after>package tail

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	fsnotify "github.com/fsnotify/fsnotify"
)

const (
	openRetryInterval = time.Second
	tailOldFileDelay  = 15 * time.Second
	linesCapacity     = 1024
	errorsCapacity    = 16
)

// Line is a line of the target file.
type Line struct {
	Text string
	Time time.Time
}

// Tail tails a file.
type Tail struct {
	Lines  <-chan *Line
	Errors <-chan error

	lines    chan<- *Line
	errors   chan<- error
	filename string
	wg       sync.WaitGroup
	ctx      context.Context
	cancel   context.CancelFunc
}

type tail struct {
	parent *Tail

	file    *os.File
	reader  *bufio.Reader
	watcher *fsnotify.Watcher
	buf     strings.Builder
	ctx     context.Context
	cancel  context.CancelFunc
}

// NewTailFile starts tailing a file
func NewTailFile(filename string) (*Tail, error) {
	filename, err := filepath.Abs(filename)
	if err != nil {
		return nil, fmt.Errorf("tail: failed to get the absolute path: %w", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	lines := make(chan *Line, linesCapacity)
	errs := make(chan error, errorsCapacity)
	parent := &Tail{
		Lines:    lines,
		Errors:   errs,
		lines:    lines,
		errors:   errs,
		filename: filename,
		ctx:      ctx,
		cancel:   cancel,
	}
	parent.wg.Add(1)
	go parent.runFile(os.SEEK_END)
	go parent.wait()
	return parent, nil
}

// NewTailReader starts tailing io.Reader
func NewTailReader(reader io.Reader) (*Tail, error) {
	ctx, cancel := context.WithCancel(context.Background())
	lines := make(chan *Line, linesCapacity)
	errs := make(chan error, errorsCapacity)
	parent := &Tail{
		Lines:  lines,
		Errors: errs,
		lines:  lines,
		errors: errs,
		ctx:    ctx,
		cancel: cancel,
	}
	t := &tail{
		parent: parent,
		reader: bufio.NewReader(reader),
		ctx:    ctx,
		cancel: cancel,
	}
	parent.wg.Add(1)
	go t.runReader()
	go parent.wait()
	return parent, nil
}

// Close stops tailing the file.
func (t *Tail) Close() error {
	t.cancel()
	t.wg.Wait()
	return nil
}

func (t *Tail) wait() {
	t.wg.Wait()
	close(t.errors)
	close(t.lines)
}

// open opens the target file.
// If it does not exist, wait for creating new file.
func (t *Tail) open(seek int) (*tail, error) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, fmt.Errorf("tail: failed to initialize fsnotify: %w", err)
	}
	for {
		file, err := os.Open(t.filename)
		if err == nil {
			// success, seek and watch the file.
			if _, err := file.Seek(0, seek); err != nil {
				file.Close()
				watcher.Close()
				return nil, fmt.Errorf("tail: failed to seek: %w", err)
			}
			if err := watcher.Add(t.filename); err != nil {
				file.Close()
				watcher.Close()
				return nil, fmt.Errorf("tail: failed to watch fsnotify event: %w", err)
			}
			ctx, cancel := context.WithCancel(t.ctx)
			return &tail{
				parent:  t,
				file:    file,
				reader:  bufio.NewReader(file),
				watcher: watcher,
				ctx:     ctx,
				cancel:  cancel,
			}, nil
		}

		// fail. retry...
		seek = os.SEEK_SET
		timer := time.NewTimer(openRetryInterval)
		select {
		case <-t.ctx.Done():
			timer.Stop()
			return nil, t.ctx.Err()
		case <-timer.C:
		}
	}
}

// runFile tails target files
func (t *Tail) runFile(seek int) {
	defer t.wg.Done()
	child, err := t.open(seek)
	if err != nil {
		if t.ctx.Err() != nil {
			// stopping tailing now. suppress the error.
			return
		}
		t.errors <- err
		return
	}
	t.wg.Add(1)
	go child.runFile()
}

// runFile tails a file
func (t *tail) runFile() {
	defer t.parent.wg.Done()
	defer t.watcher.Close()
	defer t.cancel()

	cherr := make(chan error, 1)
	ch := make(chan struct{}, 1)
	defer close(ch)

	t.parent.wg.Add(1)
	go func() {
		defer t.parent.wg.Done()
		for {
			if err := t.restrict(); err != nil {
				select {
				case cherr <- err:
				case <-t.ctx.Done():
				}
				return
			}
			err := t.tail()
			if err == nil {
				continue
			}

			// wait for writing new lines.
			select {
			case cherr <- err:
			case <-t.ctx.Done():
				return
			}
			select {
			case _, ok := <-ch:
				if !ok {
					return
				}
			case <-t.ctx.Done():
				return
			}
		}
	}()

	var renamed bool
	var waiting bool // waiting for writing new lines?
	for {
		select {
		case event := <-t.watcher.Events:
			if (event.Op & fsnotify.Remove) != 0 {
				// the target file is removed, stop tailing.
				return
			}
			if (event.Op & fsnotify.Rename) != 0 {
				// log rotation is detected.
				if !renamed {
					// start to watch creating new file.
					t.parent.wg.Add(1)
					go t.parent.runFile(os.SEEK_SET)

					// wait a little, and stop tailing old file.
					go func() {
						timer := time.NewTimer(tailOldFileDelay)
						defer timer.Stop()
						select {
						case <-timer.C:
							t.cancel()
						case <-t.ctx.Done():
						}
					}()
				}
				name, err := getFileName(t.file)
				if err != nil {
					t.parent.errors <- err
					return
				}
				if err := t.watcher.Add(name); err != nil {
					t.parent.errors <- fmt.Errorf("tail: failed to watch fsnotify event: %w", err)
					return
				}
				renamed = true
			}

			// notify new lines are wrote.
			if waiting {
				ch <- struct{}{}
				waiting = false
			}
		case err := <-cherr:
			if errors.Is(err, io.EOF) {
				waiting = true
			} else {
				t.parent.errors <- err
				return
			}
		case err := <-t.watcher.Errors:
			t.parent.errors <- err
			return
		case <-t.ctx.Done():
			return
		}
	}
}

// runReader tails io.Reader
func (t *tail) runReader() {
	defer t.parent.wg.Done()
	defer t.cancel()
	err := t.tail()
	if errors.Is(err, io.EOF) || errors.Is(err, io.ErrClosedPipe) {
		return
	}
	if err != nil {
		t.parent.errors <- err
		return
	}
}

// restrict detects a file that is truncated
func (t *tail) restrict() error {
	stat, err := t.file.Stat()
	if err != nil {
		return fmt.Errorf("tail: failed to stat the file: %w", err)
	}
	pos, err := t.file.Seek(0, os.SEEK_CUR)
	if err != nil {
		return fmt.Errorf("tail: failed to seek: %w", err)
	}
	if stat.Size() < pos {
		// file is truncated. seek to head of file.
		_, err := t.file.Seek(0, os.SEEK_SET)
		if err != nil {
			return fmt.Errorf("tail: failed to seek: %w", err)
		}
	}
	return nil
}

// tail reads lines until EOF
func (t *tail) tail() error {
	for {
		line, err := t.reader.ReadSlice('\n')
		t.buf.Write(line)
		if err == bufio.ErrBufferFull {
			// the reader cannot find EOL in its buffer.
			// continue to read a line.
			continue
		}
		if err != nil {
			return fmt.Errorf("tail: failed to read the file: %w", err)
		}
		t.parent.lines <- &Line{t.buf.String(), time.Now()}
		t.buf.Reset()
	}
}
<|endoftext|>