<commit_before>package autorollerv2

import (
	"context"
	"path"
	"sync"
	"time"

	"go.skia.org/infra/autoroll/go/autoroll_modes"
	"go.skia.org/infra/autoroll/go/autoroller"
	"go.skia.org/infra/autoroll/go/recent_rolls"
	"go.skia.org/infra/autoroll/go/repo_manager"
	"go.skia.org/infra/autoroll/go/state_machine"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

// AutoRoller is a struct which automates the merging new revisions of one
// project into another.
type AutoRoller struct {
	cqExtraTrybots  string
	currentRoll     RollImpl
	emails          []string
	emailsMtx       sync.RWMutex
	gerrit          *gerrit.Gerrit
	lastError       error
	liveness        metrics2.Liveness
	modeHistory     *autoroll_modes.ModeHistory
	modeMtx         sync.Mutex
	recent          *recent_rolls.RecentRolls
	retrieveRoll    func(*AutoRoller, int64) (RollImpl, error)
	rm              repo_manager.RepoManager
	runningMtx      sync.Mutex
	sm              *state_machine.AutoRollStateMachine
	status          *autoroller.AutoRollStatusCache
	rollIntoAndroid bool
}

// newAutoRoller returns an AutoRoller instance.
func newAutoRoller(workdir, childPath, cqExtraTrybots string, emails []string, gerrit *gerrit.Gerrit, rm repo_manager.RepoManager, retrieveRoll func(*AutoRoller, int64) (RollImpl, error)) (*AutoRoller, error) {
	recent, err := recent_rolls.NewRecentRolls(path.Join(workdir, "recent_rolls.db"))
	if err != nil {
		return nil, err
	}

	mh, err := autoroll_modes.NewModeHistory(path.Join(workdir, "autoroll_modes.db"))
	if err != nil {
		return nil, err
	}

	arb := &AutoRoller{
		cqExtraTrybots: cqExtraTrybots,
		emails:         emails,
		gerrit:         gerrit,
		liveness:       metrics2.NewLiveness("last-autoroll-landed", map[string]string{"child-path": childPath}),
		modeHistory:    mh,
		recent:         recent,
		retrieveRoll:   retrieveRoll,
		rm:             rm,
		status:         &autoroller.AutoRollStatusCache{},
	}
	sm, err := state_machine.New(arb, workdir)
	if err != nil {
		return nil, err
	}
	arb.sm = sm
	current := recent.CurrentRoll()
	if current != nil {
		roll, err := arb.retrieveRoll(arb, current.Issue)
		if err != nil {
			return nil, err
		}
		arb.currentRoll = roll
	}
	return arb, nil
}

// NewAndroidAutoRoller returns an AutoRoller instance which rolls into Android.
func NewAndroidAutoRoller(workdir, parentBranch, childPath, childBranch, cqExtraTrybots string, emails []string, gerrit *gerrit.Gerrit, strategy string) (*AutoRoller, error) {
	rm, err := repo_manager.NewAndroidRepoManager(workdir, parentBranch, childPath, childBranch, gerrit, strategy)
	if err != nil {
		return nil, err
	}
	retrieveRoll := func(arb *AutoRoller, issue int64) (RollImpl, error) {
		return newGerritAndroidRoll(arb.gerrit, arb.rm, arb.recent, issue)
	}
	return newAutoRoller(workdir, childPath, cqExtraTrybots, emails, gerrit, rm, retrieveRoll)
}

// NewDEPSAutoRoller returns an AutoRoller instance which rolls using DEPS.
func NewDEPSAutoRoller(workdir, parentRepo, parentBranch, childPath, childBranch, cqExtraTrybots string, emails []string, gerrit *gerrit.Gerrit, depot_tools string, strategy string) (*AutoRoller, error) {
	rm, err := repo_manager.NewDEPSRepoManager(workdir, parentRepo, parentBranch, childPath, childBranch, depot_tools, gerrit, strategy)
	if err != nil {
		return nil, err
	}
	retrieveRoll := func(arb *AutoRoller, issue int64) (RollImpl, error) {
		return newGerritRoll(arb.gerrit, arb.rm, arb.recent, issue)
	}
	return newAutoRoller(workdir, childPath, cqExtraTrybots, emails, gerrit, rm, retrieveRoll)
}

// NewManifestAutoRoller returns an AutoRoller instance which rolls using DEPS.
func NewManifestAutoRoller(workdir, parentRepo, parentBranch, childPath, childBranch, cqExtraTrybots string, emails []string, gerrit *gerrit.Gerrit, depot_tools string, strategy string) (*AutoRoller, error) {
	rm, err := repo_manager.NewManifestRepoManager(workdir, parentRepo, parentBranch, childPath, childBranch, depot_tools, gerrit, strategy)
	if err != nil {
		return nil, err
	}
	retrieveRoll := func(arb *AutoRoller, issue int64) (RollImpl, error) {
		return newGerritRoll(arb.gerrit, arb.rm, arb.recent, issue)
	}
	return newAutoRoller(workdir, childPath, cqExtraTrybots, emails, gerrit, rm, retrieveRoll)
}

// Start initiates the AutoRoller's loop.
func (r *AutoRoller) Start(tickFrequency, repoFrequency time.Duration, ctx context.Context) {
	sklog.Infof("Starting autoroller.")
	repo_manager.Start(r.rm, repoFrequency, ctx)
	lv := metrics2.NewLiveness("last-successful-autoroll-tick")
	go util.RepeatCtx(tickFrequency, ctx, func() {
		if err := r.Tick(); err != nil {
			sklog.Errorf("Failed to run autoroll: %s", err)
		} else {
			lv.Reset()
		}
	})
	go func() {
		for {
			select {
			case <-ctx.Done():
				util.LogErr(r.recent.Close())
				util.LogErr(r.modeHistory.Close())
			default:
			}
		}
	}()
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) GetActiveRoll() state_machine.RollCLImpl {
	return r.currentRoll
}

// GetEmails returns the list of email addresses which are copied on rolls.
func (r *AutoRoller) GetEmails() []string {
	r.emailsMtx.RLock()
	defer r.emailsMtx.RUnlock()
	rv := make([]string, len(r.emails))
	copy(rv, r.emails)
	return rv
}

// SetEmails sets the list of email addresses which are copied on rolls.
func (r *AutoRoller) SetEmails(e []string) {
	r.emailsMtx.Lock()
	defer r.emailsMtx.Unlock()
	emails := make([]string, len(e))
	copy(emails, e)
	r.emails = emails
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) GetMode() string {
	r.modeMtx.Lock()
	defer r.modeMtx.Unlock()
	return r.modeHistory.CurrentMode().Mode
}

// SetMode sets the desired mode of the bot. This forces the bot to run and
// blocks until it finishes.
func (r *AutoRoller) SetMode(m, user, message string) error {
	r.modeMtx.Lock()
	defer r.modeMtx.Unlock()
	if err := r.modeHistory.Add(m, user, message); err != nil {
		return err
	}
	return r.Tick()
}

// Return the roll-up status of the bot.
func (r *AutoRoller) GetStatus(includeError bool) *autoroller.AutoRollStatus {
	return r.status.Get(includeError)
}

// Return the AutoRoll user.
func (r *AutoRoller) GetUser() string {
	return r.rm.User()
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) UploadNewRoll(from, to string, dryRun bool) error {
	issueNum, err := r.rm.CreateNewRoll(from, to, r.GetEmails(), r.cqExtraTrybots, dryRun)
	if err != nil {
		return err
	}
	roll, err := r.retrieveRoll(r, issueNum)
	if err != nil {
		return err
	}
	if err := roll.InsertIntoDB(); err != nil {
		return err
	}
	r.currentRoll = roll
	return nil
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) GetCurrentRev() string {
	return r.rm.LastRollRev()
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) GetNextRollRev() string {
	return r.rm.NextRollRev()
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) RolledPast(rev string) (bool, error) {
	return r.rm.RolledPast(rev)
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) UpdateRepos() error {
	return r.rm.Update()
}

// Run one iteration of the roller.
func (r *AutoRoller) Tick() error {
	sklog.Infof("Running autoroller.")
	// Run the state machine.
	lastErr := r.sm.NextTransitionSequence()

	// Update the status information.
	lastErrorStr := ""
	if lastErr != nil {
		lastErrorStr = lastErr.Error()
	}
	if err := r.status.Set(&autoroller.AutoRollStatus{
		CurrentRoll: r.recent.CurrentRoll(),
		Error:       lastErrorStr,
		GerritUrl:   r.gerrit.Url(0),
		LastRoll:    r.recent.LastRoll(),
		LastRollRev: r.rm.LastRollRev(),
		Mode:        r.modeHistory.CurrentMode(),
		Recent:      r.recent.GetRecentRolls(),
		Status:      string(r.sm.Current()),
	}); err != nil {
		return err
	}
	sklog.Infof("Autoroller state %s", r.sm.Current())
	return lastErr
}
<commit_msg>[autoroll] Fix mode change mutex<commit_after>package autorollerv2

import (
	"context"
	"path"
	"sync"
	"time"

	"go.skia.org/infra/autoroll/go/autoroll_modes"
	"go.skia.org/infra/autoroll/go/autoroller"
	"go.skia.org/infra/autoroll/go/recent_rolls"
	"go.skia.org/infra/autoroll/go/repo_manager"
	"go.skia.org/infra/autoroll/go/state_machine"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

// AutoRoller is a struct which automates the merging new revisions of one
// project into another.
type AutoRoller struct {
	cqExtraTrybots  string
	currentRoll     RollImpl
	emails          []string
	emailsMtx       sync.RWMutex
	gerrit          *gerrit.Gerrit
	lastError       error
	liveness        metrics2.Liveness
	modeHistory     *autoroll_modes.ModeHistory
	recent          *recent_rolls.RecentRolls
	retrieveRoll    func(*AutoRoller, int64) (RollImpl, error)
	rm              repo_manager.RepoManager
	runningMtx      sync.Mutex
	sm              *state_machine.AutoRollStateMachine
	status          *autoroller.AutoRollStatusCache
	rollIntoAndroid bool
}

// newAutoRoller returns an AutoRoller instance.
func newAutoRoller(workdir, childPath, cqExtraTrybots string, emails []string, gerrit *gerrit.Gerrit, rm repo_manager.RepoManager, retrieveRoll func(*AutoRoller, int64) (RollImpl, error)) (*AutoRoller, error) {
	recent, err := recent_rolls.NewRecentRolls(path.Join(workdir, "recent_rolls.db"))
	if err != nil {
		return nil, err
	}

	mh, err := autoroll_modes.NewModeHistory(path.Join(workdir, "autoroll_modes.db"))
	if err != nil {
		return nil, err
	}

	arb := &AutoRoller{
		cqExtraTrybots: cqExtraTrybots,
		emails:         emails,
		gerrit:         gerrit,
		liveness:       metrics2.NewLiveness("last-autoroll-landed", map[string]string{"child-path": childPath}),
		modeHistory:    mh,
		recent:         recent,
		retrieveRoll:   retrieveRoll,
		rm:             rm,
		status:         &autoroller.AutoRollStatusCache{},
	}
	sm, err := state_machine.New(arb, workdir)
	if err != nil {
		return nil, err
	}
	arb.sm = sm
	current := recent.CurrentRoll()
	if current != nil {
		roll, err := arb.retrieveRoll(arb, current.Issue)
		if err != nil {
			return nil, err
		}
		arb.currentRoll = roll
	}
	return arb, nil
}

// NewAndroidAutoRoller returns an AutoRoller instance which rolls into Android.
func NewAndroidAutoRoller(workdir, parentBranch, childPath, childBranch, cqExtraTrybots string, emails []string, gerrit *gerrit.Gerrit, strategy string) (*AutoRoller, error) {
	rm, err := repo_manager.NewAndroidRepoManager(workdir, parentBranch, childPath, childBranch, gerrit, strategy)
	if err != nil {
		return nil, err
	}
	retrieveRoll := func(arb *AutoRoller, issue int64) (RollImpl, error) {
		return newGerritAndroidRoll(arb.gerrit, arb.rm, arb.recent, issue)
	}
	return newAutoRoller(workdir, childPath, cqExtraTrybots, emails, gerrit, rm, retrieveRoll)
}

// NewDEPSAutoRoller returns an AutoRoller instance which rolls using DEPS.
func NewDEPSAutoRoller(workdir, parentRepo, parentBranch, childPath, childBranch, cqExtraTrybots string, emails []string, gerrit *gerrit.Gerrit, depot_tools string, strategy string) (*AutoRoller, error) {
	rm, err := repo_manager.NewDEPSRepoManager(workdir, parentRepo, parentBranch, childPath, childBranch, depot_tools, gerrit, strategy)
	if err != nil {
		return nil, err
	}
	retrieveRoll := func(arb *AutoRoller, issue int64) (RollImpl, error) {
		return newGerritRoll(arb.gerrit, arb.rm, arb.recent, issue)
	}
	return newAutoRoller(workdir, childPath, cqExtraTrybots, emails, gerrit, rm, retrieveRoll)
}

// NewManifestAutoRoller returns an AutoRoller instance which rolls using DEPS.
func NewManifestAutoRoller(workdir, parentRepo, parentBranch, childPath, childBranch, cqExtraTrybots string, emails []string, gerrit *gerrit.Gerrit, depot_tools string, strategy string) (*AutoRoller, error) {
	rm, err := repo_manager.NewManifestRepoManager(workdir, parentRepo, parentBranch, childPath, childBranch, depot_tools, gerrit, strategy)
	if err != nil {
		return nil, err
	}
	retrieveRoll := func(arb *AutoRoller, issue int64) (RollImpl, error) {
		return newGerritRoll(arb.gerrit, arb.rm, arb.recent, issue)
	}
	return newAutoRoller(workdir, childPath, cqExtraTrybots, emails, gerrit, rm, retrieveRoll)
}

// Start initiates the AutoRoller's loop.
func (r *AutoRoller) Start(tickFrequency, repoFrequency time.Duration, ctx context.Context) {
	sklog.Infof("Starting autoroller.")
	repo_manager.Start(r.rm, repoFrequency, ctx)
	lv := metrics2.NewLiveness("last-successful-autoroll-tick")
	go util.RepeatCtx(tickFrequency, ctx, func() {
		if err := r.Tick(); err != nil {
			sklog.Errorf("Failed to run autoroll: %s", err)
		} else {
			lv.Reset()
		}
	})
	go func() {
		for {
			select {
			case <-ctx.Done():
				util.LogErr(r.recent.Close())
				util.LogErr(r.modeHistory.Close())
			default:
			}
		}
	}()
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) GetActiveRoll() state_machine.RollCLImpl {
	return r.currentRoll
}

// GetEmails returns the list of email addresses which are copied on rolls.
func (r *AutoRoller) GetEmails() []string {
	r.emailsMtx.RLock()
	defer r.emailsMtx.RUnlock()
	rv := make([]string, len(r.emails))
	copy(rv, r.emails)
	return rv
}

// SetEmails sets the list of email addresses which are copied on rolls.
func (r *AutoRoller) SetEmails(e []string) {
	r.emailsMtx.Lock()
	defer r.emailsMtx.Unlock()
	emails := make([]string, len(e))
	copy(emails, e)
	r.emails = emails
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) GetMode() string {
	return r.modeHistory.CurrentMode().Mode
}

// SetMode sets the desired mode of the bot. This forces the bot to run and
// blocks until it finishes.
func (r *AutoRoller) SetMode(m, user, message string) error {
	if err := r.modeHistory.Add(m, user, message); err != nil {
		return err
	}
	return r.Tick()
}

// Return the roll-up status of the bot.
func (r *AutoRoller) GetStatus(includeError bool) *autoroller.AutoRollStatus {
	return r.status.Get(includeError)
}

// Return the AutoRoll user.
func (r *AutoRoller) GetUser() string {
	return r.rm.User()
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) UploadNewRoll(from, to string, dryRun bool) error {
	issueNum, err := r.rm.CreateNewRoll(from, to, r.GetEmails(), r.cqExtraTrybots, dryRun)
	if err != nil {
		return err
	}
	roll, err := r.retrieveRoll(r, issueNum)
	if err != nil {
		return err
	}
	if err := roll.InsertIntoDB(); err != nil {
		return err
	}
	r.currentRoll = roll
	return nil
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) GetCurrentRev() string {
	return r.rm.LastRollRev()
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) GetNextRollRev() string {
	return r.rm.NextRollRev()
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) RolledPast(rev string) (bool, error) {
	return r.rm.RolledPast(rev)
}

// See documentation for state_machine.AutoRollerImpl interface.
func (r *AutoRoller) UpdateRepos() error {
	return r.rm.Update()
}

// Run one iteration of the roller.
func (r *AutoRoller) Tick() error {
	r.runningMtx.Lock()
	defer r.runningMtx.Unlock()

	sklog.Infof("Running autoroller.")
	// Run the state machine.
	lastErr := r.sm.NextTransitionSequence()

	// Update the status information.
	lastErrorStr := ""
	if lastErr != nil {
		lastErrorStr = lastErr.Error()
	}
	if err := r.status.Set(&autoroller.AutoRollStatus{
		CurrentRoll: r.recent.CurrentRoll(),
		Error:       lastErrorStr,
		GerritUrl:   r.gerrit.Url(0),
		LastRoll:    r.recent.LastRoll(),
		LastRollRev: r.rm.LastRollRev(),
		Mode:        r.modeHistory.CurrentMode(),
		Recent:      r.recent.GetRecentRolls(),
		Status:      string(r.sm.Current()),
	}); err != nil {
		return err
	}
	sklog.Infof("Autoroller state %s", r.sm.Current())
	return lastErr
}
<|endoftext|>


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

package cloud

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/pkg/errors"
)

const (
	// Default request timeout
	RequestTimeout = 10 * time.Second
	// Retry interval
	RetryInterval = 500 * time.Millisecond
	// Retry attempts
	MaxRetries = 3
)

// Client handles communication with Load Impact cloud API.
type Client struct {
	client  *http.Client
	token   string
	baseURL string
	version string

	retries       int
	retryInterval time.Duration
}

func NewClient(token, host, version string) *Client {
	c := &Client{
		client:        &http.Client{Timeout: RequestTimeout},
		token:         token,
		baseURL:       fmt.Sprintf("%s/v1", host),
		version:       version,
		retries:       MaxRetries,
		retryInterval: RetryInterval,
	}
	return c
}

func (c *Client) NewRequest(method, url string, data interface{}) (*http.Request, error) {
	var buf io.Reader

	if data != nil {
		b, err := json.Marshal(&data)
		if err != nil {
			return nil, err
		}

		buf = bytes.NewBuffer(b)
	}

	return http.NewRequest(method, url, buf)
}

func (c *Client) Do(req *http.Request, v interface{}) error {
	var originalBody []byte
	var err error

	if req.Body != nil {
		originalBody, err = ioutil.ReadAll(req.Body)
		if err != nil {
			return err
		}

		if cerr := req.Body.Close(); cerr != nil {
			err = cerr
		}
	}

	for i := 1; i <= c.retries; i++ {
		if len(originalBody) > 0 {
			req.Body = ioutil.NopCloser(bytes.NewBuffer(originalBody))
		}

		retry, err := c.do(req, v, i)

		if retry {
			time.Sleep(c.retryInterval)
			continue
		}

		return err
	}

	return err
}

func (c *Client) do(req *http.Request, v interface{}, attempt int) (retry bool, err error) {
	if req.Header.Get("Content-Type") == "" {
		req.Header.Set("Content-Type", "application/json")
	}
	if c.token != "" {
		req.Header.Set("Authorization", fmt.Sprintf("Token %s", c.token))
	}
	req.Header.Set("User-Agent", "k6cloud/"+c.version)
	resp, err := c.client.Do(req)

	defer func() {
		if resp != nil {
			if cerr := resp.Body.Close(); cerr != nil && err == nil {
				err = cerr
			}
		}
	}()

	if shouldRetry(resp, err, attempt, c.retries) {
		return true, err
	}

	if err != nil {
		return false, err
	}

	if err = checkResponse(resp); err != nil {
		return false, err
	}

	if v != nil {
		if err = json.NewDecoder(resp.Body).Decode(v); err == io.EOF {
			err = nil // Ignore EOF from empty body
		}
	}

	return false, err
}

func checkResponse(r *http.Response) error {
	if r == nil {
		return ErrUnknown
	}

	if c := r.StatusCode; c >= 200 && c <= 299 {
		return nil
	}

	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}

	var payload struct {
		Error ErrorResponse `json:"error"`
	}
	if err := json.Unmarshal(data, &payload); err != nil {
		if r.StatusCode == http.StatusUnauthorized {
			return ErrNotAuthenticated
		}
		if r.StatusCode == http.StatusForbidden {
			return ErrNotAuthorized
		}
		return errors.Errorf(
			"Unexpected HTTP error from %s: %d %s",
			r.Request.URL,
			r.StatusCode,
			http.StatusText(r.StatusCode),
		)
	}
	payload.Error.Response = r
	return payload.Error
}

func shouldRetry(resp *http.Response, err error, attempt, maxAttempts int) bool {
	if attempt >= maxAttempts {
		return false
	}

	if resp == nil || err != nil {
		return true
	}

	if resp.StatusCode >= 500 || resp.StatusCode == 429 {
		return true
	}

	return false
}
<commit_msg>stats/cloud: increase default request timeout<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2017 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cloud

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/pkg/errors"
)

const (
	// RequestTimeout is the default cloud request timeout
	RequestTimeout = 20 * time.Second
	// RetryInterval is the default cloud request retry interval
	RetryInterval = 500 * time.Millisecond
	// MaxRetries specifies max retry attempts
	MaxRetries = 3
)

// Client handles communication with Load Impact cloud API.
type Client struct {
	client  *http.Client
	token   string
	baseURL string
	version string

	retries       int
	retryInterval time.Duration
}

func NewClient(token, host, version string) *Client {
	c := &Client{
		client:        &http.Client{Timeout: RequestTimeout},
		token:         token,
		baseURL:       fmt.Sprintf("%s/v1", host),
		version:       version,
		retries:       MaxRetries,
		retryInterval: RetryInterval,
	}
	return c
}

func (c *Client) NewRequest(method, url string, data interface{}) (*http.Request, error) {
	var buf io.Reader

	if data != nil {
		b, err := json.Marshal(&data)
		if err != nil {
			return nil, err
		}

		buf = bytes.NewBuffer(b)
	}

	return http.NewRequest(method, url, buf)
}

func (c *Client) Do(req *http.Request, v interface{}) error {
	var originalBody []byte
	var err error

	if req.Body != nil {
		originalBody, err = ioutil.ReadAll(req.Body)
		if err != nil {
			return err
		}

		if cerr := req.Body.Close(); cerr != nil {
			err = cerr
		}
	}

	for i := 1; i <= c.retries; i++ {
		if len(originalBody) > 0 {
			req.Body = ioutil.NopCloser(bytes.NewBuffer(originalBody))
		}

		retry, err := c.do(req, v, i)

		if retry {
			time.Sleep(c.retryInterval)
			continue
		}

		return err
	}

	return err
}

func (c *Client) do(req *http.Request, v interface{}, attempt int) (retry bool, err error) {
	if req.Header.Get("Content-Type") == "" {
		req.Header.Set("Content-Type", "application/json")
	}
	if c.token != "" {
		req.Header.Set("Authorization", fmt.Sprintf("Token %s", c.token))
	}
	req.Header.Set("User-Agent", "k6cloud/"+c.version)
	resp, err := c.client.Do(req)

	defer func() {
		if resp != nil {
			if cerr := resp.Body.Close(); cerr != nil && err == nil {
				err = cerr
			}
		}
	}()

	if shouldRetry(resp, err, attempt, c.retries) {
		return true, err
	}

	if err != nil {
		return false, err
	}

	if err = checkResponse(resp); err != nil {
		return false, err
	}

	if v != nil {
		if err = json.NewDecoder(resp.Body).Decode(v); err == io.EOF {
			err = nil // Ignore EOF from empty body
		}
	}

	return false, err
}

func checkResponse(r *http.Response) error {
	if r == nil {
		return ErrUnknown
	}

	if c := r.StatusCode; c >= 200 && c <= 299 {
		return nil
	}

	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}

	var payload struct {
		Error ErrorResponse `json:"error"`
	}
	if err := json.Unmarshal(data, &payload); err != nil {
		if r.StatusCode == http.StatusUnauthorized {
			return ErrNotAuthenticated
		}
		if r.StatusCode == http.StatusForbidden {
			return ErrNotAuthorized
		}
		return errors.Errorf(
			"Unexpected HTTP error from %s: %d %s",
			r.Request.URL,
			r.StatusCode,
			http.StatusText(r.StatusCode),
		)
	}
	payload.Error.Response = r
	return payload.Error
}

func shouldRetry(resp *http.Response, err error, attempt, maxAttempts int) bool {
	if attempt >= maxAttempts {
		return false
	}

	if resp == nil || err != nil {
		return true
	}

	if resp.StatusCode >= 500 || resp.StatusCode == 429 {
		return true
	}

	return false
}
<|endoftext|>


<commit_before>// Copyright 2016 Stratumn SAS. All rights reserved.
// Use of this source code is governed by an Apache License 2.0
// that can be found in the LICENSE file.

// Package bcbatchfossilizer implements a fossilizer that fossilize batches of hashes on a blockchain.
package bcbatchfossilizer

import (
	"fmt"

	log "github.com/sirupsen/logrus"

	"github.com/stratumn/go/fossilizer"
	"github.com/stratumn/go/types"

	"github.com/stratumn/goprivate/batchfossilizer"
	"github.com/stratumn/goprivate/blockchain"
)

const (
	// Name is the name set in the fossilizer's information.
	Name = "bcbatch"

	// Description is the description set in the fossilizer's information.
	Description = "Stratumn Blockchain Batch Fossilizer"
)

// Config contains configuration options for the fossilizer.
type Config struct {
	HashTimestamper blockchain.HashTimestamper
}

// Info is the info returned by GetInfo.
type Info struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	Version     string `json:"version"`
	Commit      string `json:"commit"`
	Blockchain  string `json:"blockchain"`
}

// Evidence is the evidence sent to the result channel.
type Evidence struct {
	*batchfossilizer.Evidence
	TransactionID blockchain.TransactionID `json:"txid"`
}

// Fossilizer is the type that implements github.com/stratumn/go/fossilizer.Adapter.
type Fossilizer struct {
	*batchfossilizer.Fossilizer
	config            *Config
	lastRoot          *types.Bytes32
	lastTransactionID blockchain.TransactionID
}

// New creates an instance of a Fossilizer.
func New(config *Config, batchConfig *batchfossilizer.Config) (*Fossilizer, error) {
	if batchConfig.MaxSimBatches > 1 {
		return nil, fmt.Errorf("MaxSimBatches is %d want less than 2", batchConfig.MaxSimBatches)
	}

	b, err := batchfossilizer.New(batchConfig)
	if err != nil {
		return nil, err
	}

	f := Fossilizer{
		Fossilizer: b,
		config:     config,
	}

	f.SetTransformer(f.transform)

	return &f, err
}

// GetInfo implements github.com/stratumn/go/fossilizer.Adapter.GetInfo.
func (a *Fossilizer) GetInfo() (interface{}, error) {
	batchInfo, err := a.Fossilizer.GetInfo()
	if err != nil {
		return nil, err
	}

	info, ok := batchInfo.(*batchfossilizer.Info)
	if !ok {
		return nil, fmt.Errorf("unexpected batchfossilizer info %#v", batchInfo)
	}

	return &Info{
		Name:        Name,
		Description: Description,
		Version:     info.Version,
		Commit:      info.Commit,
		Blockchain:  a.config.HashTimestamper.Network().String(),
	}, nil
}

func (a *Fossilizer) transform(evidence *batchfossilizer.Evidence, data, meta []byte) (*fossilizer.Result, error) {
	var (
		root = evidence.Root
		txid blockchain.TransactionID
		err  error
	)

	if a.lastRoot == nil || *root != *a.lastRoot {
		txid, err = a.config.HashTimestamper.TimestampHash(root)
		if err != nil {
			return nil, err
		}
		log.WithFields(log.Fields{
			"txid": txid,
			"root": root,
		}).Info("Broadcasted transaction")

		a.lastRoot = root
		a.lastTransactionID = txid
	}

	evidenceWrapper := map[string]*Evidence{}
	evidenceWrapper[a.config.HashTimestamper.Network().String()] = &Evidence{
		Evidence:      evidence,
		TransactionID: a.lastTransactionID,
	}

	r := fossilizer.Result{
		Evidence: evidenceWrapper,
		Data:     data,
		Meta:     meta,
	}

	return &r, nil
}
<commit_msg>bcbatchfossilizer: sirupsen -> Sirupsen<commit_after>// Copyright 2016 Stratumn SAS. All rights reserved.
// Use of this source code is governed by an Apache License 2.0
// that can be found in the LICENSE file.

// Package bcbatchfossilizer implements a fossilizer that fossilize batches of hashes on a blockchain.
package bcbatchfossilizer

import (
	"fmt"

	log "github.com/Sirupsen/logrus"

	"github.com/stratumn/go/fossilizer"
	"github.com/stratumn/go/types"

	"github.com/stratumn/goprivate/batchfossilizer"
	"github.com/stratumn/goprivate/blockchain"
)

const (
	// Name is the name set in the fossilizer's information.
	Name = "bcbatch"

	// Description is the description set in the fossilizer's information.
	Description = "Stratumn Blockchain Batch Fossilizer"
)

// Config contains configuration options for the fossilizer.
type Config struct {
	HashTimestamper blockchain.HashTimestamper
}

// Info is the info returned by GetInfo.
type Info struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	Version     string `json:"version"`
	Commit      string `json:"commit"`
	Blockchain  string `json:"blockchain"`
}

// Evidence is the evidence sent to the result channel.
type Evidence struct {
	*batchfossilizer.Evidence
	TransactionID blockchain.TransactionID `json:"txid"`
}

// Fossilizer is the type that implements github.com/stratumn/go/fossilizer.Adapter.
type Fossilizer struct {
	*batchfossilizer.Fossilizer
	config            *Config
	lastRoot          *types.Bytes32
	lastTransactionID blockchain.TransactionID
}

// New creates an instance of a Fossilizer.
func New(config *Config, batchConfig *batchfossilizer.Config) (*Fossilizer, error) {
	if batchConfig.MaxSimBatches > 1 {
		return nil, fmt.Errorf("MaxSimBatches is %d want less than 2", batchConfig.MaxSimBatches)
	}

	b, err := batchfossilizer.New(batchConfig)
	if err != nil {
		return nil, err
	}

	f := Fossilizer{
		Fossilizer: b,
		config:     config,
	}

	f.SetTransformer(f.transform)

	return &f, err
}

// GetInfo implements github.com/stratumn/go/fossilizer.Adapter.GetInfo.
func (a *Fossilizer) GetInfo() (interface{}, error) {
	batchInfo, err := a.Fossilizer.GetInfo()
	if err != nil {
		return nil, err
	}

	info, ok := batchInfo.(*batchfossilizer.Info)
	if !ok {
		return nil, fmt.Errorf("unexpected batchfossilizer info %#v", batchInfo)
	}

	return &Info{
		Name:        Name,
		Description: Description,
		Version:     info.Version,
		Commit:      info.Commit,
		Blockchain:  a.config.HashTimestamper.Network().String(),
	}, nil
}

func (a *Fossilizer) transform(evidence *batchfossilizer.Evidence, data, meta []byte) (*fossilizer.Result, error) {
	var (
		root = evidence.Root
		txid blockchain.TransactionID
		err  error
	)

	if a.lastRoot == nil || *root != *a.lastRoot {
		txid, err = a.config.HashTimestamper.TimestampHash(root)
		if err != nil {
			return nil, err
		}
		log.WithFields(log.Fields{
			"txid": txid,
			"root": root,
		}).Info("Broadcasted transaction")

		a.lastRoot = root
		a.lastTransactionID = txid
	}

	evidenceWrapper := map[string]*Evidence{}
	evidenceWrapper[a.config.HashTimestamper.Network().String()] = &Evidence{
		Evidence:      evidence,
		TransactionID: a.lastTransactionID,
	}

	r := fossilizer.Result{
		Evidence: evidenceWrapper,
		Data:     data,
		Meta:     meta,
	}

	return &r, nil
}
<|endoftext|>


<commit_before>// Copyright 2018 The NATS Authors
// 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 stores

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	mrand "math/rand"
	"os"
	"runtime"
	"strings"
	"sync"

	"github.com/nats-io/go-nats-streaming/pb"
	"golang.org/x/crypto/chacha20poly1305"
)

// CryptStore specific errors
var (
	ErrCryptoStoreRequiresKey = errors.New("encryption key required")
	ErrCipherNotSupported     = errors.New("encryption cipher not supported")
)

const (
	// CryptoStoreEnvKeyName is the environment variable name
	// that the CryptoStore looks up if no key is passed as
	// a parameter.
	CryptoStoreEnvKeyName = "NATS_STREAMING_ENCRYPTION_KEY"

	// CryptoCipherAES is the name of the AES cipher to use for encryption
	CryptoCipherAES = "AES"

	// CryptoCipherChaChaPoly is the name of the ChaChaPoly cipher to use for encryption
	CryptoCipherChaChaPoly = "CHACHA"

	// CryptoCipherAutoSelect if passed to NewCryptoStore() will cause the cipher to
	// be auto-selected based on the platform the executable is built for.
	CryptoCipherAutoSelect = ""
)

// These constants define a code for each of the supported ciphers
const (
	CryptoCodeAES    = byte(1)
	CryptoCodeChaCha = byte(2)
)

// CryptoStore is a store wrapping a store implementation
// and adds encryption support.
type CryptoStore struct {
	sync.Mutex
	Store

	// These are set when the store is created. They are then
	// passed to a CryptoMsgStore so that there is no need
	// to reference back to these.
	// Note that nonceSize and cryptoOverhead are same for
	// those 2 ciphers. If we add more and those are different,
	// will need to be stored differently or call the appropriate
	// gcm.NonceSize() and gcm.Overhead() functions.
	cipherCode     byte
	aesgcm         cipher.AEAD
	chachagcm      cipher.AEAD
	nonceSize      int
	cryptoOverhead int
}

// CryptoMsgStore is a store wrappeing a SubStore implementation
// and adds encryption support.
type CryptoMsgStore struct {
	sync.Mutex
	MsgStore
	cipherCode     byte
	gcm            cipher.AEAD // Use this one to encrypt
	aesgcm         cipher.AEAD // This is to decrypt data encrypted with this AES cipher
	chachagcm      cipher.AEAD // This is to decrypt data encrypted with this Chacha cipher
	cryptoOverhead int
	nonce          []byte
	nonceSize      int
	nonceUsed      int64
	nonceLimit     int64
}

// CreateGCMs is creating the cipher.AEADs and return the code for
// the selected cipher, or an error if the given cipher is not supported
// or an error occurs when creating the cipher.AEADs.
// The returned cipher.AEAD are in the following order:
func CreateGCMs(encryptionCipher string, encryptionKey []byte) (byte, map[byte]cipher.AEAD, error) {
	var code byte
	// If user provides cipher, use that.
	if encryptionCipher != CryptoCipherAutoSelect {
		switch strings.ToUpper(encryptionCipher) {
		case CryptoCipherAES:
			code = CryptoCodeAES
		case CryptoCipherChaChaPoly:
			code = CryptoCodeChaCha
		default:
			return 0, nil, ErrCipherNotSupported
		}
	} else {
		// Otherwise default to AES on intel (there is hardware
		// acceleration for that) and chacha20poly1305 on ARM
		// (the two arch'es that we build docker images for).
		if runtime.GOARCH == "amd64" || runtime.GOARCH == "386" {
			code = CryptoCodeAES
		} else {
			fmt.Printf("@@IK: selected CHACHA\n")
		}
	}
	// Always check env variable first
	key := []byte(os.Getenv(CryptoStoreEnvKeyName))
	if len(key) == 0 {
		key = encryptionKey
		if len(key) == 0 {
			return 0, nil, ErrCryptoStoreRequiresKey
		}
	}

	ciphers := make(map[byte]cipher.AEAD)

	h := sha256.New()
	h.Write(key)
	keyHash := h.Sum(nil)

	block, err := aes.NewCipher(keyHash)
	if err != nil {
		return 0, nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return 0, nil, err
	}
	ciphers[CryptoCodeAES] = gcm

	gcm, err = chacha20poly1305.New(keyHash)
	if err != nil {
		return 0, nil, err
	}
	ciphers[CryptoCodeChaCha] = gcm

	// On success, erase the key
	for i := 0; i < len(encryptionKey); i++ {
		encryptionKey[i] = 'x'
	}

	return code, ciphers, nil
}

// NewCryptoStore returns a CryptoStore instance with
// given underlying store.
func NewCryptoStore(s Store, encryptionCipher string, encryptionKey []byte) (*CryptoStore, error) {
	code, ciphers, err := CreateGCMs(encryptionCipher, encryptionKey)
	if err != nil {
		return nil, err
	}
	cs := &CryptoStore{
		Store:      s,
		cipherCode: code,
		aesgcm:     ciphers[CryptoCodeAES],
		chachagcm:  ciphers[CryptoCodeChaCha],
	}
	// These values are same for the 2 ciphers we support,
	// so use any of the gcm.
	cs.cryptoOverhead = cs.aesgcm.Overhead()
	cs.nonceSize = cs.aesgcm.NonceSize()

	return cs, nil
}

func (cs *CryptoStore) newCryptoMsgStore(ms MsgStore) *CryptoMsgStore {
	cms := &CryptoMsgStore{
		MsgStore:       ms,
		cipherCode:     cs.cipherCode,
		aesgcm:         cs.aesgcm,
		chachagcm:      cs.chachagcm,
		nonceSize:      cs.nonceSize,
		cryptoOverhead: cs.cryptoOverhead,
	}
	switch cs.cipherCode {
	case CryptoCodeAES:
		cms.gcm = cs.aesgcm
	case CryptoCodeChaCha:
		cms.gcm = cs.chachagcm
	}
	cms.generateNewNonce()
	return cms
}

// Recover implements the Store interface
func (cs *CryptoStore) Recover() (*RecoveredState, error) {
	cs.Lock()
	defer cs.Unlock()
	rs, err := cs.Store.Recover()
	if rs == nil || err != nil {
		return rs, err
	}
	for _, rc := range rs.Channels {
		rc.Channel.Msgs = cs.newCryptoMsgStore(rc.Channel.Msgs)
	}
	return rs, nil
}

// CreateChannel implements the Store interface
func (cs *CryptoStore) CreateChannel(channel string) (*Channel, error) {
	cs.Lock()
	defer cs.Unlock()

	c, err := cs.Store.CreateChannel(channel)
	if err != nil {
		return nil, err
	}
	c.Msgs = cs.newCryptoMsgStore(c.Msgs)
	return c, nil
}

// Store implements the MsgStore interface
func (cms *CryptoMsgStore) Store(msg *pb.MsgProto) (uint64, error) {
	if len(msg.Data) == 0 {
		return cms.MsgStore.Store(msg)
	}
	ed, err := cms.encrypt(msg.Data)
	if err != nil {
		return 0, err
	}
	msg.Data = ed
	return cms.MsgStore.Store(msg)
}

func (cms *CryptoMsgStore) generateNewNonce() error {
	cms.nonce = make([]byte, cms.nonceSize)
	if _, err := io.ReadFull(rand.Reader, cms.nonce); err != nil {
		return err
	}
	cms.nonceUsed = 0
	cms.nonceLimit = mrand.Int63n(1e6) + 100000
	return nil
}

func (cms *CryptoMsgStore) encrypt(data []byte) ([]byte, error) {
	// We can't reuse a buffer since when we pass the data to
	// the underlying store, we don't know if this is retained
	// in some cache, etc..
	buf := make([]byte, 1+cms.nonceSize+cms.cryptoOverhead+len(data))
	cms.Lock()
	buf[0] = cms.cipherCode
	copy(buf[1:], cms.nonce)
	copy(buf[1+cms.nonceSize:], data)
	dst := buf[1+cms.nonceSize : 1+cms.nonceSize+len(data)]
	ed := cms.gcm.Seal(dst[:0], cms.nonce, dst, nil)
	cms.nonceUsed++
	if cms.nonceUsed >= cms.nonceLimit {
		cms.generateNewNonce()
	}
	cms.Unlock()
	return buf[:1+cms.nonceSize+len(ed)], nil
}

func (cms *CryptoMsgStore) decryptedMsg(m *pb.MsgProto) (*pb.MsgProto, error) {
	var gcm cipher.AEAD
	if len(m.Data) > 0 {
		switch m.Data[0] {
		case CryptoCodeAES:
			gcm = cms.aesgcm
		case CryptoCodeChaCha:
			gcm = cms.chachagcm
		default:
			// Anything else, assume no algo or something we don't know how to decrypt.
			return m, nil
		}
	}
	if len(m.Data) <= 1+cms.nonceSize {
		return nil, fmt.Errorf("trying to decrypt data that is not (len=%v)", len(m.Data))
	}
	// When decrypting we can't do it in the original buffer because
	// the store's copy may be in a cache and so this would decipher
	// the encrypted copy and during the next call to decryptedMsg()
	// for the same message, there would be attempt to decrypt something
	// that is not, which would fail.
	dd, err := gcm.Open(nil, m.Data[1:1+cms.nonceSize], m.Data[1+cms.nonceSize:], nil)
	if err != nil {
		return nil, err
	}
	// Store owns the message, so make a copy before returning
	retMsg := *m
	retMsg.Data = dd
	return &retMsg, nil
}

// Lookup implements the MsgStore interface
func (cms *CryptoMsgStore) Lookup(seq uint64) (*pb.MsgProto, error) {
	m, err := cms.MsgStore.Lookup(seq)
	if m == nil || m.Data == nil || err != nil {
		return m, err
	}
	return cms.decryptedMsg(m)
}

// FirstMsg implements the MsgStore interface
func (cms *CryptoMsgStore) FirstMsg() (*pb.MsgProto, error) {
	m, err := cms.MsgStore.FirstMsg()
	if m == nil || m.Data == nil || err != nil {
		return m, err
	}
	return cms.decryptedMsg(m)
}

// LastMsg implements the MsgStore interface
func (cms *CryptoMsgStore) LastMsg() (*pb.MsgProto, error) {
	m, err := cms.MsgStore.LastMsg()
	if m == nil || m.Data == nil || err != nil {
		return m, err
	}
	return cms.decryptedMsg(m)
}
<commit_msg>Removed debug print statement and put back actual code!<commit_after>// Copyright 2018 The NATS Authors
// 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 stores

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	mrand "math/rand"
	"os"
	"runtime"
	"strings"
	"sync"

	"github.com/nats-io/go-nats-streaming/pb"
	"golang.org/x/crypto/chacha20poly1305"
)

// CryptStore specific errors
var (
	ErrCryptoStoreRequiresKey = errors.New("encryption key required")
	ErrCipherNotSupported     = errors.New("encryption cipher not supported")
)

const (
	// CryptoStoreEnvKeyName is the environment variable name
	// that the CryptoStore looks up if no key is passed as
	// a parameter.
	CryptoStoreEnvKeyName = "NATS_STREAMING_ENCRYPTION_KEY"

	// CryptoCipherAES is the name of the AES cipher to use for encryption
	CryptoCipherAES = "AES"

	// CryptoCipherChaChaPoly is the name of the ChaChaPoly cipher to use for encryption
	CryptoCipherChaChaPoly = "CHACHA"

	// CryptoCipherAutoSelect if passed to NewCryptoStore() will cause the cipher to
	// be auto-selected based on the platform the executable is built for.
	CryptoCipherAutoSelect = ""
)

// These constants define a code for each of the supported ciphers
const (
	CryptoCodeAES    = byte(1)
	CryptoCodeChaCha = byte(2)
)

// CryptoStore is a store wrapping a store implementation
// and adds encryption support.
type CryptoStore struct {
	sync.Mutex
	Store

	// These are set when the store is created. They are then
	// passed to a CryptoMsgStore so that there is no need
	// to reference back to these.
	// Note that nonceSize and cryptoOverhead are same for
	// those 2 ciphers. If we add more and those are different,
	// will need to be stored differently or call the appropriate
	// gcm.NonceSize() and gcm.Overhead() functions.
	cipherCode     byte
	aesgcm         cipher.AEAD
	chachagcm      cipher.AEAD
	nonceSize      int
	cryptoOverhead int
}

// CryptoMsgStore is a store wrappeing a SubStore implementation
// and adds encryption support.
type CryptoMsgStore struct {
	sync.Mutex
	MsgStore
	cipherCode     byte
	gcm            cipher.AEAD // Use this one to encrypt
	aesgcm         cipher.AEAD // This is to decrypt data encrypted with this AES cipher
	chachagcm      cipher.AEAD // This is to decrypt data encrypted with this Chacha cipher
	cryptoOverhead int
	nonce          []byte
	nonceSize      int
	nonceUsed      int64
	nonceLimit     int64
}

// CreateGCMs is creating the cipher.AEADs and return the code for
// the selected cipher, or an error if the given cipher is not supported
// or an error occurs when creating the cipher.AEADs.
// The returned cipher.AEAD are in the following order:
func CreateGCMs(encryptionCipher string, encryptionKey []byte) (byte, map[byte]cipher.AEAD, error) {
	var code byte
	// If user provides cipher, use that.
	if encryptionCipher != CryptoCipherAutoSelect {
		switch strings.ToUpper(encryptionCipher) {
		case CryptoCipherAES:
			code = CryptoCodeAES
		case CryptoCipherChaChaPoly:
			code = CryptoCodeChaCha
		default:
			return 0, nil, ErrCipherNotSupported
		}
	} else {
		// Otherwise default to AES on intel (there is hardware
		// acceleration for that) and chacha20poly1305 on ARM
		// (the two arch'es that we build docker images for).
		if runtime.GOARCH == "amd64" || runtime.GOARCH == "386" {
			code = CryptoCodeAES
		} else {
			code = CryptoCodeChaCha
		}
	}
	// Always check env variable first
	key := []byte(os.Getenv(CryptoStoreEnvKeyName))
	if len(key) == 0 {
		key = encryptionKey
		if len(key) == 0 {
			return 0, nil, ErrCryptoStoreRequiresKey
		}
	}

	ciphers := make(map[byte]cipher.AEAD)

	h := sha256.New()
	h.Write(key)
	keyHash := h.Sum(nil)

	block, err := aes.NewCipher(keyHash)
	if err != nil {
		return 0, nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return 0, nil, err
	}
	ciphers[CryptoCodeAES] = gcm

	gcm, err = chacha20poly1305.New(keyHash)
	if err != nil {
		return 0, nil, err
	}
	ciphers[CryptoCodeChaCha] = gcm

	// On success, erase the key
	for i := 0; i < len(encryptionKey); i++ {
		encryptionKey[i] = 'x'
	}

	return code, ciphers, nil
}

// NewCryptoStore returns a CryptoStore instance with
// given underlying store.
func NewCryptoStore(s Store, encryptionCipher string, encryptionKey []byte) (*CryptoStore, error) {
	code, ciphers, err := CreateGCMs(encryptionCipher, encryptionKey)
	if err != nil {
		return nil, err
	}
	cs := &CryptoStore{
		Store:      s,
		cipherCode: code,
		aesgcm:     ciphers[CryptoCodeAES],
		chachagcm:  ciphers[CryptoCodeChaCha],
	}
	// These values are same for the 2 ciphers we support,
	// so use any of the gcm.
	cs.cryptoOverhead = cs.aesgcm.Overhead()
	cs.nonceSize = cs.aesgcm.NonceSize()

	return cs, nil
}

func (cs *CryptoStore) newCryptoMsgStore(ms MsgStore) *CryptoMsgStore {
	cms := &CryptoMsgStore{
		MsgStore:       ms,
		cipherCode:     cs.cipherCode,
		aesgcm:         cs.aesgcm,
		chachagcm:      cs.chachagcm,
		nonceSize:      cs.nonceSize,
		cryptoOverhead: cs.cryptoOverhead,
	}
	switch cs.cipherCode {
	case CryptoCodeAES:
		cms.gcm = cs.aesgcm
	case CryptoCodeChaCha:
		cms.gcm = cs.chachagcm
	}
	cms.generateNewNonce()
	return cms
}

// Recover implements the Store interface
func (cs *CryptoStore) Recover() (*RecoveredState, error) {
	cs.Lock()
	defer cs.Unlock()
	rs, err := cs.Store.Recover()
	if rs == nil || err != nil {
		return rs, err
	}
	for _, rc := range rs.Channels {
		rc.Channel.Msgs = cs.newCryptoMsgStore(rc.Channel.Msgs)
	}
	return rs, nil
}

// CreateChannel implements the Store interface
func (cs *CryptoStore) CreateChannel(channel string) (*Channel, error) {
	cs.Lock()
	defer cs.Unlock()

	c, err := cs.Store.CreateChannel(channel)
	if err != nil {
		return nil, err
	}
	c.Msgs = cs.newCryptoMsgStore(c.Msgs)
	return c, nil
}

// Store implements the MsgStore interface
func (cms *CryptoMsgStore) Store(msg *pb.MsgProto) (uint64, error) {
	if len(msg.Data) == 0 {
		return cms.MsgStore.Store(msg)
	}
	ed, err := cms.encrypt(msg.Data)
	if err != nil {
		return 0, err
	}
	msg.Data = ed
	return cms.MsgStore.Store(msg)
}

func (cms *CryptoMsgStore) generateNewNonce() error {
	cms.nonce = make([]byte, cms.nonceSize)
	if _, err := io.ReadFull(rand.Reader, cms.nonce); err != nil {
		return err
	}
	cms.nonceUsed = 0
	cms.nonceLimit = mrand.Int63n(1e6) + 100000
	return nil
}

func (cms *CryptoMsgStore) encrypt(data []byte) ([]byte, error) {
	// We can't reuse a buffer since when we pass the data to
	// the underlying store, we don't know if this is retained
	// in some cache, etc..
	buf := make([]byte, 1+cms.nonceSize+cms.cryptoOverhead+len(data))
	cms.Lock()
	buf[0] = cms.cipherCode
	copy(buf[1:], cms.nonce)
	copy(buf[1+cms.nonceSize:], data)
	dst := buf[1+cms.nonceSize : 1+cms.nonceSize+len(data)]
	ed := cms.gcm.Seal(dst[:0], cms.nonce, dst, nil)
	cms.nonceUsed++
	if cms.nonceUsed >= cms.nonceLimit {
		cms.generateNewNonce()
	}
	cms.Unlock()
	return buf[:1+cms.nonceSize+len(ed)], nil
}

func (cms *CryptoMsgStore) decryptedMsg(m *pb.MsgProto) (*pb.MsgProto, error) {
	var gcm cipher.AEAD
	if len(m.Data) > 0 {
		switch m.Data[0] {
		case CryptoCodeAES:
			gcm = cms.aesgcm
		case CryptoCodeChaCha:
			gcm = cms.chachagcm
		default:
			// Anything else, assume no algo or something we don't know how to decrypt.
			return m, nil
		}
	}
	if len(m.Data) <= 1+cms.nonceSize {
		return nil, fmt.Errorf("trying to decrypt data that is not (len=%v)", len(m.Data))
	}
	// When decrypting we can't do it in the original buffer because
	// the store's copy may be in a cache and so this would decipher
	// the encrypted copy and during the next call to decryptedMsg()
	// for the same message, there would be attempt to decrypt something
	// that is not, which would fail.
	dd, err := gcm.Open(nil, m.Data[1:1+cms.nonceSize], m.Data[1+cms.nonceSize:], nil)
	if err != nil {
		return nil, err
	}
	// Store owns the message, so make a copy before returning
	retMsg := *m
	retMsg.Data = dd
	return &retMsg, nil
}

// Lookup implements the MsgStore interface
func (cms *CryptoMsgStore) Lookup(seq uint64) (*pb.MsgProto, error) {
	m, err := cms.MsgStore.Lookup(seq)
	if m == nil || m.Data == nil || err != nil {
		return m, err
	}
	return cms.decryptedMsg(m)
}

// FirstMsg implements the MsgStore interface
func (cms *CryptoMsgStore) FirstMsg() (*pb.MsgProto, error) {
	m, err := cms.MsgStore.FirstMsg()
	if m == nil || m.Data == nil || err != nil {
		return m, err
	}
	return cms.decryptedMsg(m)
}

// LastMsg implements the MsgStore interface
func (cms *CryptoMsgStore) LastMsg() (*pb.MsgProto, error) {
	m, err := cms.MsgStore.LastMsg()
	if m == nil || m.Data == nil || err != nil {
		return m, err
	}
	return cms.decryptedMsg(m)
}
<|endoftext|>


<commit_before>package assert

import "testing"

func TestThatStringIsEqualToPrintsNothing(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Alice").IsEqualTo("Alice").IsEqualTo("Alice")
	mockT.HasNoErrors()
}

func TestThatStringIsEqualToPrintsMessage(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Alice").IsEqualTo("Bob").IsEqualTo("Charlie")
	mockT.HasErrorMessages(
		"Expected <Bob>, but was <Alice>.",
		"Expected <Charlie>, but was <Alice>.",
	)
}

func TestThatStringIsNotEqualToPrintsNothing(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Alice").IsNotEqualTo("Bob").IsNotEqualTo("Bob")
	mockT.HasNoErrors()
}

func TestThatStringIsNotEqualToPrintsMessage(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Alice").IsNotEqualTo("Alice")
	mockT.HasErrorMessages(
		"Expected string not equal to <Alice>, but was equal.",
	)
}

func TestThatStringIsEmptyPrintsNothing(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "").IsEmpty()
	mockT.HasNoErrors()
}

func TestThatStringIsEmptyPrintsMessage(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Marta").IsEmpty()
	mockT.HasErrorMessages(
		"Expected string to be empty, but was <Marta>.",
	)
}

func TestThatStringIsInSlicePrintsNothing(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Marta").IsInSlice([]string{"Marta"})
	mockT.HasNoErrors()
}

func TestThatStringIsInSlicePrintsMessage(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Marta").IsInSlice([]string{"Michal", "Macy"})
	mockT.HasErrorMessages(
		"Expected string to be in slice <[Michal Macy]>, but wasn't.",
	)
}

func TestThatStringContainsePrintsNothing(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Marta").Contains("art")
	mockT.HasNoErrors()
}

func TestThatStringContainsPrintsMessage(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Marta").Contains("Michal")
	mockT.HasErrorMessages(
		"Expected string to contain <Michal>, but didn't.",
	)
}
<commit_msg>fixed typo in method name<commit_after>package assert

import "testing"

func TestThatStringIsEqualToPrintsNothing(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Alice").IsEqualTo("Alice").IsEqualTo("Alice")
	mockT.HasNoErrors()
}

func TestThatStringIsEqualToPrintsMessage(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Alice").IsEqualTo("Bob").IsEqualTo("Charlie")
	mockT.HasErrorMessages(
		"Expected <Bob>, but was <Alice>.",
		"Expected <Charlie>, but was <Alice>.",
	)
}

func TestThatStringIsNotEqualToPrintsNothing(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Alice").IsNotEqualTo("Bob").IsNotEqualTo("Bob")
	mockT.HasNoErrors()
}

func TestThatStringIsNotEqualToPrintsMessage(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Alice").IsNotEqualTo("Alice")
	mockT.HasErrorMessages(
		"Expected string not equal to <Alice>, but was equal.",
	)
}

func TestThatStringIsEmptyPrintsNothing(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "").IsEmpty()
	mockT.HasNoErrors()
}

func TestThatStringIsEmptyPrintsMessage(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Marta").IsEmpty()
	mockT.HasErrorMessages(
		"Expected string to be empty, but was <Marta>.",
	)
}

func TestThatStringIsInSlicePrintsNothing(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Marta").IsInSlice([]string{"Marta"})
	mockT.HasNoErrors()
}

func TestThatStringIsInSlicePrintsMessage(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Marta").IsInSlice([]string{"Michal", "Macy"})
	mockT.HasErrorMessages(
		"Expected string to be in slice <[Michal Macy]>, but wasn't.",
	)
}

func TestThatStringContainsPrintsNothing(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Marta").Contains("art")
	mockT.HasNoErrors()
}

func TestThatStringContainsPrintsMessage(t *testing.T) {
	mockT := newMockT(t)
	ThatString(mockT, "Marta").Contains("Michal")
	mockT.HasErrorMessages(
		"Expected string to contain <Michal>, but didn't.",
	)
}
<|endoftext|>


<commit_before>// Copyright 2016 the DestructHub Authors. All rights reserved
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

// solution for the problem 27 of Project Euler
// concurrent solution: take about 3.5~4 seconds (about 3x faster from python version)

package main

import (
	"fmt"
	"math"
	"time"
)

const (
	limit    = 1000 // a, b limits to search
	nThreads = 16   // number of nThreads
)

type result struct {
	primes int
	a      int
	b      int
}

var (
	resultchan = make(chan result, nThreads) // used to get results concurrently
	memory     = make(map[int]bool)          // memory for primes numbers
)

// quadratic function generator
// n² + an + b
func quadratic(a, b int) func(int) int {
	return func(n int) int {
		return n*n + a*n + b
	}
}

// decorator workaround on Go
// check if the value was be computed before
// return if exists, else compute
func memo(f func(int) bool) func(int) bool {
	return func(n int) bool {
		if val, exists := memory[n]; exists {
			return val
		} else {
			return f(n)
		}
	}
}

// simple function
// to check if is a prime
func isPrime(n int) bool {
	if n < 2 {
		return false
	}
	for q := 2; float64(q) < math.Floor(math.Sqrt(float64(n)))+1; q++ {
		if n%q == 0 {
			return false
		}
	}

	return true
}

// func evaluation n² + an + b
// return the n primes generate without gaps
// between [0, n]
func evalFunc(f func(int) int) int {
	n := 0
	for {
		if memo(isPrime)(f(n)) {
			n += 1
		} else {
			break
		}
	}
	return n
}

// make concurrent evaluation
// of the equation whose have
// generate more primes
func worker(split int) {
	limitSplited := split * limit / nThreads
	a, b, nPrimes := 0, 0, 0
	for i := -limitSplited; i < limitSplited; i++ {
		for j := -limit; j < limit; j++ {
			if n := evalFunc(quadratic(i, j)); n > nPrimes {
				a, b = i, j
				nPrimes = n
			}
		}
	}

	resultchan <- result{nPrimes, a, b}
}

// receiver the results
// from workers
func receiver() (a int, b int, nPrimes int) {
	a, b, nPrimes = 0, 0, 0
	for i := 0; i < nThreads; i++ {
		result := <-resultchan
		if result.primes > nPrimes {
			a, b = result.a, result.b
			nPrimes = result.primes
		}
	}
	return a, b, nPrimes
}

func main() {
	start := time.Now()
	for part := 1; part <= nThreads; part++ {
		go worker(part)
	}

	a, b, nPrimes := receiver()
	elapsed := time.Since(start)
	fmt.Printf("Equation: n² + %vn + %v\n", a, b)
	fmt.Printf("Generate: primes on range [0, %v]\n", nPrimes)
	fmt.Printf("Answer: %v * %v => %v\n", a, b, a*b)
	fmt.Printf("Time execution: %v\n", elapsed)
}
<commit_msg>Uniq solution made in Golang checked about #2<commit_after>// Copyright 2016 the DestructHub Authors. All rights reserved
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

// solution for the problem 27 of Project Euler
// concurrent solution: take about 3.5~4 seconds (about 3x faster from python version)

package main

import (
	"fmt"
	"math"
)

const (
	limit    = 1000 // a, b limits to search
	nThreads = 16   // number of nThreads
)

type result struct {
	primes int
	a      int
	b      int
}

var (
	resultchan = make(chan result, nThreads) // used to get results concurrently
	memory     = make(map[int]bool)          // memory for primes numbers
)

// quadratic function generator
// n² + an + b
func quadratic(a, b int) func(int) int {
	return func(n int) int {
		return n*n + a*n + b
	}
}

// decorator workaround on Go
// check if the value was be computed before
// return if exists, else compute
func memo(f func(int) bool) func(int) bool {
	return func(n int) bool {
		if val, exists := memory[n]; exists {
			return val
		} else {
			return f(n)
		}
	}
}

// simple function
// to check if is a prime
func isPrime(n int) bool {
	if n < 2 {
		return false
	}
	for q := 2; float64(q) < math.Floor(math.Sqrt(float64(n)))+1; q++ {
		if n%q == 0 {
			return false
		}
	}

	return true
}

// func evaluation n² + an + b
// return the n primes generate without gaps
// between [0, n]
func evalFunc(f func(int) int) int {
	n := 0
	for {
		if memo(isPrime)(f(n)) {
			n += 1
		} else {
			break
		}
	}
	return n
}

// make concurrent evaluation
// of the equation whose have
// generate more primes
func worker(split int) {
	limitSplited := split * limit / nThreads
	a, b, nPrimes := 0, 0, 0
	for i := -limitSplited; i < limitSplited; i++ {
		for j := -limit; j < limit; j++ {
			if n := evalFunc(quadratic(i, j)); n > nPrimes {
				a, b = i, j
				nPrimes = n
			}
		}
	}

	resultchan <- result{nPrimes, a, b}
}

// receiver the results
// from workers
func receiver() (a int, b int, nPrimes int) {
	a, b, nPrimes = 0, 0, 0
	for i := 0; i < nThreads; i++ {
		result := <-resultchan
		if result.primes > nPrimes {
			a, b = result.a, result.b
			nPrimes = result.primes
		}
	}
	return a, b, nPrimes
}

func main() {
	for part := 1; part <= nThreads; part++ {
		go worker(part)
	}

	a, b, _ := receiver()
	fmt.Printf("%v\n", a*b)
}
<|endoftext|>


<commit_before>//
// daemon.go
//
package srnd
import (
  "log"
  "net"
  "strconv"
  "strings"
  "net/textproto"
  "os"
  "time"
)

type NNTPDaemon struct {
  instance_name string
  bind_addr string
  conf *SRNdConfig
  store ArticleStore
  database Database
  mod ModEngine
  expire ExpirationCore
  listener net.Listener
  debug bool
  sync_on_start bool
  running bool
  // http frontend
  frontend Frontend

  // thumbnail generator for images
  img_thm ThumbnailGenerator

  // nntp feeds map, feed, isoutbound
  feeds map[NNTPConnection]bool
  infeed chan NNTPMessage
  // channel to load messages to infeed given their message id
  infeed_load chan string
  // channel for broadcasting a message to all feeds given their newsgroup, message_id
  send_all_feeds chan ArticleEntry
  // channel for broadcasting an ARTICLE command to all feeds in reader mode
  ask_for_article chan ArticleEntry
}

func (self *NNTPDaemon) End() {
  self.listener.Close()
}

// register a new connection
// can be either inbound or outbound
func (self *NNTPDaemon) newConnection(conn net.Conn, inbound bool, policy *FeedPolicy) NNTPConnection {
  allow_tor := self.conf.daemon["allow_tor"]
  allow_tor_attachments := self.conf.daemon["allow_tor_attachments"]
  feed := NNTPConnection{conn, textproto.NewConn(conn), inbound, self.debug, new(ConnectionInfo), policy,  make(chan ArticleEntry, 128), self.store, self.store, allow_tor == "1", allow_tor_attachments == "1"}
  self.feeds[feed] = ! inbound
  return feed
}

func (self *NNTPDaemon) persistFeed(conf FeedConfig, mode string) {
  for {
    if self.running {

      var conn net.Conn
      var err error
      proxy_type := strings.ToLower(conf.proxy_type)

      if proxy_type ==  "" || proxy_type == "none" {
        // connect out without proxy
        log.Println("dial out to ", conf.addr)
        conn, err = net.Dial("tcp", conf.addr)
        if err != nil {
          log.Println("cannot connect to outfeed", conf.addr, err)
					time.Sleep(time.Second * 5)
          continue
        }
      } else if proxy_type == "socks4a" {
        // connect via socks4a
        log.Println("dial out via proxy", conf.proxy_addr)
        conn, err = net.Dial("tcp", conf.proxy_addr)
        if err != nil {
          log.Println("cannot connect to proxy", conf.proxy_addr)
					time.Sleep(time.Second * 5)
          continue
        }
        // generate request
        idx := strings.LastIndex(conf.addr, ":")
        if idx == -1 {
          log.Fatal("invalid outfeed address")
        }
        var port uint64
        addr := conf.addr[:idx]
        port, err = strconv.ParseUint(conf.addr[idx+1:], 10, 16)
        if port >= 25536 {
          log.Fatal("bad proxy port" , port)
        }
        var proxy_port uint16
        proxy_port = uint16(port)
        proxy_ident := "srndv2"
        req_len := len(addr) + 1 + len(proxy_ident) + 1 + 8

        req := make([]byte, req_len)
        // pack request
        req[0] = '\x04'
        req[1] = '\x01'
        req[2] = byte(proxy_port & 0xff00 >> 8)
        req[3] = byte(proxy_port & 0x00ff)
        req[7] = '\x01'
        idx = 8

        proxy_ident_b := []byte(proxy_ident)
        addr_b := []byte(addr)

        var bi int
        for bi = range proxy_ident_b {
          req[idx] = proxy_ident_b[bi]
          idx += 1
        }
        idx += 1
        for bi = range addr_b {
          req[idx] = addr_b[bi]
          idx += 1
        }

        // send request
        conn.Write(req)
        resp := make([]byte, 8)

        // receive response
        conn.Read(resp)
        if resp[1] == '\x5a' {
          // success
          log.Println("connected to", conf.addr)
        } else {
          log.Println("failed to connect to", conf.addr)
					time.Sleep(time.Second * 5)
          continue
        }
      }
      policy := &conf.policy
      nntp := self.newConnection(conn, false, policy)
      // start syncing in background if streaming mode
      if mode == "stream" {
        go func() {
          if self.sync_on_start {
            log.Println("sync on start")
            // get every article
            articles := self.database.GetAllArticles()
            // wait 5 seconds for feed to handshake
            time.Sleep(5 * time.Second)
            log.Println("outfeed begin sync")
            for _, result := range articles {
              if policy.AllowsNewsgroup(result.Newsgroup()) {
                //XXX: will this crash if interrupted?
                nntp.sync <- result
              }
            }
            log.Println("outfeed end sync")
          }
        }()
      }
      nntp.HandleOutbound(self, conf.quarks, mode)
      log.Println("remove outfeed")
      delete(self.feeds, nntp)
      close(nntp.sync)
    }
  }
  time.Sleep(1 * time.Second)
}

// run daemon
func (self *NNTPDaemon) Run() {
  defer self.listener.Close()
  // run expiration mainloop
  go self.expire.Mainloop()
  // we are now running
  self.running = true

  // persist outfeeds
  for idx := range self.conf.feeds {
    go self.persistFeed(self.conf.feeds[idx], "stream")
    go self.persistFeed(self.conf.feeds[idx], "reader")
  }

  // start accepting incoming connections
  go self.acceptloop()

  go func () {
    // if we have no initial posts create one
    if self.database.ArticleCount() == 0 {
      nntp := newPlaintextArticle("welcome to nntpchan, this post was inserted on startup automatically", "system@"+self.instance_name, "Welcome to NNTPChan", "system", self.instance_name, "overchan.test")
      nntp.Pack()
      file := self.store.CreateTempFile(nntp.MessageID())
      if file != nil {
        err := self.store.WriteMessage(nntp, file)
        file.Close()
        if err == nil {
          self.infeed <- nntp
        } else {
          log.Println("failed to create startup messge?", err)
        }
      }
    }
  }()

  // get all pending articles from infeed and load them
  go func() {
    f, err := os.Open(self.store.TempDir())
    if err == nil {
      names, err := f.Readdirnames(0)
      if err == nil {
        for _, name := range names {
          self.infeed_load <- name
        }
      }
    }

  }()

  // if we have no frontend this does nothing
  if self.frontend != nil {
    go self.pollfrontend()
  }
  go self.pollinfeed()
  go self.polloutfeeds()
  self.pollmessages()
}

func (self *NNTPDaemon) pollfrontend() {
  chnl := self.frontend.NewPostsChan()
  for {
    nntp := <- chnl
    // new post from frontend
    log.Println("frontend post", nntp.MessageID())
    self.infeed <- nntp
  }
}
func (self *NNTPDaemon) pollinfeed() {
  for {
    msgid := <- self.infeed_load
    log.Println("load from infeed", msgid)
    msg := self.store.ReadTempMessage(msgid)
    if msg != nil {
      self.infeed <- msg
    }
  }
}

func (self *NNTPDaemon) polloutfeeds() {
  for {
    select {
    case nntp := <- self.send_all_feeds:
      for feed , use := range self.feeds {
        if use && feed.policy != nil {
          if feed.policy.AllowsNewsgroup(nntp.Newsgroup()) {
            if feed.info.mode == "stream" {
              feed.sync <- nntp
            }
          }
        }
      }
    case nntp := <- self.ask_for_article:
      for feed, use := range self.feeds {
        if use && feed.policy != nil {
          if feed.policy.AllowsNewsgroup(nntp.Newsgroup()) {
            if feed.info.mode == "reader" {
              log.Println("ask for", nntp.MesaageID())
              feed.sync <- nntp
            }
          }
        }
      }
    }
  }
}

func (self *NNTPDaemon) pollmessages() {
  var chnl chan NNTPMessage
  modchnl := self.mod.MessageChan()
  if self.frontend != nil {
    chnl = self.frontend.PostsChan()
  }
  for {

    nntp := <- self.infeed
    // ammend path
    nntp.AppendPath(self.instance_name)
    msgid := nntp.MessageID()
    log.Println("daemon got", msgid)

    // store article and attachments
    // register with database
    // this also generates thumbnails
    go self.store.StorePost(nntp)

    // prepare for content rollover
    // fallback rollover
    rollover := 100

    group := nntp.Newsgroup()
    tpp, err := self.database.GetThreadsPerPage(group)
    ppb, err := self.database.GetPagesPerBoard(group)
    if err == nil {
      rollover = tpp * ppb
    }

    // roll over old content
    self.expire.ExpireGroup(group, rollover)
    // handle mod events
    if group == "ctl" {
      modchnl <- nntp
    }

    // queue to all outfeeds
    // XXX: blocking ?
    self.send_all_feeds <- ArticleEntry{msgid, group}
    // tell frontend
    // XXX: blocking ?
    if chnl != nil {
      if self.frontend.AllowNewsgroup(group) {
        chnl <- nntp
      }
    }
  }
}

func (self *NNTPDaemon) acceptloop() {
  for {
    // accept
    conn, err := self.listener.Accept()
    if err != nil {
      log.Fatal(err)
    }
    // make a new inbound nntp connection handler
    nntp := self.newConnection(conn, true, nil)
    go self.RunInbound(nntp)
  }
}

func (self *NNTPDaemon) RunInbound(nntp NNTPConnection) {
  nntp.HandleInbound(self)
  delete(self.feeds, nntp)
}

func (self *NNTPDaemon) Setup() {
  log.Println("checking for configs...")
  // check that are configs exist
  CheckConfig()
  log.Println("loading config...")
  // read the config
  self.conf = ReadConfig()
  if self.conf == nil {
    log.Fatal("failed to load config")
  }
  // validate the config
  log.Println("validating configs...")
  self.conf.Validate()
  log.Println("configs are valid")

  db_host := self.conf.database["host"]
  db_port := self.conf.database["port"]
  db_user := self.conf.database["user"]
  db_passwd := self.conf.database["password"]

  // set up database stuff
  log.Println("connecting to database...")
  self.database = NewDatabase(self.conf.database["type"], self.conf.database["schema"], db_host, db_port, db_user, db_passwd)
  log.Println("ensure that the database is created...")
  self.database.CreateTables()

  r_url := self.conf.worker["url"]

  // set up store
  log.Println("set up article store...")
  self.store = createArticleStore(self.conf.store, r_url, self.database)

  self.mod = modEngine{
    store: self.store,
    database:  self.database,
    chnl: make(chan NNTPMessage),
  }
}

// bind to address
func (self *NNTPDaemon) Bind() error {
  listener , err := net.Listen("tcp", self.bind_addr)
  if err != nil {
    log.Println("failed to bind to", self.bind_addr, err)
    return err
  }
  self.listener = listener
  log.Printf("SRNd NNTPD bound at %s", listener.Addr())
  return nil
}

// load configuration
// bind to interface
func (self *NNTPDaemon) Init() bool {

  // set up daemon configs
  self.Setup()

  self.infeed = make(chan NNTPMessage, 8)
  self.infeed_load = make(chan string)
  self.send_all_feeds = make(chan ArticleEntry, 64)
  self.feeds = make(map[NNTPConnection]bool)
  self.ask_for_article = make(chan ArticleEntry, 64)

  self.bind_addr = self.conf.daemon["bind"]

  err := self.Bind()
  if err != nil {
    log.Println("failed to bind:", err)
    return false
  }

  self.expire = createExpirationCore(self.database, self.store)
  self.sync_on_start = self.conf.daemon["sync_on_start"] == "1"
  self.debug = self.conf.daemon["log"] == "debug"
  self.instance_name = self.conf.daemon["instance_name"]
  if self.debug {
    log.Println("debug mode activated")
  }

  // do we enable the frontend?
  if self.conf.frontend["enable"] == "1" {
    log.Printf("frontend %s enabled", self.conf.frontend["name"])
    http_frontend := NewHTTPFrontend(self, self.conf.frontend, self.conf.worker["url"])
    nntp_frontend := NewNNTPFrontend(self, self.conf.frontend["nntp"])
    self.frontend = MuxFrontends(http_frontend, nntp_frontend)
    go self.frontend.Mainloop()
  }

  // set up admin user if it's specified in the config
  pubkey , ok := self.conf.frontend["admin_key"]
  if ok {
    // TODO: check for valid format
    log.Println("add admin key", pubkey)
    err = self.database.MarkModPubkeyGlobal(pubkey)
    if err != nil {
      log.Printf("failed to add admin mod key, %s", err)
    }
  }
  go RunModEngine(self.mod)
  return true
}
<commit_msg>mehhhh<commit_after>//
// daemon.go
//
package srnd
import (
  "log"
  "net"
  "strconv"
  "strings"
  "net/textproto"
  "os"
  "time"
)

type NNTPDaemon struct {
  instance_name string
  bind_addr string
  conf *SRNdConfig
  store ArticleStore
  database Database
  mod ModEngine
  expire ExpirationCore
  listener net.Listener
  debug bool
  sync_on_start bool
  running bool
  // http frontend
  frontend Frontend

  // thumbnail generator for images
  img_thm ThumbnailGenerator

  // nntp feeds map, feed, isoutbound
  feeds map[NNTPConnection]bool
  infeed chan NNTPMessage
  // channel to load messages to infeed given their message id
  infeed_load chan string
  // channel for broadcasting a message to all feeds given their newsgroup, message_id
  send_all_feeds chan ArticleEntry
  // channel for broadcasting an ARTICLE command to all feeds in reader mode
  ask_for_article chan ArticleEntry
}

func (self *NNTPDaemon) End() {
  self.listener.Close()
}

// register a new connection
// can be either inbound or outbound
func (self *NNTPDaemon) newConnection(conn net.Conn, inbound bool, policy *FeedPolicy) NNTPConnection {
  allow_tor := self.conf.daemon["allow_tor"]
  allow_tor_attachments := self.conf.daemon["allow_tor_attachments"]
  feed := NNTPConnection{conn, textproto.NewConn(conn), inbound, self.debug, new(ConnectionInfo), policy,  make(chan ArticleEntry, 128), self.store, self.store, allow_tor == "1", allow_tor_attachments == "1"}
  self.feeds[feed] = ! inbound
  return feed
}

func (self *NNTPDaemon) persistFeed(conf FeedConfig, mode string) {
  for {
    if self.running {

      var conn net.Conn
      var err error
      proxy_type := strings.ToLower(conf.proxy_type)

      if proxy_type ==  "" || proxy_type == "none" {
        // connect out without proxy
        log.Println("dial out to ", conf.addr)
        conn, err = net.Dial("tcp", conf.addr)
        if err != nil {
          log.Println("cannot connect to outfeed", conf.addr, err)
					time.Sleep(time.Second * 5)
          continue
        }
      } else if proxy_type == "socks4a" {
        // connect via socks4a
        log.Println("dial out via proxy", conf.proxy_addr)
        conn, err = net.Dial("tcp", conf.proxy_addr)
        if err != nil {
          log.Println("cannot connect to proxy", conf.proxy_addr)
					time.Sleep(time.Second * 5)
          continue
        }
        // generate request
        idx := strings.LastIndex(conf.addr, ":")
        if idx == -1 {
          log.Fatal("invalid outfeed address")
        }
        var port uint64
        addr := conf.addr[:idx]
        port, err = strconv.ParseUint(conf.addr[idx+1:], 10, 16)
        if port >= 25536 {
          log.Fatal("bad proxy port" , port)
        }
        var proxy_port uint16
        proxy_port = uint16(port)
        proxy_ident := "srndv2"
        req_len := len(addr) + 1 + len(proxy_ident) + 1 + 8

        req := make([]byte, req_len)
        // pack request
        req[0] = '\x04'
        req[1] = '\x01'
        req[2] = byte(proxy_port & 0xff00 >> 8)
        req[3] = byte(proxy_port & 0x00ff)
        req[7] = '\x01'
        idx = 8

        proxy_ident_b := []byte(proxy_ident)
        addr_b := []byte(addr)

        var bi int
        for bi = range proxy_ident_b {
          req[idx] = proxy_ident_b[bi]
          idx += 1
        }
        idx += 1
        for bi = range addr_b {
          req[idx] = addr_b[bi]
          idx += 1
        }

        // send request
        conn.Write(req)
        resp := make([]byte, 8)

        // receive response
        conn.Read(resp)
        if resp[1] == '\x5a' {
          // success
          log.Println("connected to", conf.addr)
        } else {
          log.Println("failed to connect to", conf.addr)
					time.Sleep(time.Second * 5)
          continue
        }
      }
      policy := &conf.policy
      nntp := self.newConnection(conn, false, policy)
      // start syncing in background if streaming mode
      if mode == "stream" {
        go func() {
          if self.sync_on_start {
            log.Println("sync on start")
            // get every article
            articles := self.database.GetAllArticles()
            // wait 5 seconds for feed to handshake
            time.Sleep(5 * time.Second)
            log.Println("outfeed begin sync")
            for _, result := range articles {
              if policy.AllowsNewsgroup(result.Newsgroup()) {
                //XXX: will this crash if interrupted?
                nntp.sync <- result
              }
            }
            log.Println("outfeed end sync")
          }
        }()
      }
      nntp.HandleOutbound(self, conf.quarks, mode)
      log.Println("remove outfeed")
      delete(self.feeds, nntp)
      close(nntp.sync)
    }
  }
  time.Sleep(1 * time.Second)
}

// run daemon
func (self *NNTPDaemon) Run() {
  defer self.listener.Close()
  // run expiration mainloop
  go self.expire.Mainloop()
  // we are now running
  self.running = true

  // persist outfeeds
  for idx := range self.conf.feeds {
    go self.persistFeed(self.conf.feeds[idx], "stream")
    go self.persistFeed(self.conf.feeds[idx], "reader")
  }

  // start accepting incoming connections
  go self.acceptloop()

  go func () {
    // if we have no initial posts create one
    if self.database.ArticleCount() == 0 {
      nntp := newPlaintextArticle("welcome to nntpchan, this post was inserted on startup automatically", "system@"+self.instance_name, "Welcome to NNTPChan", "system", self.instance_name, "overchan.test")
      nntp.Pack()
      file := self.store.CreateTempFile(nntp.MessageID())
      if file != nil {
        err := self.store.WriteMessage(nntp, file)
        file.Close()
        if err == nil {
          self.infeed <- nntp
        } else {
          log.Println("failed to create startup messge?", err)
        }
      }
    }
  }()

  // get all pending articles from infeed and load them
  go func() {
    f, err := os.Open(self.store.TempDir())
    if err == nil {
      names, err := f.Readdirnames(0)
      if err == nil {
        for _, name := range names {
          self.infeed_load <- name
        }
      }
    }

  }()

  // if we have no frontend this does nothing
  if self.frontend != nil {
    go self.pollfrontend()
  }
  go self.pollinfeed()
  go self.polloutfeeds()
  self.pollmessages()
}

func (self *NNTPDaemon) pollfrontend() {
  chnl := self.frontend.NewPostsChan()
  for {
    nntp := <- chnl
    // new post from frontend
    log.Println("frontend post", nntp.MessageID())
    self.infeed <- nntp
  }
}
func (self *NNTPDaemon) pollinfeed() {
  for {
    msgid := <- self.infeed_load
    log.Println("load from infeed", msgid)
    msg := self.store.ReadTempMessage(msgid)
    if msg != nil {
      self.infeed <- msg
    }
  }
}

func (self *NNTPDaemon) polloutfeeds() {
  for {
    select {
    case nntp := <- self.send_all_feeds:
      for feed , use := range self.feeds {
        if use && feed.policy != nil {
          if feed.policy.AllowsNewsgroup(nntp.Newsgroup()) {
            if feed.info.mode == "stream" {
              feed.sync <- nntp
            }
          }
        }
      }
    case nntp := <- self.ask_for_article:
      for feed, use := range self.feeds {
        if use && feed.policy != nil {
          if feed.policy.AllowsNewsgroup(nntp.Newsgroup()) {
            if feed.info.mode == "reader" {
              log.Println("ask for", nntp.MessageID())
              feed.sync <- nntp
            }
          }
        }
      }
    }
  }
}

func (self *NNTPDaemon) pollmessages() {
  var chnl chan NNTPMessage
  modchnl := self.mod.MessageChan()
  if self.frontend != nil {
    chnl = self.frontend.PostsChan()
  }
  for {

    nntp := <- self.infeed
    // ammend path
    nntp.AppendPath(self.instance_name)
    msgid := nntp.MessageID()
    log.Println("daemon got", msgid)

    // store article and attachments
    // register with database
    // this also generates thumbnails
    go self.store.StorePost(nntp)

    // prepare for content rollover
    // fallback rollover
    rollover := 100

    group := nntp.Newsgroup()
    tpp, err := self.database.GetThreadsPerPage(group)
    ppb, err := self.database.GetPagesPerBoard(group)
    if err == nil {
      rollover = tpp * ppb
    }

    // roll over old content
    self.expire.ExpireGroup(group, rollover)
    // handle mod events
    if group == "ctl" {
      modchnl <- nntp
    }

    // queue to all outfeeds
    // XXX: blocking ?
    self.send_all_feeds <- ArticleEntry{msgid, group}
    // tell frontend
    // XXX: blocking ?
    if chnl != nil {
      if self.frontend.AllowNewsgroup(group) {
        chnl <- nntp
      }
    }
  }
}

func (self *NNTPDaemon) acceptloop() {
  for {
    // accept
    conn, err := self.listener.Accept()
    if err != nil {
      log.Fatal(err)
    }
    // make a new inbound nntp connection handler
    nntp := self.newConnection(conn, true, nil)
    go self.RunInbound(nntp)
  }
}

func (self *NNTPDaemon) RunInbound(nntp NNTPConnection) {
  nntp.HandleInbound(self)
  delete(self.feeds, nntp)
}

func (self *NNTPDaemon) Setup() {
  log.Println("checking for configs...")
  // check that are configs exist
  CheckConfig()
  log.Println("loading config...")
  // read the config
  self.conf = ReadConfig()
  if self.conf == nil {
    log.Fatal("failed to load config")
  }
  // validate the config
  log.Println("validating configs...")
  self.conf.Validate()
  log.Println("configs are valid")

  db_host := self.conf.database["host"]
  db_port := self.conf.database["port"]
  db_user := self.conf.database["user"]
  db_passwd := self.conf.database["password"]

  // set up database stuff
  log.Println("connecting to database...")
  self.database = NewDatabase(self.conf.database["type"], self.conf.database["schema"], db_host, db_port, db_user, db_passwd)
  log.Println("ensure that the database is created...")
  self.database.CreateTables()

  r_url := self.conf.worker["url"]

  // set up store
  log.Println("set up article store...")
  self.store = createArticleStore(self.conf.store, r_url, self.database)

  self.mod = modEngine{
    store: self.store,
    database:  self.database,
    chnl: make(chan NNTPMessage),
  }
}

// bind to address
func (self *NNTPDaemon) Bind() error {
  listener , err := net.Listen("tcp", self.bind_addr)
  if err != nil {
    log.Println("failed to bind to", self.bind_addr, err)
    return err
  }
  self.listener = listener
  log.Printf("SRNd NNTPD bound at %s", listener.Addr())
  return nil
}

// load configuration
// bind to interface
func (self *NNTPDaemon) Init() bool {

  // set up daemon configs
  self.Setup()

  self.infeed = make(chan NNTPMessage, 8)
  self.infeed_load = make(chan string)
  self.send_all_feeds = make(chan ArticleEntry, 64)
  self.feeds = make(map[NNTPConnection]bool)
  self.ask_for_article = make(chan ArticleEntry, 64)

  self.bind_addr = self.conf.daemon["bind"]

  err := self.Bind()
  if err != nil {
    log.Println("failed to bind:", err)
    return false
  }

  self.expire = createExpirationCore(self.database, self.store)
  self.sync_on_start = self.conf.daemon["sync_on_start"] == "1"
  self.debug = self.conf.daemon["log"] == "debug"
  self.instance_name = self.conf.daemon["instance_name"]
  if self.debug {
    log.Println("debug mode activated")
  }

  // do we enable the frontend?
  if self.conf.frontend["enable"] == "1" {
    log.Printf("frontend %s enabled", self.conf.frontend["name"])
    http_frontend := NewHTTPFrontend(self, self.conf.frontend, self.conf.worker["url"])
    nntp_frontend := NewNNTPFrontend(self, self.conf.frontend["nntp"])
    self.frontend = MuxFrontends(http_frontend, nntp_frontend)
    go self.frontend.Mainloop()
  }

  // set up admin user if it's specified in the config
  pubkey , ok := self.conf.frontend["admin_key"]
  if ok {
    // TODO: check for valid format
    log.Println("add admin key", pubkey)
    err = self.database.MarkModPubkeyGlobal(pubkey)
    if err != nil {
      log.Printf("failed to add admin mod key, %s", err)
    }
  }
  go RunModEngine(self.mod)
  return true
}
<|endoftext|>