<commit_before>package main

import (
	"context"
	"database/sql"
	"log"
	"net/http"
	"sync"
	"time"
)

type RoomNameMap map[RoomID]string
type RoomGroupMap map[BuildingName]map[FloorID][]RoomID

const (
	INTERVAL     = 1 * time.Minute
	CACHE_EXPIRE = 3 * time.Minute
)

type RoomStatus struct {
	RoomID  RoomID         `json:"id"`
	Sensors []SensorStatus `json:"sensors"`

	Hot     uint64 `json:"hot"`
	Comfort uint64 `json:"comfort"`
	Cold    uint64 `json:"cold"`
	lock    sync.RWMutex
}

type MyVote struct {
	Vote      VoteChoice `json:"vote"`
	Timestamp int64      `json:"timestamp"`
}

type RoomStatusManager struct {
	db        *sql.DB
	thingworx *ThingWorxClient

	sensorCache map[RoomID]map[ThingName]SensorStatus
	cacheLock   sync.RWMutex
}

type RoomStatusTx struct {
	rsm *RoomStatusManager
	tx  *sql.Tx

	// nilになる場合があるため、使用前に必ずnilチェックを行うこと。
	s *Session
}

type SensorStatus struct {
	Temperature float64 `json:"temperature"`
	Humidity    float64 `json:"humidity"`
	IsConnected bool    `json:"isConnected"`

	expire time.Time
}

func NewRoomStatusManager(db *sql.DB, thingworx *ThingWorxClient, ctx context.Context) *RoomStatusManager {
	// create RSM
	rs := &RoomStatusManager{}
	rs.db = db
	rs.thingworx = thingworx
	rs.sensorCache = make(map[RoomID]map[ThingName]SensorStatus)

	go rs.cacheUpdater(ctx)
	return rs
}

func (rsm *RoomStatusManager) GetTx(w http.ResponseWriter, req *http.Request, new bool) (*RoomStatusTx, error) {
	tx, err := rsm.db.Begin()
	if err != nil {
		return nil, err
	}
	s := GetSession(w, req, tx)
	if s == nil && new {
		s, err = NewSession(w, req, tx)
		if err != nil {
			return nil, err
		}
	}
	return &RoomStatusTx{
		rsm: rsm,
		tx:  tx,
		s:   s,
	}, nil
}

func (rst *RoomStatusTx) Rollback() error {
	return rst.tx.Rollback()
}

func (rst *RoomStatusTx) Commit() error {
	return rst.tx.Commit()
}

func (rst *RoomStatusTx) GetRoomName(id RoomID) (name string, err error) {
	err = rst.tx.QueryRow(
		`SELECT name FROM room
		WHERE room_id=?`,
		id,
	).Scan(&name)
	return
}

// 投票内容を取得する。未投票の場合はnilを返す
func (rst *RoomStatusTx) GetMyVote(id RoomID) (vote *MyVote, err error) {
	var v Vote

	if rst.s == nil {
		// セッションがnilなので、未投票とみなす
		return nil, nil
	}

	if err = rst.tx.QueryRow(
		`SELECT choice, timestamp FROM vote
			WHERE session_id=? AND room_id=?`,
		rst.s.SessionID, id,
	).Scan((*string)(&v.Choice), &v.Timestamp); err != nil {
		if err == sql.ErrNoRows {
			// 未投票の状態。
			err = nil
			return
		}
		return
	}

	vote = &MyVote{
		Vote:      v.Choice,
		Timestamp: v.Timestamp.Unix(),
	}
	return vote, err
}

func (rst *RoomStatusTx) GetStatus(id RoomID) (*RoomStatus, error) {
	rs := &RoomStatus{
		RoomID: id,
	}

	var ok bool
	rs.Sensors, ok = rst.rsm.getSensorStatusFromCache(id)
	if !ok {
		// センサーの状態を更新できていない状態。
		rs.Sensors = []SensorStatus{}
	}

	rows, err := rst.tx.Query(
		`SELECT vote.choice, count(vote.vote_id) FROM vote
		NATURAL JOIN session
		WHERE vote.room_id=? AND session.expire>=?
		GROUP BY vote.choice`,
		id, time.Now(),
	)
	if err != nil {
		return nil, err
	}
	for rows.Next() {
		var choice VoteChoice
		var count uint64
		if err := rows.Scan((*string)(&choice), &count); err != nil {
			return nil, err
		}
		switch choice {
		case Hot:
			rs.Hot = count
		case Comfort:
			rs.Comfort = count
		case Cold:
			rs.Cold = count
		}
	}
	return rs, nil
}

func (rst *RoomStatusTx) Vote(id RoomID, choice VoteChoice) error {
	if rst.s == nil {
		panic("session must not nil")
	}

	vote := Vote{
		RoomID: id,
		S:      rst.s,
	}

	if err := rst.tx.QueryRow(
		`SELECT vote_id FROM vote
		WHERE session_id=? AND room_id=?`,
		rst.s.SessionID, id,
	).Scan(&vote.VoteID); err != nil {
		switch err {
		case sql.ErrNoRows:
			// 未投票であることを表す、0を代入
			vote.VoteID = VoteID(0)
		default:
			return err
		}
	}

	return vote.UpdateChoice(rst.tx, choice)
}

func (rst *RoomStatusTx) GetAllRoomsInfo() (names RoomNameMap, groups RoomGroupMap, err error) {
	// NOTE: roomテーブルの行数は少ないことを想定しているため、テーブルスキャンをしている。
	{
		names = make(RoomNameMap)
		var rows *sql.Rows
		rows, err = rst.tx.Query(`
			SELECT room_id, name FROM room
		`)
		if err != nil {
			return
		}
		defer rows.Close()
		for rows.Next() {
			var id RoomID
			var name string
			if err = rows.Scan(&id, &name); err != nil {
				return
			}
			names[id] = name
		}
	}

	{
		groups = make(RoomGroupMap)
		var rows *sql.Rows
		rows, err = rst.tx.Query(`
			SELECT building_name, floor, room_id FROM room
			GROUP BY building_name, floor, room_id
		`)
		defer rows.Close()
		for rows.Next() {
			var bname BuildingName
			var floor FloorID
			var id RoomID
			if err = rows.Scan((*string)(&bname), &floor, &id); err != nil {
				return
			}

			if _, ok := groups[bname]; !ok {
				groups[bname] = make(map[FloorID][]RoomID)
			}
			if _, ok := groups[bname][floor]; !ok {
				groups[bname][floor] = make([]RoomID, 0)
			}
			groups[bname][floor] = append(groups[bname][floor], id)
		}
	}
	return
}

func (rsm *RoomStatusManager) getSensorStatusFromCache(id RoomID) ([]SensorStatus, bool) {
	rsm.cacheLock.RLock()
	defer rsm.cacheLock.RUnlock()

	cache, ok := rsm.sensorCache[id]
	if ok {
		array := make([]SensorStatus, 0, len(cache))
		for i := range cache {
			if cache[i].expire.After(time.Now()) {
				array = append(array, cache[i])
			}
		}
		return array, len(array) > 0
	}
	return []SensorStatus{}, false
}

// すべてのセンサーの状態をキャッシュする
func (rsm *RoomStatusManager) cacheUpdater(ctx context.Context) {
	log.Println("starting cacheUpdater")

	tick := time.NewTicker(INTERVAL)
	for {
		log.Println("update all sensor statuses")
		for _, err := range rsm.updateAllSensorStatuses() {
			log.Println(err)
		}

		log.Println("clean up expired sessions")
		if err := rsm.cleanUpExpiredSessions(); err != nil {
			log.Println(err)
		}

		select {
		case <-ctx.Done():
			return
		case <-tick.C:
		}
	}
}

func (rsm *RoomStatusManager) updateAllSensorStatuses() []error {
	errCh := make(chan error)
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		defer wg.Done()
		tx, err := rsm.db.Begin()
		if err != nil {
			errCh <- err
			return
		}
		defer tx.Rollback()

		rows, err := tx.Query(
			`SELECT room_id, thing_name FROM thing`,
		)
		if err != nil {
			errCh <- err
			return
		}

		for rows.Next() {
			var id RoomID
			var name ThingName
			rows.Scan(&id, (*string)(&name))

			// start async update
			wg.Add(1)
			go func(id RoomID, name ThingName) {
				defer wg.Done()
				if err := rsm.updateSensorStatus(id, name); err != nil {
					errCh <- err
					return
				}
			}(id, name)
		}
	}()

	go func() {
		wg.Wait()
		close(errCh)
	}()

	errs := []error{}
	for err := range errCh {
		errs = append(errs, err)
	}
	return errs
}

// センサーで測定した部屋の状態を、DBに反映する。
func (rsm *RoomStatusManager) updateSensorStatus(id RoomID, thingName ThingName) error {
	var stat SensorStatus

	prop, err := rsm.thingworx.Properties(thingName)
	if err != nil {
		return err
	}
	stat.Temperature, err = prop.M("temperature").Float64()
	if err != nil {
		return err
	}
	stat.Humidity, err = prop.M("humidity").Float64()
	if err != nil {
		return err
	}
	stat.IsConnected, err = prop.M("isConnected").Bool()
	if err != nil {
		return err
	}
	stat.expire = time.Now().Add(CACHE_EXPIRE)

	if !stat.IsConnected {
		log.Printf("WARN: \"%s\" is not connected", thingName)
		return nil
	}

	rsm.cacheLock.Lock()
	defer rsm.cacheLock.Unlock()
	if _, ok := rsm.sensorCache[id]; !ok {
		rsm.sensorCache[id] = map[ThingName]SensorStatus{}
	}
	rsm.sensorCache[id][thingName] = stat
	return nil
}

func (rsm *RoomStatusManager) cleanUpExpiredSessions() error {
	tx, err := rsm.db.Begin()
	if err != nil {
		return err
	}
	if _, err := tx.Exec(
		`DELETE FROM session WHERE expire<?`,
		time.Now(),
	); err != nil {
		return err
	}
	if err := tx.Commit(); err != nil {
		return err
	}
	return nil
}
<commit_msg>call Rollback() when failed<commit_after>package main

import (
	"context"
	"database/sql"
	"log"
	"net/http"
	"sync"
	"time"
)

type RoomNameMap map[RoomID]string
type RoomGroupMap map[BuildingName]map[FloorID][]RoomID

const (
	INTERVAL     = 1 * time.Minute
	CACHE_EXPIRE = 3 * time.Minute
)

type RoomStatus struct {
	RoomID  RoomID         `json:"id"`
	Sensors []SensorStatus `json:"sensors"`

	Hot     uint64 `json:"hot"`
	Comfort uint64 `json:"comfort"`
	Cold    uint64 `json:"cold"`
	lock    sync.RWMutex
}

type MyVote struct {
	Vote      VoteChoice `json:"vote"`
	Timestamp int64      `json:"timestamp"`
}

type RoomStatusManager struct {
	db        *sql.DB
	thingworx *ThingWorxClient

	sensorCache map[RoomID]map[ThingName]SensorStatus
	cacheLock   sync.RWMutex
}

type RoomStatusTx struct {
	rsm *RoomStatusManager
	tx  *sql.Tx

	// nilになる場合があるため、使用前に必ずnilチェックを行うこと。
	s *Session
}

type SensorStatus struct {
	Temperature float64 `json:"temperature"`
	Humidity    float64 `json:"humidity"`
	IsConnected bool    `json:"isConnected"`

	expire time.Time
}

func NewRoomStatusManager(db *sql.DB, thingworx *ThingWorxClient, ctx context.Context) *RoomStatusManager {
	// create RSM
	rs := &RoomStatusManager{}
	rs.db = db
	rs.thingworx = thingworx
	rs.sensorCache = make(map[RoomID]map[ThingName]SensorStatus)

	go rs.cacheUpdater(ctx)
	return rs
}

func (rsm *RoomStatusManager) GetTx(w http.ResponseWriter, req *http.Request, new bool) (*RoomStatusTx, error) {
	tx, err := rsm.db.Begin()
	if err != nil {
		return nil, err
	}
	s := GetSession(w, req, tx)
	if s == nil && new {
		s, err = NewSession(w, req, tx)
		if err != nil {
			defer tx.Rollback()
			return nil, err
		}
	}
	return &RoomStatusTx{
		rsm: rsm,
		tx:  tx,
		s:   s,
	}, nil
}

func (rst *RoomStatusTx) Rollback() error {
	return rst.tx.Rollback()
}

func (rst *RoomStatusTx) Commit() error {
	return rst.tx.Commit()
}

func (rst *RoomStatusTx) GetRoomName(id RoomID) (name string, err error) {
	err = rst.tx.QueryRow(
		`SELECT name FROM room
		WHERE room_id=?`,
		id,
	).Scan(&name)
	return
}

// 投票内容を取得する。未投票の場合はnilを返す
func (rst *RoomStatusTx) GetMyVote(id RoomID) (vote *MyVote, err error) {
	var v Vote

	if rst.s == nil {
		// セッションがnilなので、未投票とみなす
		return nil, nil
	}

	if err = rst.tx.QueryRow(
		`SELECT choice, timestamp FROM vote
			WHERE session_id=? AND room_id=?`,
		rst.s.SessionID, id,
	).Scan((*string)(&v.Choice), &v.Timestamp); err != nil {
		if err == sql.ErrNoRows {
			// 未投票の状態。
			err = nil
			return
		}
		return
	}

	vote = &MyVote{
		Vote:      v.Choice,
		Timestamp: v.Timestamp.Unix(),
	}
	return vote, err
}

func (rst *RoomStatusTx) GetStatus(id RoomID) (*RoomStatus, error) {
	rs := &RoomStatus{
		RoomID: id,
	}

	var ok bool
	rs.Sensors, ok = rst.rsm.getSensorStatusFromCache(id)
	if !ok {
		// センサーの状態を更新できていない状態。
		rs.Sensors = []SensorStatus{}
	}

	rows, err := rst.tx.Query(
		`SELECT vote.choice, count(vote.vote_id) FROM vote
		NATURAL JOIN session
		WHERE vote.room_id=? AND session.expire>=?
		GROUP BY vote.choice`,
		id, time.Now(),
	)
	if err != nil {
		return nil, err
	}
	for rows.Next() {
		var choice VoteChoice
		var count uint64
		if err := rows.Scan((*string)(&choice), &count); err != nil {
			return nil, err
		}
		switch choice {
		case Hot:
			rs.Hot = count
		case Comfort:
			rs.Comfort = count
		case Cold:
			rs.Cold = count
		}
	}
	return rs, nil
}

func (rst *RoomStatusTx) Vote(id RoomID, choice VoteChoice) error {
	if rst.s == nil {
		panic("session must not nil")
	}

	vote := Vote{
		RoomID: id,
		S:      rst.s,
	}

	if err := rst.tx.QueryRow(
		`SELECT vote_id FROM vote
		WHERE session_id=? AND room_id=?`,
		rst.s.SessionID, id,
	).Scan(&vote.VoteID); err != nil {
		switch err {
		case sql.ErrNoRows:
			// 未投票であることを表す、0を代入
			vote.VoteID = VoteID(0)
		default:
			return err
		}
	}

	return vote.UpdateChoice(rst.tx, choice)
}

func (rst *RoomStatusTx) GetAllRoomsInfo() (names RoomNameMap, groups RoomGroupMap, err error) {
	// NOTE: roomテーブルの行数は少ないことを想定しているため、テーブルスキャンをしている。
	{
		names = make(RoomNameMap)
		var rows *sql.Rows
		rows, err = rst.tx.Query(`
			SELECT room_id, name FROM room
		`)
		if err != nil {
			return
		}
		defer rows.Close()
		for rows.Next() {
			var id RoomID
			var name string
			if err = rows.Scan(&id, &name); err != nil {
				return
			}
			names[id] = name
		}
	}

	{
		groups = make(RoomGroupMap)
		var rows *sql.Rows
		rows, err = rst.tx.Query(`
			SELECT building_name, floor, room_id FROM room
			GROUP BY building_name, floor, room_id
		`)
		defer rows.Close()
		for rows.Next() {
			var bname BuildingName
			var floor FloorID
			var id RoomID
			if err = rows.Scan((*string)(&bname), &floor, &id); err != nil {
				return
			}

			if _, ok := groups[bname]; !ok {
				groups[bname] = make(map[FloorID][]RoomID)
			}
			if _, ok := groups[bname][floor]; !ok {
				groups[bname][floor] = make([]RoomID, 0)
			}
			groups[bname][floor] = append(groups[bname][floor], id)
		}
	}
	return
}

func (rsm *RoomStatusManager) getSensorStatusFromCache(id RoomID) ([]SensorStatus, bool) {
	rsm.cacheLock.RLock()
	defer rsm.cacheLock.RUnlock()

	cache, ok := rsm.sensorCache[id]
	if ok {
		array := make([]SensorStatus, 0, len(cache))
		for i := range cache {
			if cache[i].expire.After(time.Now()) {
				array = append(array, cache[i])
			}
		}
		return array, len(array) > 0
	}
	return []SensorStatus{}, false
}

// すべてのセンサーの状態をキャッシュする
func (rsm *RoomStatusManager) cacheUpdater(ctx context.Context) {
	log.Println("starting cacheUpdater")

	tick := time.NewTicker(INTERVAL)
	for {
		log.Println("update all sensor statuses")
		for _, err := range rsm.updateAllSensorStatuses() {
			log.Println(err)
		}

		log.Println("clean up expired sessions")
		if err := rsm.cleanUpExpiredSessions(); err != nil {
			log.Println(err)
		}

		select {
		case <-ctx.Done():
			return
		case <-tick.C:
		}
	}
}

func (rsm *RoomStatusManager) updateAllSensorStatuses() []error {
	errCh := make(chan error)
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		defer wg.Done()
		tx, err := rsm.db.Begin()
		if err != nil {
			errCh <- err
			return
		}
		defer tx.Rollback()

		rows, err := tx.Query(
			`SELECT room_id, thing_name FROM thing`,
		)
		if err != nil {
			errCh <- err
			return
		}

		for rows.Next() {
			var id RoomID
			var name ThingName
			rows.Scan(&id, (*string)(&name))

			// start async update
			wg.Add(1)
			go func(id RoomID, name ThingName) {
				defer wg.Done()
				if err := rsm.updateSensorStatus(id, name); err != nil {
					errCh <- err
					return
				}
			}(id, name)
		}
	}()

	go func() {
		wg.Wait()
		close(errCh)
	}()

	errs := []error{}
	for err := range errCh {
		errs = append(errs, err)
	}
	return errs
}

// センサーで測定した部屋の状態を、DBに反映する。
func (rsm *RoomStatusManager) updateSensorStatus(id RoomID, thingName ThingName) error {
	var stat SensorStatus

	prop, err := rsm.thingworx.Properties(thingName)
	if err != nil {
		return err
	}
	stat.Temperature, err = prop.M("temperature").Float64()
	if err != nil {
		return err
	}
	stat.Humidity, err = prop.M("humidity").Float64()
	if err != nil {
		return err
	}
	stat.IsConnected, err = prop.M("isConnected").Bool()
	if err != nil {
		return err
	}
	stat.expire = time.Now().Add(CACHE_EXPIRE)

	if !stat.IsConnected {
		log.Printf("WARN: \"%s\" is not connected", thingName)
		return nil
	}

	rsm.cacheLock.Lock()
	defer rsm.cacheLock.Unlock()
	if _, ok := rsm.sensorCache[id]; !ok {
		rsm.sensorCache[id] = map[ThingName]SensorStatus{}
	}
	rsm.sensorCache[id][thingName] = stat
	return nil
}

func (rsm *RoomStatusManager) cleanUpExpiredSessions() error {
	tx, err := rsm.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if _, err := tx.Exec(
		`DELETE FROM session WHERE expire<?`,
		time.Now(),
	); err != nil {
		return err
	}
	if err := tx.Commit(); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package util

import (
	"compress/bzip2"
	"compress/gzip"
	"io"
	"os"
	"strings"
)

type ReadCloser struct {
	r       io.Reader
	fp      *os.File
	gz      *gzip.Reader
	isOpen  bool
	isGzip  bool
	isBzip2 bool
}

// Opens for reading a plain file, gzip'ed file (extension .gz), or bzip2'ed file (extension .bz2)
func Open(filename string) (rc io.ReadCloser, err error) {
	fp, err := os.Open(filename)
	if err != nil {
		return nil, err
	}

	if strings.HasSuffix(filename, ".gz") {
		r := ReadCloser{
			fp:     fp,
			isGzip: true,
			isOpen: true,
		}
		r.gz, err = gzip.NewReader(fp)
		if err != nil {
			fp.Close()
			return nil, err
		}
		r.r = r.gz
		return r, nil
	}

	if strings.HasSuffix(filename, ".bz2") {
		r := ReadCloser{
			fp:      fp,
			r:       bzip2.NewReader(fp),
			isBzip2: true,
			isOpen:  true,
		}
		return r, nil
	}

	return fp, nil
}

func (rc ReadCloser) Read(p []byte) (n int, err error) {
	if !rc.isOpen {
		panic("ReadCloser is closed")
	}
	return rc.r.Read(p)
}

func (rc ReadCloser) Close() (err error) {
	if rc.isOpen {
		rc.isOpen = false

		if rc.isGzip {
			err = rc.gz.Close()
		}

		e := rc.fp.Close()
		if err == nil {
			err = e
		}
	}
	return
}
<commit_msg>runtime.SetFinalizer in ReadCloser<commit_after>package util

import (
	"compress/bzip2"
	"compress/gzip"
	"io"
	"os"
	"runtime"
	"strings"
)

type ReadCloser struct {
	r       io.Reader
	fp      *os.File
	gz      *gzip.Reader
	isOpen  bool
	isGzip  bool
	isBzip2 bool
}

// Opens for reading a plain file, gzip'ed file (extension .gz), or bzip2'ed file (extension .bz2)
func Open(filename string) (rc io.ReadCloser, err error) {
	fp, err := os.Open(filename)
	if err != nil {
		return nil, err
	}

	if strings.HasSuffix(filename, ".gz") {
		r := ReadCloser{
			fp:     fp,
			isGzip: true,
			isOpen: true,
		}
		r.gz, err = gzip.NewReader(fp)
		if err != nil {
			fp.Close()
			return nil, err
		}
		r.r = r.gz
		runtime.SetFinalizer(&r, (*ReadCloser).Close)
		return r, nil
	}

	if strings.HasSuffix(filename, ".bz2") {
		r := ReadCloser{
			fp:      fp,
			r:       bzip2.NewReader(fp),
			isBzip2: true,
			isOpen:  true,
		}
		runtime.SetFinalizer(&r, (*ReadCloser).Close)
		return r, nil
	}

	return fp, nil
}

func (rc ReadCloser) Read(p []byte) (n int, err error) {
	if !rc.isOpen {
		panic("ReadCloser is closed")
	}
	return rc.r.Read(p)
}

func (rc ReadCloser) Close() (err error) {
	if rc.isOpen {
		rc.isOpen = false

		if rc.isGzip {
			err = rc.gz.Close()
		}

		e := rc.fp.Close()
		if err == nil {
			err = e
		}
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Copyright 2014 Ninja Blocks Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rpc

import (
	"reflect"
	"unicode"
	"unicode/utf8"

	"git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"
)

// ----------------------------------------------------------------------------
// Codec
// ----------------------------------------------------------------------------

// Codec creates a CodecRequest to process each request.
type Codec interface {
	NewRequest(topic string, message mqtt.Message) (CodecRequest, error)
	SendNotification(c *mqtt.MqttClient, topic string, payload ...interface{}) error
}

// CodecRequest decodes a request and encodes a response using a specific
// serialization scheme.
type CodecRequest interface {
	// Reads the request and returns the RPC method name.
	Method() (string, error)
	// Reads the request filling the RPC method args.
	ReadRequest(interface{}) error
	// Writes the response using the RPC method reply.
	WriteResponse(c *mqtt.MqttClient, response interface{})
	// Writes an error produced by the server.
	WriteError(c *mqtt.MqttClient, err error)
}

// ----------------------------------------------------------------------------
// Server
// ----------------------------------------------------------------------------

// NewServer returns a new RPC server.
func NewServer(client *mqtt.MqttClient, codec Codec) *Server {
	return &Server{
		client:   client,
		codec:    codec,
		services: new(serviceMap),
	}
}

// Server serves registered RPC services using registered codecs.
type Server struct {
	client   *mqtt.MqttClient
	codec    Codec
	services *serviceMap
}

type ExportedService struct {
	Methods []string
	topic   string
	server  *Server
}

func (s *ExportedService) SendEvent(event string, payload interface{}) error {
	return s.server.SendNotification(s.topic+"/event/"+event, payload)
}

// RegisterService adds a new service to the server.
//
// The name parameter is optional: if empty it will be inferred from
// the receiver type name.
//
// Methods from the receiver will be extracted if these rules are satisfied:
//
//    - The receiver is exported (begins with an upper case letter) or local
//      (defined in the package registering the service).
//    - The method name is exported.
//    - The method has three arguments: *mqtt.Message, *args, *reply.
//    - All three arguments are pointers.
//    - The second and third arguments are exported or local.
//    - The method has return type error.
//
// All other methods are ignored.
func (s *Server) RegisterService(receiver interface{}, topic string) (service *ExportedService, err error) {

	filter, err := mqtt.NewTopicFilter(topic, 0)
	if err != nil {
		return nil, err
	}

	receipt, err := s.client.StartSubscription(func(client *mqtt.MqttClient, message mqtt.Message) {
		s.serveRequest(topic, message)
	}, filter)

	if err != nil {
		return nil, err
	}

	<-receipt

	exportedMethods, err := s.services.register(receiver, topic)

	var exportedMethodsLower []string

	for _, m := range exportedMethods {
		exportedMethodsLower = append(exportedMethodsLower, lowerFirst(m))
	}

	return &ExportedService{Methods: exportedMethodsLower, topic: topic, server: s}, err
}

func lowerFirst(s string) string {
	if s == "" {
		return ""
	}
	r, n := utf8.DecodeRuneInString(s)
	return string(unicode.ToLower(r)) + s[n:]
}

// SendNotification sends a one-way notification. Perhaps shouldn't be in Server....
func (s *Server) SendNotification(topic string, params ...interface{}) error {
	return s.codec.SendNotification(s.client, topic, params...)
}

// HasMethod returns true if the given method is registered on a topic
func (s *Server) HasMethod(topic string, method string) bool {
	if _, _, err := s.services.get(topic, method); err == nil {
		return true
	}
	return false
}

type Message struct {
	Payload []byte
	Topic   string
}

// ServeRequest handles an incoming Json-RPC MQTT message
func (s *Server) serveRequest(topic string, message mqtt.Message) {

	log.Debugf("Serving request to %s", topic)

	// Create a new codec request.
	codecReq, err := s.codec.NewRequest(topic, message)

	if err != nil {
		codecReq.WriteError(s.client, err)
		return
	}

	// Get service method to be called.
	method, errMethod := codecReq.Method()
	if errMethod != nil {
		codecReq.WriteError(s.client, errMethod)
		return
	}

	serviceSpec, methodSpec, errGet := s.services.get(topic, method)
	if errGet != nil {
		codecReq.WriteError(s.client, errGet)
		return
	}
	// Decode the args.
	args := reflect.New(methodSpec.argsType)
	if errRead := codecReq.ReadRequest(args.Interface()); errRead != nil {
		codecReq.WriteError(s.client, errRead)
		return
	}
	// Call the service method.

	params := []reflect.Value{
		serviceSpec.rcvr,
		reflect.ValueOf(&Message{
			Payload: message.Payload(),
			Topic:   topic,
		}),
	}

	if methodSpec.argsType != nil {
		params = append(params, args)
	}

	/*var reply reflect.Value
	if methodSpec.replyType != nil {
		reply = reflect.New(methodSpec.replyType)
		params = append(params, reply)
	}*/

	retVals := methodSpec.method.Func.Call(params)
	// Cast the last result to error if needed.
	var errResult error
	errInter := retVals[len(retVals)-1].Interface()
	if errInter != nil {
		errResult = errInter.(error)
	}

	var reply reflect.Value
	if methodSpec.replyType != nil {
		reply = retVals[0]
	}

	// Encode the response.
	if errResult == nil {
		if methodSpec.replyType != nil {
			codecReq.WriteResponse(s.client, reply.Interface())
		} else {
			codecReq.WriteResponse(s.client, nil)
		}
	} else {
		codecReq.WriteError(s.client, errResult)
	}
}
<commit_msg>Update the RPC method signature rules documentation<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Copyright 2014 Ninja Blocks Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rpc

import (
	"reflect"
	"unicode"
	"unicode/utf8"

	"git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"
)

// ----------------------------------------------------------------------------
// Codec
// ----------------------------------------------------------------------------

// Codec creates a CodecRequest to process each request.
type Codec interface {
	NewRequest(topic string, message mqtt.Message) (CodecRequest, error)
	SendNotification(c *mqtt.MqttClient, topic string, payload ...interface{}) error
}

// CodecRequest decodes a request and encodes a response using a specific
// serialization scheme.
type CodecRequest interface {
	// Reads the request and returns the RPC method name.
	Method() (string, error)
	// Reads the request filling the RPC method args.
	ReadRequest(interface{}) error
	// Writes the response using the RPC method reply.
	WriteResponse(c *mqtt.MqttClient, response interface{})
	// Writes an error produced by the server.
	WriteError(c *mqtt.MqttClient, err error)
}

// ----------------------------------------------------------------------------
// Server
// ----------------------------------------------------------------------------

// NewServer returns a new RPC server.
func NewServer(client *mqtt.MqttClient, codec Codec) *Server {
	return &Server{
		client:   client,
		codec:    codec,
		services: new(serviceMap),
	}
}

// Server serves registered RPC services using registered codecs.
type Server struct {
	client   *mqtt.MqttClient
	codec    Codec
	services *serviceMap
}

type ExportedService struct {
	Methods []string
	topic   string
	server  *Server
}

func (s *ExportedService) SendEvent(event string, payload interface{}) error {
	return s.server.SendNotification(s.topic+"/event/"+event, payload)
}

// RegisterService adds a new service to the server.
//
// The name parameter is optional: if empty it will be inferred from
// the receiver type name.
//
// Methods from the receiver will be extracted if these rules are satisfied:
//
//    - The receiver is exported (begins with an upper case letter) or local
//      (defined in the package registering the service).
//    - The method name is exported.
//    - The method's first argument is *mqtt.Message
//    - If there is a second argument (the RPC params value) it must be exported and a pointer
//    - If there is a return value, it must be first, exported and a pointer
//    - The method's last return value is an error
//
// All other methods are ignored.
func (s *Server) RegisterService(receiver interface{}, topic string) (service *ExportedService, err error) {

	filter, err := mqtt.NewTopicFilter(topic, 0)
	if err != nil {
		return nil, err
	}

	receipt, err := s.client.StartSubscription(func(client *mqtt.MqttClient, message mqtt.Message) {
		s.serveRequest(topic, message)
	}, filter)

	if err != nil {
		return nil, err
	}

	<-receipt

	exportedMethods, err := s.services.register(receiver, topic)

	var exportedMethodsLower []string

	for _, m := range exportedMethods {
		exportedMethodsLower = append(exportedMethodsLower, lowerFirst(m))
	}

	return &ExportedService{Methods: exportedMethodsLower, topic: topic, server: s}, err
}

func lowerFirst(s string) string {
	if s == "" {
		return ""
	}
	r, n := utf8.DecodeRuneInString(s)
	return string(unicode.ToLower(r)) + s[n:]
}

// SendNotification sends a one-way notification. Perhaps shouldn't be in Server....
func (s *Server) SendNotification(topic string, params ...interface{}) error {
	return s.codec.SendNotification(s.client, topic, params...)
}

// HasMethod returns true if the given method is registered on a topic
func (s *Server) HasMethod(topic string, method string) bool {
	if _, _, err := s.services.get(topic, method); err == nil {
		return true
	}
	return false
}

type Message struct {
	Payload []byte
	Topic   string
}

// ServeRequest handles an incoming Json-RPC MQTT message
func (s *Server) serveRequest(topic string, message mqtt.Message) {

	log.Debugf("Serving request to %s", topic)

	// Create a new codec request.
	codecReq, err := s.codec.NewRequest(topic, message)

	if err != nil {
		codecReq.WriteError(s.client, err)
		return
	}

	// Get service method to be called.
	method, errMethod := codecReq.Method()
	if errMethod != nil {
		codecReq.WriteError(s.client, errMethod)
		return
	}

	serviceSpec, methodSpec, errGet := s.services.get(topic, method)
	if errGet != nil {
		codecReq.WriteError(s.client, errGet)
		return
	}
	// Decode the args.
	args := reflect.New(methodSpec.argsType)
	if errRead := codecReq.ReadRequest(args.Interface()); errRead != nil {
		codecReq.WriteError(s.client, errRead)
		return
	}
	// Call the service method.

	params := []reflect.Value{
		serviceSpec.rcvr,
		reflect.ValueOf(&Message{
			Payload: message.Payload(),
			Topic:   topic,
		}),
	}

	if methodSpec.argsType != nil {
		params = append(params, args)
	}

	/*var reply reflect.Value
	if methodSpec.replyType != nil {
		reply = reflect.New(methodSpec.replyType)
		params = append(params, reply)
	}*/

	retVals := methodSpec.method.Func.Call(params)
	// Cast the last result to error if needed.
	var errResult error
	errInter := retVals[len(retVals)-1].Interface()
	if errInter != nil {
		errResult = errInter.(error)
	}

	var reply reflect.Value
	if methodSpec.replyType != nil {
		reply = retVals[0]
	}

	// Encode the response.
	if errResult == nil {
		if methodSpec.replyType != nil {
			codecReq.WriteResponse(s.client, reply.Interface())
		} else {
			codecReq.WriteResponse(s.client, nil)
		}
	} else {
		codecReq.WriteError(s.client, errResult)
	}
}
<|endoftext|>


<commit_before>package rpc

import (
	"code.google.com/p/go-uuid/uuid"
	zmq "github.com/pebbe/zmq4"
)

type Server struct {
	id            string
	uri           string
	running       bool
	remoteSocket  *zmq.Socket
	workersSocket *zmq.Socket
}

func NewServer(uri string) *Server {
	server := &Server{
		id:  uuid.New(),
		uri: uri,
	}

	return server
}

func (s *Server) Start() error {
	var err error

	s.remoteSocket, err = zmq.NewSocket(zmq.ROUTER)
	if err != nil {
		return err
	}

	s.remoteSocket.Bind(s.uri)
	if err != nil {
		return err
	}

	s.workersSocket, err = zmq.NewSocket(zmq.ROUTER)
	if err != nil {
		return err
	}

	err = s.workersSocket.Bind("inproc://" + s.id)
	if err != nil {
		return err
	}

	poller := zmq.NewPoller()
	poller.Add(s.remoteSocket, zmq.POLLIN)
	poller.Add(s.workersSocket, zmq.POLLIN)

	s.running = true

	for s.running {
		polled, err := poller.Poll(100)
		if err != nil {
			// TODO Log a warning
		}

		if len(polled) > 0 {
			s.handleSockets(polled)
		}
	}

	return nil
}

func (s *Server) Stop() {
	if s.running {
		s.running = false
	}
}

func (s *Server) handleSockets(polled []zmq.Polled) {
	for _, p := range polled {
		switch socket := p.Socket; socket {
		case s.remoteSocket:
			s.handleRemoteSocket()
		case s.workersSocket:
			s.handleWorkersSocket()
		}
	}
}

func (s *Server) handleRemoteSocket() {
	// TODO Receive message and route it to one of available workers
}

func (s *Server) handleWorkersSocket() {
	// TODO Receive message and route it back to the remote client
}
<commit_msg>Server: add error as return value.<commit_after>package rpc

import (
	"code.google.com/p/go-uuid/uuid"
	"errors"
	zmq "github.com/pebbe/zmq4"
)

type Server struct {
	id            string
	uri           string
	running       bool
	remoteSocket  *zmq.Socket
	workersSocket *zmq.Socket
}

func NewServer(uri string) *Server {
	server := &Server{
		id:  uuid.New(),
		uri: uri,
	}

	return server
}

func (s *Server) Start() error {
	var err error

	s.remoteSocket, err = zmq.NewSocket(zmq.ROUTER)
	if err != nil {
		return err
	}

	s.remoteSocket.Bind(s.uri)
	if err != nil {
		return err
	}

	s.workersSocket, err = zmq.NewSocket(zmq.ROUTER)
	if err != nil {
		return err
	}

	err = s.workersSocket.Bind("inproc://" + s.id)
	if err != nil {
		return err
	}

	poller := zmq.NewPoller()
	poller.Add(s.remoteSocket, zmq.POLLIN)
	poller.Add(s.workersSocket, zmq.POLLIN)

	s.running = true

	for s.running {
		polled, err := poller.Poll(100)
		if err != nil {
			// TODO Log a warning
		}

		if len(polled) > 0 {
			s.handleSockets(polled)
		}
	}

	return nil
}

func (s *Server) Stop() error {
	if !s.running {
		return errors.New("Server is already stopped")
	}

	s.running = false
	return nil
}

func (s *Server) handleSockets(polled []zmq.Polled) {
	for _, p := range polled {
		switch socket := p.Socket; socket {
		case s.remoteSocket:
			s.handleRemoteSocket()
		case s.workersSocket:
			s.handleWorkersSocket()
		}
	}
}

func (s *Server) handleRemoteSocket() {
	// TODO Receive message and route it to one of available workers
}

func (s *Server) handleWorkersSocket() {
	// TODO Receive message and route it back to the remote client
}
<|endoftext|>


<commit_before>package logpeck

import (
	"github.com/Sirupsen/logrus"
	"strconv"
)

type AggregatorConfig struct {
	Tags         []string `json:"tags"`
	Aggregations []string `json:"aggregations"`
	Target       string   `json:"target"`
	Time         string   `json:"time"`
}

type Aggregator struct {
	interval          int64
	name              string
	aggregatorConfigs map[string]AggregatorConfig
	buckets           map[string]map[string][]int64
	postTime          int64
}

func NewAggregator(interval int64, name string, aggregators *map[string]AggregatorConfig) *Aggregator {
	aggregator := &Aggregator{
		interval:          interval,
		name:              name,
		aggregatorConfigs: *aggregators,
		buckets:           make(map[string]map[string][]int64),
		postTime:          0,
	}
	return aggregator
}

func getSampleTime(ts int64, interval int64) int64 {
	return ts / interval
}

func (p *Aggregator) IsDeadline(timestamp int64) bool {
	interval := p.interval
	nowTime := getSampleTime(timestamp, interval)
	if p.postTime != nowTime {
		return true
	}
	return false
}

func (p *Aggregator) Record(fields map[string]interface{}) int64 {
	//get sender
	//influxDbConfig := p.Config.SenderConfig.Config.(InfluxDbConfig)
	bucketName := fields[p.name].(string)
	bucketTag := ""
	aggregatorConfig := p.aggregatorConfigs[bucketName]
	tags := aggregatorConfig.Tags
	aggregations := aggregatorConfig.Aggregations
	target := aggregatorConfig.Target
	time := aggregatorConfig.Time
	for i := 0; i < len(tags); i++ {
		bucketTag += "," + tags[i] + "=" + fields[tags[i]].(string)
	}
	int_bool := false
	for i := 0; i < len(aggregations); i++ {
		if aggregations[i] != "cnt" {
			int_bool = true
		}
	}
	aggValue := fields[target].(string)

	//get time
	now, err := strconv.ParseInt(fields[time].(string), 10, 64)
	if err != nil {
		logrus.Debug("[Record] timestamp:%v can't use strconv.ParseInt", fields[time].(string))
		return now
	}

	if _, ok := p.buckets[bucketName]; !ok {
		p.buckets[bucketName] = make(map[string][]int64)
	}
	if int_bool == false {
		p.buckets[bucketName][bucketTag] = append(p.buckets[bucketName][bucketTag], 1)
	} else {
		aggValue, err := strconv.ParseInt(aggValue, 10, 64)
		if err != nil {
			logrus.Debug("[Record] target:%v can't use strconv.ParseInt", aggValue)
			return now
		}
		p.buckets[bucketName][bucketTag] = append(p.buckets[bucketName][bucketTag], aggValue)
	}
	return now
}

func quickSort(values []int64, left, right int) {

	temp := values[left]
	p := left
	i, j := left, right

	for i <= j {
		for j >= p && values[j] >= temp {
			j--
		}
		if j >= p {
			values[p] = values[j]
			p = j
		}

		for i <= p && values[i] <= temp {
			i++
		}
		if i <= p {
			values[p] = values[i]
			p = i
		}

	}

	values[p] = temp

	if p-left > 1 {
		quickSort(values, left, p-1)
	}
	if right-p > 1 {
		quickSort(values, p+1, right)
	}
}

func getAggregation(targetValue []int64, aggregations []string) map[string]int64 {
	aggregationResults := map[string]int64{}
	cnt := int64(len(targetValue))
	avg := int64(0)
	sum := int64(0)
	quickSort(targetValue, 0, len(targetValue)-1)
	for _, value := range targetValue {
		sum += value
	}
	avg = sum / cnt
	for i := 0; i < len(aggregations); i++ {
		switch aggregations[i] {
		case "cnt":
			aggregationResults["cnt"] = int64(len(targetValue))
		case "avg":
			aggregationResults["avg"] = avg
		default:
			if aggregations[i][0] == 'p' {
				proportion, err := strconv.ParseInt(aggregations[i][1:], 10, 64)
				if err != nil {
					panic(aggregations[i])
				}
				percentile := targetValue[cnt*proportion/100-1]
				aggregationResults[aggregations[i]] = percentile
			}
		}
	}
	return aggregationResults
}

func (p *Aggregator) Dump(timestamp int64) map[string]interface{} {
	fields := map[string]interface{}{}
	//now := strconv.FormatInt(timestamp, 10)
	for bucketName, bucketTag_value := range p.buckets {
		for bucketTag, targetValue := range bucketTag_value {
			aggregations := p.aggregatorConfigs[bucketName].Aggregations
			fields[bucketName+bucketTag] = getAggregation(targetValue, aggregations)
		}
	}
	fields["timestamp"] = timestamp
	p.postTime = getSampleTime(timestamp, p.interval)
	p.buckets = map[string]map[string][]int64{}
	return fields
}
<commit_msg>Modify cnt<commit_after>package logpeck

import (
	"github.com/Sirupsen/logrus"
	"strconv"
)

type AggregatorConfig struct {
	Tags         []string `json:"tags"`
	Aggregations []string `json:"aggregations"`
	Target       string   `json:"target"`
	Time         string   `json:"time"`
}

type Aggregator struct {
	interval          int64
	name              string
	aggregatorConfigs map[string]AggregatorConfig
	buckets           map[string]map[string][]int64
	postTime          int64
}

func NewAggregator(interval int64, name string, aggregators *map[string]AggregatorConfig) *Aggregator {
	aggregator := &Aggregator{
		interval:          interval,
		name:              name,
		aggregatorConfigs: *aggregators,
		buckets:           make(map[string]map[string][]int64),
		postTime:          0,
	}
	return aggregator
}

func getSampleTime(ts int64, interval int64) int64 {
	return ts / interval
}

func (p *Aggregator) IsDeadline(timestamp int64) bool {
	interval := p.interval
	nowTime := getSampleTime(timestamp, interval)
	if p.postTime != nowTime {
		return true
	}
	return false
}

func (p *Aggregator) Record(fields map[string]interface{}) int64 {
	//get sender
	//influxDbConfig := p.Config.SenderConfig.Config.(InfluxDbConfig)
	bucketName := fields[p.name].(string)
	bucketTag := ""
	aggregatorConfig := p.aggregatorConfigs[bucketName]
	tags := aggregatorConfig.Tags
	aggregations := aggregatorConfig.Aggregations
	target := aggregatorConfig.Target
	time := aggregatorConfig.Time
	for i := 0; i < len(tags); i++ {
		bucketTag += "," + tags[i] + "=" + fields[tags[i]].(string)
	}
	int_bool := false
	for i := 0; i < len(aggregations); i++ {
		if aggregations[i] != "cnt" {
			int_bool = true
		}
	}
	aggValue := fields[target].(string)

	//get time
	now, err := strconv.ParseInt(fields[time].(string), 10, 64)
	if err != nil {
		logrus.Debug("[Record] timestamp:%v can't use strconv.ParseInt", fields[time].(string))
		return now
	}

	if _, ok := p.buckets[bucketName]; !ok {
		p.buckets[bucketName] = make(map[string][]int64)
	}
	if int_bool == false {
		aggValue, err := strconv.ParseInt(aggValue, 10, 64)
		if err != nil {
			p.buckets[bucketName][bucketTag] = append(p.buckets[bucketName][bucketTag], 1)
		}else {
			p.buckets[bucketName][bucketTag] = append(p.buckets[bucketName][bucketTag], aggValue)
		}
	} else {
		aggValue, err := strconv.ParseInt(aggValue, 10, 64)
		if err != nil {
			logrus.Debug("[Record] target:%v can't use strconv.ParseInt", aggValue)
			return now
		}
		p.buckets[bucketName][bucketTag] = append(p.buckets[bucketName][bucketTag], aggValue)
	}
	return now
}

func quickSort(values []int64, left, right int) {

	temp := values[left]
	p := left
	i, j := left, right

	for i <= j {
		for j >= p && values[j] >= temp {
			j--
		}
		if j >= p {
			values[p] = values[j]
			p = j
		}

		for i <= p && values[i] <= temp {
			i++
		}
		if i <= p {
			values[p] = values[i]
			p = i
		}

	}

	values[p] = temp

	if p-left > 1 {
		quickSort(values, left, p-1)
	}
	if right-p > 1 {
		quickSort(values, p+1, right)
	}
}

func getAggregation(targetValue []int64, aggregations []string) map[string]int64 {
	aggregationResults := map[string]int64{}
	cnt := int64(len(targetValue))
	avg := int64(0)
	sum := int64(0)
	quickSort(targetValue, 0, len(targetValue)-1)
	for _, value := range targetValue {
		sum += value
	}
	avg = sum / cnt
	for i := 0; i < len(aggregations); i++ {
		switch aggregations[i] {
		case "cnt":
			aggregationResults["cnt"] = int64(len(targetValue))
		case "avg":
			aggregationResults["avg"] = avg
		default:
			if aggregations[i][0] == 'p' {
				proportion, err := strconv.ParseInt(aggregations[i][1:], 10, 64)
				if err != nil {
					panic(aggregations[i])
				}
				percentile := targetValue[cnt*proportion/100-1]
				aggregationResults[aggregations[i]] = percentile
			}
		}
	}
	return aggregationResults
}

func (p *Aggregator) Dump(timestamp int64) map[string]interface{} {
	fields := map[string]interface{}{}
	//now := strconv.FormatInt(timestamp, 10)
	for bucketName, bucketTag_value := range p.buckets {
		for bucketTag, targetValue := range bucketTag_value {
			aggregations := p.aggregatorConfigs[bucketName].Aggregations
			fields[bucketName+bucketTag] = getAggregation(targetValue, aggregations)
		}
	}
	fields["timestamp"] = timestamp
	p.postTime = getSampleTime(timestamp, p.interval)
	p.buckets = map[string]map[string][]int64{}
	return fields
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/andrew-d/go-termutil"
	"github.com/codegangsta/cli"
	"github.com/garyburd/redigo/redis"
)

//Pop all values from the given redis list and write the values to stdout
func readAll(list string, client redis.Conn) {
	for {
		read(list, client, 64)
	}
}

func read(list string, conn redis.Conn, count int) {
	for i := 0; i < count; i++ {
		conn.Send("LPOP", list)
	}
	conn.Flush()
	for i := 0; i < count; i++ {
		value, err := redis.String(conn.Receive())

		if err != nil {
			os.Exit(0)
		}
		if value != "" {
			fmt.Println(value)
		}
	}
}

//Push all values from stdin to a given redis list
func write(list string, conn redis.Conn) {
	scanner := bufio.NewScanner(os.Stdin)
	for i := 0; i < 64; i++ {
		for scanner.Scan() {
			value := strings.TrimSpace(scanner.Text())
			if value != "" {
				conn.Send("LPUSH", list, value)
			}
		}
	}
	err := conn.Flush()
	if err != nil {
		log.Fatalf("Could not LPUSH: %s \"%s\"\n", list, err.Error())
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "redis-pipe"
	app.Usage = "connect stdin with LPUSH and LPOP with stout."
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "host",
			Value:  "localhost",
			Usage:  "Redis host",
			EnvVar: "REDIS_HOST",
		},
		cli.IntFlag{
			Name:   "port",
			Value:  6379,
			Usage:  "Redis port",
			EnvVar: "REDIS_PORT",
		},
		cli.IntFlag{
			Name:  "count",
			Value: -1,
			Usage: "Redis Stop reading from list after count",
		},
	}

	app.Action = func(c *cli.Context) {
		list := c.Args().First()
		if list == "" {
			fmt.Println("Please provide name of the list")
			os.Exit(1)
		}

		conn, err := redis.Dial("tcp", fmt.Sprintf("%s:%d", c.String("host"), c.Int("port")))
		if err != nil {
			log.Fatalf("Connecting to Redis failed: %s\n", err.Error())
		}

		if termutil.Isatty(os.Stdin.Fd()) {
			count := c.Int("count")
			if count > 0 {
				read(list, conn, count)
			} else {
				readAll(list, conn)
			}
		} else {
			write(list, conn)
		}
		conn.Close()
	}

	app.Run(os.Args)
}
<commit_msg>Fixed bug with writing to Redis list<commit_after>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/andrew-d/go-termutil"
	"github.com/codegangsta/cli"
	"github.com/garyburd/redigo/redis"
)

//Pop all values from the given redis list and write the values to stdout
func readAll(list string, client redis.Conn) {
	for {
		read(list, client, 64)
	}
}

func read(list string, conn redis.Conn, count int) {
	for i := 0; i < count; i++ {
		conn.Send("LPOP", list)
	}
	conn.Flush()
	for i := 0; i < count; i++ {
		value, err := redis.String(conn.Receive())

		if err != nil {
			os.Exit(0)
		}
		if value != "" {
			fmt.Println(value)
		}
	}
}

func writeAll(list string, client redis.Conn) {
	for {
		write(list, client, 16)
	}
}

//Push all values from stdin to a given redis list
func write(list string, conn redis.Conn, count int) {
	scanner := bufio.NewScanner(os.Stdin)
	for i := 0; i < count; i++ {
		for scanner.Scan() {
			value := strings.TrimSpace(scanner.Text())
			if value != "" {
				conn.Send("LPUSH", list, value)
			}
		}
	}
	conn.Flush()
	for i := 0; i < count; i++ {
		_, err := redis.Int(conn.Receive())
		if err != nil {
			log.Fatalf("Could not LPUSH: %s \"%s\"\n", list, err.Error())
		}
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "redis-pipe"
	app.Usage = "connect stdin with LPUSH and LPOP with stout."
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "host",
			Value:  "localhost",
			Usage:  "Redis host",
			EnvVar: "REDIS_HOST",
		},
		cli.IntFlag{
			Name:   "port",
			Value:  6379,
			Usage:  "Redis port",
			EnvVar: "REDIS_PORT",
		},
		cli.IntFlag{
			Name:  "count",
			Value: -1,
			Usage: "Redis Stop reading from list after count",
		},
	}

	app.Action = func(c *cli.Context) {
		list := c.Args().First()
		if list == "" {
			fmt.Println("Please provide name of the list")
			os.Exit(1)
		}

		conn, err := redis.Dial("tcp", fmt.Sprintf("%s:%d", c.String("host"), c.Int("port")))
		if err != nil {
			log.Fatalf("Connecting to Redis failed: %s\n", err.Error())
		}

		if termutil.Isatty(os.Stdin.Fd()) {
			count := c.Int("count")
			if count > 0 {
				read(list, conn, count)
			} else {
				readAll(list, conn)
			}
		} else {
			writeAll(list, conn)
		}
		conn.Close()
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package srtp

import (
	"bytes"
	"testing"

	"github.com/pkg/errors"
)

const cipherContextAlgo = "SRTP_AES128_CM_SHA1_80"
const defaultSsrc = 0

func TestKeyLen(t *testing.T) {
	if _, err := CreateContext([]byte{}, make([]byte, saltLen), cipherContextAlgo, defaultSsrc); err == nil {
		t.Errorf("CreateContext accepted a 0 length key")
	}

	if _, err := CreateContext(make([]byte, keyLen), []byte{}, cipherContextAlgo, defaultSsrc); err == nil {
		t.Errorf("CreateContext accepted a 0 length salt")
	}

	if _, err := CreateContext(make([]byte, keyLen), make([]byte, saltLen), cipherContextAlgo, defaultSsrc); err != nil {
		t.Error(errors.Wrap(err, "CreateContext failed with a valid length key and salt"))
	}
}

func TestValidSessionKeys(t *testing.T) {
	masterKey := []byte{0xE1, 0xF9, 0x7A, 0x0D, 0x3E, 0x01, 0x8B, 0xE0, 0xD6, 0x4F, 0xA3, 0x2C, 0x06, 0xDE, 0x41, 0x39}
	masterSalt := []byte{0x0E, 0xC6, 0x75, 0xAD, 0x49, 0x8A, 0xFE, 0xEB, 0xB6, 0x96, 0x0B, 0x3A, 0xAB, 0xE6}

	expectedSessionKey := []byte{0xC6, 0x1E, 0x7A, 0x93, 0x74, 0x4F, 0x39, 0xEE, 0x10, 0x73, 0x4A, 0xFE, 0x3F, 0xF7, 0xA0, 0x87}
	expectedSessionSalt := []byte{0x30, 0xCB, 0xBC, 0x08, 0x86, 0x3D, 0x8C, 0x85, 0xD4, 0x9D, 0xB3, 0x4A, 0x9A, 0xE1}
	expectedSessionAuthTag := []byte{0xCE, 0xBE, 0x32, 0x1F, 0x6F, 0xF7, 0x71, 0x6B, 0x6F, 0xD4, 0xAB, 0x49, 0xAF, 0x25, 0x6A, 0x15, 0x6D, 0x38, 0xBA, 0xA4}

	c, err := CreateContext(masterKey, masterSalt, cipherContextAlgo, defaultSsrc)
	if err != nil {
		t.Error(errors.Wrap(err, "CreateContext failed"))
	}

	sessionKey, err := c.generateSessionKey()
	if err != nil {
		t.Error(errors.Wrap(err, "generateSessionKey failed"))
	} else if !bytes.Equal(sessionKey, expectedSessionKey) {
		t.Errorf("Session Key % 02x does not match expected % 02x", sessionKey, expectedSessionKey)
	}

	sessionSalt, err := c.generateSessionSalt()
	if err != nil {
		t.Error(errors.Wrap(err, "generateSessionSalt failed"))
	} else if !bytes.Equal(sessionSalt, expectedSessionSalt) {
		t.Errorf("Session Salt % 02x does not match expected % 02x", sessionSalt, expectedSessionSalt)
	}

	sessionAuthTag, err := c.generateSessionAuthTag()
	if err != nil {
		t.Error(errors.Wrap(err, "generateSessionAuthTag failed"))
	} else if !bytes.Equal(sessionAuthTag, expectedSessionAuthTag) {
		t.Errorf("Session Auth Tag % 02x does not match expected % 02x", sessionAuthTag, expectedSessionAuthTag)
	}

}

func TestValidPacketCounter(t *testing.T) {
	masterKey := []byte{0x0d, 0xcd, 0x21, 0x3e, 0x4c, 0xbc, 0xf2, 0x8f, 0x01, 0x7f, 0x69, 0x94, 0x40, 0x1e, 0x28, 0x89}
	masterSalt := []byte{0x62, 0x77, 0x60, 0x38, 0xc0, 0x6d, 0xc9, 0x41, 0x9f, 0x6d, 0xd9, 0x43, 0x3e, 0x7c}

	c, err := CreateContext(masterKey, masterSalt, cipherContextAlgo, defaultSsrc)
	if err != nil {
		t.Error(errors.Wrap(err, "CreateContext failed"))
	}

	c.ssrc = 4160032510
	expectedCounter := []byte{0xcf, 0x90, 0x1e, 0xa5, 0xda, 0xd3, 0x2c, 0x15, 0x00, 0xa2, 0x24, 0xae, 0xae, 0xaf, 0x00, 0x00}
	counter := c.generateCounter(32846)
	if !bytes.Equal(counter, expectedCounter) {
		t.Errorf("Session Key % 02x does not match expected % 02x", counter, expectedCounter)
	}
}

func TestRolloverCount(t *testing.T) {
	masterKey := []byte{0x0d, 0xcd, 0x21, 0x3e, 0x4c, 0xbc, 0xf2, 0x8f, 0x01, 0x7f, 0x69, 0x94, 0x40, 0x1e, 0x28, 0x89}
	masterSalt := []byte{0x62, 0x77, 0x60, 0x38, 0xc0, 0x6d, 0xc9, 0x41, 0x9f, 0x6d, 0xd9, 0x43, 0x3e, 0x7c}

	c, err := CreateContext(masterKey, masterSalt, cipherContextAlgo, defaultSsrc)
	if err != nil {
		t.Error(errors.Wrap(err, "CreateContext failed"))
	}

	// Set initial seqnum
	c.updateRolloverCount(65530)

	// We rolled over to 0
	c.updateRolloverCount(0)
	if c.rolloverCounter != 1 {
		t.Errorf("rolloverCounter was not updated after it crossed 0")
	}

	c.updateRolloverCount(65530)
	if c.rolloverCounter != 0 {
		t.Errorf("rolloverCounter was not updated when it rolled back, failed to handle out of order")
	}

	c.updateRolloverCount(5)
	if c.rolloverCounter != 1 {
		t.Errorf("rolloverCounter was not updated when it rolled over initial, to handle out of order")
	}

	c.updateRolloverCount(6)
	c.updateRolloverCount(7)
	c.updateRolloverCount(8)
	if c.rolloverCounter != 1 {
		t.Errorf("rolloverCounter was improperly updated for non-significant packets")
	}
}
<commit_msg>Add full SRTP Lifecycle tests<commit_after>package srtp

import (
	"bytes"
	"testing"

	"github.com/pions/webrtc/pkg/rtp"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
)

const cipherContextAlgo = "SRTP_AES128_CM_SHA1_80"
const defaultSsrc = 0

type rtpTestCase struct {
	sequenceNumber uint16
	encrypted      []byte
}

func TestKeyLen(t *testing.T) {
	if _, err := CreateContext([]byte{}, make([]byte, saltLen), cipherContextAlgo, defaultSsrc); err == nil {
		t.Errorf("CreateContext accepted a 0 length key")
	}

	if _, err := CreateContext(make([]byte, keyLen), []byte{}, cipherContextAlgo, defaultSsrc); err == nil {
		t.Errorf("CreateContext accepted a 0 length salt")
	}

	if _, err := CreateContext(make([]byte, keyLen), make([]byte, saltLen), cipherContextAlgo, defaultSsrc); err != nil {
		t.Error(errors.Wrap(err, "CreateContext failed with a valid length key and salt"))
	}
}

func TestValidSessionKeys(t *testing.T) {
	masterKey := []byte{0xE1, 0xF9, 0x7A, 0x0D, 0x3E, 0x01, 0x8B, 0xE0, 0xD6, 0x4F, 0xA3, 0x2C, 0x06, 0xDE, 0x41, 0x39}
	masterSalt := []byte{0x0E, 0xC6, 0x75, 0xAD, 0x49, 0x8A, 0xFE, 0xEB, 0xB6, 0x96, 0x0B, 0x3A, 0xAB, 0xE6}

	expectedSessionKey := []byte{0xC6, 0x1E, 0x7A, 0x93, 0x74, 0x4F, 0x39, 0xEE, 0x10, 0x73, 0x4A, 0xFE, 0x3F, 0xF7, 0xA0, 0x87}
	expectedSessionSalt := []byte{0x30, 0xCB, 0xBC, 0x08, 0x86, 0x3D, 0x8C, 0x85, 0xD4, 0x9D, 0xB3, 0x4A, 0x9A, 0xE1}
	expectedSessionAuthTag := []byte{0xCE, 0xBE, 0x32, 0x1F, 0x6F, 0xF7, 0x71, 0x6B, 0x6F, 0xD4, 0xAB, 0x49, 0xAF, 0x25, 0x6A, 0x15, 0x6D, 0x38, 0xBA, 0xA4}

	c, err := CreateContext(masterKey, masterSalt, cipherContextAlgo, defaultSsrc)
	if err != nil {
		t.Error(errors.Wrap(err, "CreateContext failed"))
	}

	sessionKey, err := c.generateSessionKey()
	if err != nil {
		t.Error(errors.Wrap(err, "generateSessionKey failed"))
	} else if !bytes.Equal(sessionKey, expectedSessionKey) {
		t.Errorf("Session Key % 02x does not match expected % 02x", sessionKey, expectedSessionKey)
	}

	sessionSalt, err := c.generateSessionSalt()
	if err != nil {
		t.Error(errors.Wrap(err, "generateSessionSalt failed"))
	} else if !bytes.Equal(sessionSalt, expectedSessionSalt) {
		t.Errorf("Session Salt % 02x does not match expected % 02x", sessionSalt, expectedSessionSalt)
	}

	sessionAuthTag, err := c.generateSessionAuthTag()
	if err != nil {
		t.Error(errors.Wrap(err, "generateSessionAuthTag failed"))
	} else if !bytes.Equal(sessionAuthTag, expectedSessionAuthTag) {
		t.Errorf("Session Auth Tag % 02x does not match expected % 02x", sessionAuthTag, expectedSessionAuthTag)
	}

}

func TestValidPacketCounter(t *testing.T) {
	masterKey := []byte{0x0d, 0xcd, 0x21, 0x3e, 0x4c, 0xbc, 0xf2, 0x8f, 0x01, 0x7f, 0x69, 0x94, 0x40, 0x1e, 0x28, 0x89}
	masterSalt := []byte{0x62, 0x77, 0x60, 0x38, 0xc0, 0x6d, 0xc9, 0x41, 0x9f, 0x6d, 0xd9, 0x43, 0x3e, 0x7c}

	c, err := CreateContext(masterKey, masterSalt, cipherContextAlgo, defaultSsrc)
	if err != nil {
		t.Error(errors.Wrap(err, "CreateContext failed"))
	}

	c.ssrc = 4160032510
	expectedCounter := []byte{0xcf, 0x90, 0x1e, 0xa5, 0xda, 0xd3, 0x2c, 0x15, 0x00, 0xa2, 0x24, 0xae, 0xae, 0xaf, 0x00, 0x00}
	counter := c.generateCounter(32846)
	if !bytes.Equal(counter, expectedCounter) {
		t.Errorf("Session Key % 02x does not match expected % 02x", counter, expectedCounter)
	}
}

func TestRolloverCount(t *testing.T) {
	masterKey := []byte{0x0d, 0xcd, 0x21, 0x3e, 0x4c, 0xbc, 0xf2, 0x8f, 0x01, 0x7f, 0x69, 0x94, 0x40, 0x1e, 0x28, 0x89}
	masterSalt := []byte{0x62, 0x77, 0x60, 0x38, 0xc0, 0x6d, 0xc9, 0x41, 0x9f, 0x6d, 0xd9, 0x43, 0x3e, 0x7c}

	c, err := CreateContext(masterKey, masterSalt, cipherContextAlgo, defaultSsrc)
	if err != nil {
		t.Error(errors.Wrap(err, "CreateContext failed"))
	}

	// Set initial seqnum
	c.updateRolloverCount(65530)

	// We rolled over to 0
	c.updateRolloverCount(0)
	if c.rolloverCounter != 1 {
		t.Errorf("rolloverCounter was not updated after it crossed 0")
	}

	c.updateRolloverCount(65530)
	if c.rolloverCounter != 0 {
		t.Errorf("rolloverCounter was not updated when it rolled back, failed to handle out of order")
	}

	c.updateRolloverCount(5)
	if c.rolloverCounter != 1 {
		t.Errorf("rolloverCounter was not updated when it rolled over initial, to handle out of order")
	}

	c.updateRolloverCount(6)
	c.updateRolloverCount(7)
	c.updateRolloverCount(8)
	if c.rolloverCounter != 1 {
		t.Errorf("rolloverCounter was improperly updated for non-significant packets")
	}
}

func TestRTPLifecyle(t *testing.T) {
	assert := assert.New(t)
	masterKey := []byte{0x0d, 0xcd, 0x21, 0x3e, 0x4c, 0xbc, 0xf2, 0x8f, 0x01, 0x7f, 0x69, 0x94, 0x40, 0x1e, 0x28, 0x89}
	masterSalt := []byte{0x62, 0x77, 0x60, 0x38, 0xc0, 0x6d, 0xc9, 0x41, 0x9f, 0x6d, 0xd9, 0x43, 0x3e, 0x7c}

	encryptContext, err := CreateContext(masterKey, masterSalt, cipherContextAlgo, defaultSsrc)
	if err != nil {
		t.Error(errors.Wrap(err, "CreateContext failed"))
	}

	decryptContext, err := CreateContext(masterKey, masterSalt, cipherContextAlgo, defaultSsrc)
	if err != nil {
		t.Error(errors.Wrap(err, "CreateContext failed"))
	}

	decrypted := []byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05}
	var testCases = []rtpTestCase{
		{
			sequenceNumber: 5000,
			encrypted:      []byte{0x6d, 0xd3, 0x7e, 0xd5, 0x99, 0xb7, 0x2d, 0x28, 0xb1, 0xf3, 0xa1, 0xf0, 0xc, 0xfb, 0xfd, 0x8},
		},
		{
			sequenceNumber: 5001,
			encrypted:      []byte{0xda, 0x47, 0xb, 0x2a, 0x74, 0x53, 0x65, 0xbd, 0x2f, 0xeb, 0xdc, 0x4b, 0x6d, 0x23, 0xf3, 0xde},
		},
		{
			sequenceNumber: 5002,
			encrypted:      []byte{0x6e, 0xa7, 0x69, 0x8d, 0x24, 0x6d, 0xdc, 0xbf, 0xec, 0x2, 0x1c, 0xd1, 0x60, 0x76, 0xc1, 0xe},
		},
		{
			sequenceNumber: 5003,
			encrypted:      []byte{0x24, 0x7e, 0x96, 0xc8, 0x7d, 0x33, 0xa2, 0x92, 0x8d, 0x13, 0x8d, 0xe0, 0x76, 0x9f, 0x8, 0xdc},
		},
		{
			sequenceNumber: 5004,
			encrypted:      []byte{0x75, 0x43, 0x28, 0xe4, 0x3a, 0x77, 0x59, 0x9b, 0x2e, 0xdf, 0x7b, 0x12, 0x68, 0xb, 0x57, 0x49},
		},
	}

	for _, testCase := range testCases {
		pkt := &rtp.Packet{Payload: append([]byte{}, decrypted...), SequenceNumber: testCase.sequenceNumber}
		if !encryptContext.EncryptPacket(pkt) {
			t.Errorf("Failed to encrypt RTP packet with SeqNum: %d", testCase.sequenceNumber)
		}
		assert.Equalf(pkt.Payload, testCase.encrypted, "RTP packet with SeqNum invalid encryption: %d", testCase.sequenceNumber)

		if !decryptContext.DecryptPacket(pkt) {
			t.Errorf("Failed to decrypt RTP packet with SeqNum: %d", testCase.sequenceNumber)
		}
		assert.Equalf(pkt.Payload, decrypted, "RTP packet with SeqNum invalid decryption: %d")

	}
}
<|endoftext|>


<commit_before>//  Inspiration from
//  https://gist.github.com/jpillora/b480fde82bff51a06238
//  Thanks to Jamie Pillora for the detailed gist.
//
// A small SSH daemon providing bash sessions
//
// Server:
// cd my/new/dir/
// #generate server keypair
// ssh-keygen -t rsa
// go get -v .
// go run sshd.go
//
// Client:
// ssh foo@localhost -p 2200 #pass=bar

package main

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"os/exec"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"github.com/kr/pty"
	pelican "github.com/mailgun/pelican-protocol"
	"golang.org/x/crypto/ssh"
)

type User struct {
	PubKey     string
	AcctId     string // 32-byte one-way crypto hash of PubKey. The 'user' who is ssh-ing in.
	Email      string
	SMS        string
	FirstName  string
	MiddleName string
	LastName   string
	FirstIp    string
	FirstTm    time.Time
	LastIp     string
	LastTm     time.Time
	Banned     bool
}

type Users struct {
	KnownClientPubKey map[string]User
}

func NewUsers() *Users {
	u := &Users{
		KnownClientPubKey: make(map[string]User),
	}
	return u
}

// should be constant-time to avoid side-channel timing attacks.
func (u *Users) PermitClientConnection(clientUser string, clientAddr net.Addr, clientPubKey ssh.PublicKey) (bool, error) {

	pubBytes := ssh.MarshalAuthorizedKey(clientPubKey)
	strPubBytes := string(chomp(pubBytes))

	fmt.Printf("\n in PermitClientConnection(): clientUser = '%s', clientAddr = '%#v', clientPubKey = '%s'\n", clientUser, clientAddr, strPubBytes)

	if strPubBytes == pelican.GetNewAcctPublicKey() {
		fmt.Printf("PermitClientConnection detected NewAcct public key, returning true\n")
		return true, fmt.Errorf("new-account")
	}
	fmt.Printf("equal? %v \n pelican.GetNewAcctPublicKey() = \n'%s'\n and strPubBytes = \n'%s'\n", pelican.GetNewAcctPublicKey() == strPubBytes, pelican.GetNewAcctPublicKey(), strPubBytes)

	if clientUser == "newacct" {
		fmt.Printf("PermitClientConnection detected user 'newacct', returning true\n")
		return true, fmt.Errorf("new-account")
	}
	// the username is issued by the server upon the completion of the
	// new account protocol, and is the hmac of the client's public key
	// signed with a secret only the server knows. The secret used
	// to sign the hmac should be preserved as long as the service is
	// alive, since if you loose it none of the account names can
	// be validated.

	/* not so strict at first!

	secretServerId, err := FetchSecretIdForService(".secret_id_for_service")
	panicOn(err)
	hmac := Sha1HMAC(pubBytes, []byte(secretServerId))
	acctid := encodeSha1HmacAsUsername(hmac)

	if clientUser != acctid {
		// somebody is trying to use a public key we've seen before, but
		// they did not receive our signature (acctid) for it, so likely
		// what happened is they didn't complete the new-account-creation
		// protocol. Hence we reject their connection.
		// The final step in completion of the new account protocol is that
		// we give the client their acctid (equivalent of a username);
		// which is an hmac of their public key, signed with a secret
		// known only to the server.
		return false, fmt.Errorf("bad-account-id")
	}
	*/

	user, ok := u.KnownClientPubKey[strPubBytes]
	if ok {
		if user.Banned {
			fmt.Printf("PermitClientConnection returning banned-user\n")
			return false, fmt.Errorf("banned-user")
		}
		now := time.Now()
		user.LastIp = clientAddr.String()
		user.LastTm = now
		if user.FirstIp == "" {
			user.FirstIp = user.LastIp
			user.FirstTm = now
		}
		fmt.Printf("PermitClientConnection true, user known.\n")
		return true, nil
	}

	fmt.Printf("PermitClientConnection false: user-unknown\n")
	return false, fmt.Errorf("user-unknown")
}

func main() {

	users := NewUsers()

	config := &ssh.ServerConfig{
		// must have keys
		NoClientAuth: false,

		// pki based login only
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			ok, err := users.PermitClientConnection(conn.User(), conn.RemoteAddr(), key)
			if !ok {
				return nil, err
			}
			perm := &ssh.Permissions{Extensions: map[string]string{
				"pubkey": string(key.Marshal()),
			}}
			return perm, nil
		},
		// no passwords
		PasswordCallback: nil,
	}

	err := GetOrGenServerKey("./host-key-id-rsa", config)
	panicOn(err)

	// Once a ServerConfig has been configured, connections can be accepted.
	listener, err := net.Listen("tcp", "0.0.0.0:2200")
	if err != nil {
		log.Fatalf("Failed to listen on 2200 (%s)", err)
	}

	// Accept all connections
	log.Print("Listening on 2200...")
	for {
		tcpConn, err := listener.Accept()
		if err != nil {
			log.Printf("Failed to accept incoming connection (%s)", err)
			continue
		}
		// Before use, a handshake must be performed on the incoming net.Conn.
		sshConn, chans, reqs, err := ssh.NewServerConn(tcpConn, config)
		if err != nil {
			log.Printf("Failed to handshake (%s)", err)
			continue
		}

		log.Printf("New SSH connection from %s (%s)", sshConn.RemoteAddr(), sshConn.ClientVersion())
		// Discard all global out-of-band Requests
		go ssh.DiscardRequests(reqs)
		// Accept all channels
		go handleChannels(chans)
	}
}

func handleChannels(chans <-chan ssh.NewChannel) {
	// Service the incoming Channel channel in go routine
	for newChannel := range chans {
		go handleChannel(newChannel)
	}
}

func handleChannel(newChannel ssh.NewChannel) {
	// Since we're handling a shell, we expect a
	// channel type of "session". The also describes
	// "x11", "direct-tcpip" and "forwarded-tcpip"
	// channel types.
	if t := newChannel.ChannelType(); t != "session" {
		newChannel.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
		return
	}

	// At this point, we have the opportunity to reject the client's
	// request for another logical connection
	connection, requests, err := newChannel.Accept()
	if err != nil {
		log.Printf("Could not accept channel (%s)", err)
		return
	}

	// Fire up bash for this session
	bash := exec.Command("bash")

	// Prepare teardown function
	close := func() {
		connection.Close()
		_, err := bash.Process.Wait()
		if err != nil {
			log.Printf("Failed to exit bash (%s)", err)
		}
		log.Printf("Session closed")
	}

	// Allocate a terminal for this channel
	log.Print("Creating pty...")
	bashf, err := pty.Start(bash)
	if err != nil {
		log.Printf("Could not start pty (%s)", err)
		close()
		return
	}

	//pipe session to bash and visa-versa
	var once sync.Once
	go func() {
		io.Copy(connection, bashf)
		once.Do(close)
	}()
	go func() {
		io.Copy(bashf, connection)
		once.Do(close)
	}()

	// Sessions have out-of-band requests such as "shell", "pty-req" and "env"
	go func() {
		for req := range requests {
			switch req.Type {
			case "shell":
				// We only accept the default shell
				// (i.e. no command in the Payload)
				if len(req.Payload) == 0 {
					req.Reply(true, nil)
				}
			case "pty-req":
				termLen := req.Payload[3]
				w, h := parseDims(req.Payload[termLen+4:])
				SetWinsize(bashf.Fd(), w, h)
				// Responding true (OK) here will let the client
				// know we have a pty ready for input
				req.Reply(true, nil)
			case "window-change":
				w, h := parseDims(req.Payload)
				SetWinsize(bashf.Fd(), w, h)
			}
		}
	}()
}

// =======================

// parseDims extracts terminal dimensions (width x height) from the provided buffer.
func parseDims(b []byte) (uint32, uint32) {
	w := binary.BigEndian.Uint32(b)
	h := binary.BigEndian.Uint32(b[4:])
	return w, h
}

// ======================

// Winsize stores the Height and Width of a terminal.
type Winsize struct {
	Height uint16
	Width  uint16
	x      uint16 // unused
	y      uint16 // unused
}

// SetWinsize sets the size of the given pty.
//  from https://github.com/creack/termios/blob/master/win/win.go
func SetWinsize(fd uintptr, w, h uint32) {
	ws := &Winsize{Width: uint16(w), Height: uint16(h)}
	syscall.Syscall(syscall.SYS_IOCTL, fd, uintptr(syscall.TIOCSWINSZ), uintptr(unsafe.Pointer(ws)))
}

func chomp(b []byte) []byte {
	for i := len(b) - 1; i >= 0; i-- {
		if b[i] != '\n' {
			return b[:i+1]
		}
	}
	return b[:0]
}
<commit_msg>trim sshd down<commit_after>package main

import (
	"fmt"
	"log"
	"net"
	"time"

	pelican "github.com/mailgun/pelican-protocol"
	"golang.org/x/crypto/ssh"
)

type User struct {
	PubKey     string
	AcctId     string // 32-byte one-way crypto hash of PubKey. The 'user' who is ssh-ing in.
	Email      string
	SMS        string
	FirstName  string
	MiddleName string
	LastName   string
	FirstIp    string
	FirstTm    time.Time
	LastIp     string
	LastTm     time.Time
	Banned     bool
	Completed  bool // signup complete marked here.
}

type Users struct {
	KnownClientPubKey map[string]User
}

func NewUsers() *Users {
	u := &Users{
		KnownClientPubKey: make(map[string]User),
	}
	return u
}

// should be constant-time to avoid side-channel timing attacks.
func (u *Users) PermitClientConnection(clientUser string, clientAddr net.Addr, clientPubKey ssh.PublicKey) (bool, error) {

	pubBytes := ssh.MarshalAuthorizedKey(clientPubKey)
	strPubBytes := string(chomp(pubBytes))

	fmt.Printf("\n in PermitClientConnection(): clientUser = '%s', clientAddr = '%#v', clientPubKey = '%s'\n", clientUser, clientAddr, strPubBytes)

	if strPubBytes == pelican.GetNewAcctPublicKey() {
		fmt.Printf("PermitClientConnection detected NewAcct public key, returning true\n")
		return true, fmt.Errorf("new-account")
	}
	fmt.Printf("equal? %v \n pelican.GetNewAcctPublicKey() = \n'%s'\n and strPubBytes = \n'%s'\n", pelican.GetNewAcctPublicKey() == strPubBytes, pelican.GetNewAcctPublicKey(), strPubBytes)

	if clientUser == "newacct" {
		fmt.Printf("PermitClientConnection detected user 'newacct', returning true\n")
		return true, fmt.Errorf("new-account")
	}
	// the username is issued by the server upon the completion of the
	// new account protocol, and is the hmac of the client's public key
	// signed with a secret only the server knows. The secret used
	// to sign the hmac should be preserved as long as the service is
	// alive, since if you loose it none of the account names can
	// be validated.

	secretServerId, err := FetchSecretIdForService(".secret_id_for_service")
	panicOn(err)
	hmac := Sha1HMAC(pubBytes, []byte(secretServerId))
	acctid := encodeSha1HmacAsUsername(hmac)

	if clientUser != acctid {
		// somebody is trying to use a public key we've seen before, but
		// they did not receive our signature (acctid) for it, so likely
		// what happened is they didn't complete the new-account-creation
		// protocol. Hence we reject their connection.
		// The final step in completion of the new account protocol is that
		// we give the client their acctid (equivalent of a username);
		// which is an hmac of their public key, signed with a secret
		// known only to the server.
		return false, fmt.Errorf("bad-account-id")
	}

	user, ok := u.KnownClientPubKey[strPubBytes]
	if ok {
		if user.Banned {
			fmt.Printf("PermitClientConnection returning banned-user\n")
			return false, fmt.Errorf("banned-user")
		}
		now := time.Now()
		user.LastIp = clientAddr.String()
		user.LastTm = now
		if user.FirstIp == "" {
			user.FirstIp = user.LastIp
			user.FirstTm = now
		}
		fmt.Printf("PermitClientConnection true, user known.\n")
		return true, nil
	}

	fmt.Printf("PermitClientConnection false: user-unknown\n")
	return false, fmt.Errorf("user-unknown")
}

func main() {

	users := NewUsers()

	config := &ssh.ServerConfig{
		// must have keys
		NoClientAuth: false,

		// pki based login only
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			ok, err := users.PermitClientConnection(conn.User(), conn.RemoteAddr(), key)
			if !ok {
				return nil, err
			}
			perm := &ssh.Permissions{Extensions: map[string]string{
				"pubkey": string(key.Marshal()),
			}}
			return perm, nil
		},
		// no passwords
		PasswordCallback: nil,
	}

	err := GetOrGenServerKey("./host-key-id-rsa", config)
	panicOn(err)

	// Once a ServerConfig has been configured, connections can be accepted.
	listener, err := net.Listen("tcp", "0.0.0.0:2200")
	if err != nil {
		log.Fatalf("Failed to listen on 2200 (%s)", err)
	}

	// Accept all connections
	log.Print("Listening on 2200...")
	for {
		tcpConn, err := listener.Accept()
		if err != nil {
			log.Printf("Failed to accept incoming connection (%s)", err)
			continue
		}
		// Before use, a handshake must be performed on the incoming net.Conn.
		sshConn, chans, reqs, err := ssh.NewServerConn(tcpConn, config)
		if err != nil {
			log.Printf("Failed to handshake (%s)", err)
			continue
		}

		log.Printf("New SSH connection from %s (%s)", sshConn.RemoteAddr(), sshConn.ClientVersion())
		// Discard all global out-of-band Requests
		go ssh.DiscardRequests(reqs)
		// Accept all channels
		go handleChannels(chans)
	}
}

func handleChannels(chans <-chan ssh.NewChannel) {
	// Service the incoming Channel channel in go routine
	for newChannel := range chans {
		go handleChannel(newChannel)
	}
}

// this was for bash channels, which aren't desired for pelican-protocol
/*
func handleChannel(newChannel ssh.NewChannel) {

	fmt.Printf("\n in handleChannle() with channel type = '%s'\n", newChannel.ChannelType())

	// Since we're handling a shell, we expect a
	// channel type of "session". The also describes
	// "x11", "direct-tcpip" and "forwarded-tcpip"
	// channel types.
	if t := newChannel.ChannelType(); t != "session" {
		newChannel.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
		return
	}

	// At this point, we have the opportunity to reject the client's
	// request for another logical connection
	connection, requests, err := newChannel.Accept()
	if err != nil {
		log.Printf("Could not accept channel (%s)", err)
		return
	}

	// Fire up bash for this session
	bash := exec.Command("bash")

	// Prepare teardown function
	close := func() {
		connection.Close()
		_, err := bash.Process.Wait()
		if err != nil {
			log.Printf("Failed to exit bash (%s)", err)
		}
		log.Printf("Session closed")
	}

	// Allocate a terminal for this channel
	log.Print("Creating pty...")
	bashf, err := pty.Start(bash)
	if err != nil {
		log.Printf("Could not start pty (%s)", err)
		close()
		return
	}

	//pipe session to bash and visa-versa
	var once sync.Once
	go func() {
		io.Copy(connection, bashf)
		once.Do(close)
	}()
	go func() {
		io.Copy(bashf, connection)
		once.Do(close)
	}()

	// Sessions have out-of-band requests such as "shell", "pty-req" and "env"
	go func() {
		for req := range requests {
			switch req.Type {
			case "shell":
				// We only accept the default shell
				// (i.e. no command in the Payload)
				if len(req.Payload) == 0 {
					req.Reply(true, nil)
				}
			case "pty-req":
				termLen := req.Payload[3]
				w, h := parseDims(req.Payload[termLen+4:])
				SetWinsize(bashf.Fd(), w, h)
				// Responding true (OK) here will let the client
				// know we have a pty ready for input
				req.Reply(true, nil)
			case "window-change":
				w, h := parseDims(req.Payload)
				SetWinsize(bashf.Fd(), w, h)
			}
		}
	}()
}
*/
// =======================

func handleChannel(newChannel ssh.NewChannel) {

	fmt.Printf("\n in handleChannle() with channel type = '%s'\n", newChannel.ChannelType())

	// Since we're handling a shell, we expect a
	// channel type of "session". The also describes
	// "x11", "direct-tcpip" and "forwarded-tcpip"
	// channel types.
	if t := newChannel.ChannelType(); t != "forwarded-tcpip" && t != "direct-tcpip" {
		newChannel.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
		return
	}

	// At this point, we have the opportunity to reject the client's
	// request for another logical connection
	connection, requests, err := newChannel.Accept()
	if err != nil {
		log.Printf("Could not accept channel (%s)", err)
		return
	}

	fmt.Printf("connection = '%#v'  requests = '%#v'\n", connection, requests)

	// Sessions have out-of-band requests such as "shell", "pty-req" and "env"
	go func() {
		for req := range requests {
			fmt.Printf("\n ignoring req.Type = '%v'\n", req.Type)
		}
	}()
}

func chomp(b []byte) []byte {
	for i := len(b) - 1; i >= 0; i-- {
		if b[i] != '\n' {
			return b[:i+1]
		}
	}
	return b[:0]
}
<|endoftext|>


<commit_before>package phraseapp

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"os"

	"github.com/bgentry/speakeasy"
)

var Debug bool

func EnableDebug() {
	Debug = true
}

type DefaultParams Action
type Action map[string]map[string]interface{}

type Client struct {
	http.Client
	Credentials *Credentials
}

type Credentials struct {
	Username string `cli:"opt --username -u desc='username used for authentication'"`
	Token    string `cli:"opt --access-token -t desc='access token used for authentication'"`
	TFA      bool   `cli:"opt --tfa desc='use Two-Factor Authentication'"`
	Host     string `cli:"opt --host desc='Host to send Request to'"`
	Debug    bool   `cli:"opt --verbose -v desc='Verbose output'"`
}

func NewClient(credentials *Credentials) (*Client, error) {
	client := &Client{Credentials: &Credentials{}}

	envToken := os.Getenv("PHRASEAPP_ACCESS_TOKEN")

	if credentials.Token != "" && client.Credentials.Token == "" && client.Credentials.Username == "" {
		client.Credentials.Token = credentials.Token
	} else if credentials.Username != "" && client.Credentials.Username == "" {
		client.Credentials.Username = credentials.Username
	} else if envToken != "" && credentials.Token == "" && credentials.Username == "" && client.Credentials.Username == "" {
		client.Credentials.Token = envToken
	}

	if credentials.TFA {
		client.Credentials.TFA = credentials.TFA
	}

	if credentials.Debug == true {
		EnableDebug()
	}

	if credentials.Host != "" {
		client.Credentials.Host = credentials.Host
	} else {
		client.Credentials.Host = "https://api.phraseapp.com"
	}

	return client, nil
}

func (client *Client) authenticate(req *http.Request) error {
	if client.Credentials == nil {
		return fmt.Errorf("no auth handler registered")
	}

	if err := client.Credentials.validate(); err != nil {
		return err
	}

	req.Header.Set("User-Agent", GetUserAgent())
	switch {
	case client.Credentials.Token != "":
		req.Header.Set("Authorization", "token "+client.Credentials.Token)
	case client.Credentials.Username != "":
		pwd, err := speakeasy.Ask("Password: ")
		if err != nil {
			return err
		}
		req.SetBasicAuth(client.Credentials.Username, pwd)

		if client.Credentials.TFA { // TFA only required for username+password based login.
			token, err := speakeasy.Ask("TFA-Token: ")
			if err != nil {
				return err
			}
			req.Header.Set("X-PhraseApp-OTP", token)
		}
	}

	return nil
}

func (ah *Credentials) validate() error {
	switch {
	case ah.Username == "" && ah.Token == "":
		return fmt.Errorf("either username or token must be given")
	default:
		return nil
	}
}

func (client *Client) sendRequestPaginated(method, rawurl, ctype string, r io.Reader, status, page, perPage int) (io.ReadCloser, error) {
	endpointUrl := client.Credentials.Host + rawurl
	u, err := url.Parse(endpointUrl)
	if err != nil {
		return nil, err
	}

	query := u.Query()
	query.Add("page", strconv.Itoa(page))
	query.Add("per_page", strconv.Itoa(perPage))

	u.RawQuery = query.Encode()

	if Debug {
		fmt.Fprintln(os.Stderr, method, u)
		if r != nil {
			bytes, err := ioutil.ReadAll(r)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s", err.Error())
			}
			str := string(bytes)
			fmt.Fprintln(os.Stderr, str)
			r = strings.NewReader(str)
		}
	}

	req, err := http.NewRequest(method, u.String(), r)
	if err != nil {
		return nil, err
	}

	if ctype != "" {
		req.Header.Add("Content-Type", ctype)
	}

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

	return resp.Body, nil
}

func (client *Client) sendRequest(method, url, ctype string, r io.Reader, status int) (io.ReadCloser, error) {
	endpointUrl := client.Credentials.Host + url
	if Debug {
		fmt.Fprintln(os.Stderr, method, url)
		if r != nil {
			bytes, err := ioutil.ReadAll(r)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s", err.Error())
			}
			str := string(bytes)
			fmt.Fprintln(os.Stderr, str)
			r = strings.NewReader(str)
		}
	}
	req, err := http.NewRequest(method, endpointUrl, r)
	if err != nil {
		return nil, err
	}

	if ctype != "" {
		req.Header.Add("Content-Type", ctype)
	}

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

func (client *Client) send(req *http.Request, status int) (*http.Response, error) {
	err := client.authenticate(req)
	if err != nil {
		return nil, err
	}

	if Debug {
		b := new(bytes.Buffer)
		err = req.Header.Write(b)
		if err != nil {
			return nil, err
		}

		bytes, err := ioutil.ReadAll(b)
		if err != nil {
			return nil, err
		}
		fmt.Fprintln(os.Stderr, string(bytes))

		r := req.Body
		if r != nil {
			var by []byte
			_, err = r.Read(by)
			fmt.Fprintln(os.Stderr, string(by))
		}
	}
	resp, err := client.Client.Do(req)
	if err != nil {
		return nil, err
	}

	if Debug {
		fmt.Fprintf(os.Stderr, "\nResponse HTTP Status Code: %s\n", resp.Status)
	}

	err = handleResponseStatus(resp, status)
	if err != nil {
		resp.Body.Close()
	}
	return resp, err
}
<commit_msg>client: remove legacy types<commit_after>package phraseapp

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"os"

	"github.com/bgentry/speakeasy"
)

var Debug bool

func EnableDebug() {
	Debug = true
}

type Client struct {
	http.Client
	Credentials *Credentials
}

type Credentials struct {
	Username string `cli:"opt --username -u desc='username used for authentication'"`
	Token    string `cli:"opt --access-token -t desc='access token used for authentication'"`
	TFA      bool   `cli:"opt --tfa desc='use Two-Factor Authentication'"`
	Host     string `cli:"opt --host desc='Host to send Request to'"`
	Debug    bool   `cli:"opt --verbose -v desc='Verbose output'"`
}

func NewClient(credentials *Credentials) (*Client, error) {
	client := &Client{Credentials: &Credentials{}}

	envToken := os.Getenv("PHRASEAPP_ACCESS_TOKEN")

	if credentials.Token != "" && client.Credentials.Token == "" && client.Credentials.Username == "" {
		client.Credentials.Token = credentials.Token
	} else if credentials.Username != "" && client.Credentials.Username == "" {
		client.Credentials.Username = credentials.Username
	} else if envToken != "" && credentials.Token == "" && credentials.Username == "" && client.Credentials.Username == "" {
		client.Credentials.Token = envToken
	}

	if credentials.TFA {
		client.Credentials.TFA = credentials.TFA
	}

	if credentials.Debug == true {
		EnableDebug()
	}

	if credentials.Host != "" {
		client.Credentials.Host = credentials.Host
	} else {
		client.Credentials.Host = "https://api.phraseapp.com"
	}

	return client, nil
}

func (client *Client) authenticate(req *http.Request) error {
	if client.Credentials == nil {
		return fmt.Errorf("no auth handler registered")
	}

	if err := client.Credentials.validate(); err != nil {
		return err
	}

	req.Header.Set("User-Agent", GetUserAgent())
	switch {
	case client.Credentials.Token != "":
		req.Header.Set("Authorization", "token "+client.Credentials.Token)
	case client.Credentials.Username != "":
		pwd, err := speakeasy.Ask("Password: ")
		if err != nil {
			return err
		}
		req.SetBasicAuth(client.Credentials.Username, pwd)

		if client.Credentials.TFA { // TFA only required for username+password based login.
			token, err := speakeasy.Ask("TFA-Token: ")
			if err != nil {
				return err
			}
			req.Header.Set("X-PhraseApp-OTP", token)
		}
	}

	return nil
}

func (ah *Credentials) validate() error {
	switch {
	case ah.Username == "" && ah.Token == "":
		return fmt.Errorf("either username or token must be given")
	default:
		return nil
	}
}

func (client *Client) sendRequestPaginated(method, rawurl, ctype string, r io.Reader, status, page, perPage int) (io.ReadCloser, error) {
	endpointUrl := client.Credentials.Host + rawurl
	u, err := url.Parse(endpointUrl)
	if err != nil {
		return nil, err
	}

	query := u.Query()
	query.Add("page", strconv.Itoa(page))
	query.Add("per_page", strconv.Itoa(perPage))

	u.RawQuery = query.Encode()

	if Debug {
		fmt.Fprintln(os.Stderr, method, u)
		if r != nil {
			bytes, err := ioutil.ReadAll(r)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s", err.Error())
			}
			str := string(bytes)
			fmt.Fprintln(os.Stderr, str)
			r = strings.NewReader(str)
		}
	}

	req, err := http.NewRequest(method, u.String(), r)
	if err != nil {
		return nil, err
	}

	if ctype != "" {
		req.Header.Add("Content-Type", ctype)
	}

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

	return resp.Body, nil
}

func (client *Client) sendRequest(method, url, ctype string, r io.Reader, status int) (io.ReadCloser, error) {
	endpointUrl := client.Credentials.Host + url
	if Debug {
		fmt.Fprintln(os.Stderr, method, url)
		if r != nil {
			bytes, err := ioutil.ReadAll(r)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s", err.Error())
			}
			str := string(bytes)
			fmt.Fprintln(os.Stderr, str)
			r = strings.NewReader(str)
		}
	}
	req, err := http.NewRequest(method, endpointUrl, r)
	if err != nil {
		return nil, err
	}

	if ctype != "" {
		req.Header.Add("Content-Type", ctype)
	}

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

func (client *Client) send(req *http.Request, status int) (*http.Response, error) {
	err := client.authenticate(req)
	if err != nil {
		return nil, err
	}

	if Debug {
		b := new(bytes.Buffer)
		err = req.Header.Write(b)
		if err != nil {
			return nil, err
		}

		bytes, err := ioutil.ReadAll(b)
		if err != nil {
			return nil, err
		}
		fmt.Fprintln(os.Stderr, string(bytes))

		r := req.Body
		if r != nil {
			var by []byte
			_, err = r.Read(by)
			fmt.Fprintln(os.Stderr, string(by))
		}
	}
	resp, err := client.Client.Do(req)
	if err != nil {
		return nil, err
	}

	if Debug {
		fmt.Fprintf(os.Stderr, "\nResponse HTTP Status Code: %s\n", resp.Status)
	}

	err = handleResponseStatus(resp, status)
	if err != nil {
		resp.Body.Close()
	}
	return resp, err
}
<|endoftext|>


<commit_before>package phraseapp

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"os"

	"github.com/bgentry/speakeasy"
)

var Debug bool

func EnableDebug() {
	Debug = true
}

type DefaultParams Action
type Action map[string]map[string]interface{}

type Client struct {
	http.Client
	Credentials *Credentials
}

type Credentials struct {
	Username string
	Token    string
	TFA      bool
	Host     string
	Debug    bool
}

func NewClient(credentials Credentials, defaultCredentials *Credentials) (*Client, error) {
	client := &Client{Credentials: &Credentials{}}

	envToken := os.Getenv("PHRASEAPP_ACCESS_TOKEN")

	if credentials.Token != "" && client.Credentials.Token == "" && client.Credentials.Username == "" {
		client.Credentials.Token = credentials.Token
	} else if credentials.Username != "" && client.Credentials.Username == "" {
		client.Credentials.Username = credentials.Username
	} else if envToken != "" && credentials.Token == "" && credentials.Username == "" && client.Credentials.Username == "" {
		client.Credentials.Token = envToken
	}

	if credentials.TFA {
		client.Credentials.TFA = credentials.TFA
	}

	if credentials.Debug == true || ((defaultCredentials != nil) && defaultCredentials.Debug == true) {
		EnableDebug()
	}

	if credentials.Host != "" {
		client.Credentials.Host = credentials.Host
	} else {
		if defaultCredentials != nil && defaultCredentials.Host != "" {
			client.Credentials.Host = defaultCredentials.Host
		}
	}

	if client.Credentials.Host == "" {
		client.Credentials.Host = "https://api.phraseapp.com"
	}

	notSet := client.Credentials.Token == "" && client.Credentials.Username == ""
	if notSet && defaultCredentials != nil && defaultCredentials.Token != "" {
		client.Credentials.Token = defaultCredentials.Token
	}
	if notSet && defaultCredentials != nil && defaultCredentials.Username != "" {
		client.Credentials.Username = defaultCredentials.Username
	}

	return client, nil
}

func (client *Client) authenticate(req *http.Request) error {
	if client.Credentials == nil {
		return fmt.Errorf("no auth handler registered")
	}

	if err := client.Credentials.validate(); err != nil {
		return err
	}

	req.Header.Set("User-Agent", GetUserAgent())
	switch {
	case client.Credentials.Token != "":
		req.Header.Set("Authorization", "token "+client.Credentials.Token)
	case client.Credentials.Username != "":
		pwd, err := speakeasy.Ask("Password: ")
		if err != nil {
			return err
		}
		req.SetBasicAuth(client.Credentials.Username, pwd)

		if client.Credentials.TFA { // TFA only required for username+password based login.
			token, err := speakeasy.Ask("TFA-Token: ")
			if err != nil {
				return err
			}
			req.Header.Set("X-PhraseApp-OTP", token)
		}
	}

	return nil
}

func (ah *Credentials) validate() error {
	switch {
	case ah.Username == "" && ah.Token == "":
		return fmt.Errorf("either username or token must be given")
	default:
		return nil
	}
}

func (client *Client) sendRequestPaginated(method, rawurl, ctype string, r io.Reader, status, page, perPage int) (io.ReadCloser, error) {
	endpointUrl := client.Credentials.Host + rawurl
	u, err := url.Parse(endpointUrl)
	if err != nil {
		return nil, err
	}

	query := u.Query()
	query.Add("page", strconv.Itoa(page))
	query.Add("per_page", strconv.Itoa(perPage))

	u.RawQuery = query.Encode()

	if Debug {
		fmt.Fprintln(os.Stderr, method, u)
		if r != nil {
			bytes, err := ioutil.ReadAll(r)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s", err.Error())
			}
			str := string(bytes)
			fmt.Fprintln(os.Stderr, str)
			r = strings.NewReader(str)
		}
	}

	req, err := http.NewRequest(method, u.String(), r)
	if err != nil {
		return nil, err
	}

	if ctype != "" {
		req.Header.Add("Content-Type", ctype)
	}

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

	return resp.Body, nil
}

func (client *Client) sendRequest(method, url, ctype string, r io.Reader, status int) (io.ReadCloser, error) {
	endpointUrl := client.Credentials.Host + url
	if Debug {
		fmt.Fprintln(os.Stderr, method, url)
		if r != nil {
			bytes, err := ioutil.ReadAll(r)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s", err.Error())
			}
			str := string(bytes)
			fmt.Fprintln(os.Stderr, str)
			r = strings.NewReader(str)
		}
	}
	req, err := http.NewRequest(method, endpointUrl, r)
	if err != nil {
		return nil, err
	}

	if ctype != "" {
		req.Header.Add("Content-Type", ctype)
	}

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

func (client *Client) send(req *http.Request, status int) (*http.Response, error) {
	err := client.authenticate(req)
	if err != nil {
		return nil, err
	}

	if Debug {
		b := new(bytes.Buffer)
		err = req.Header.Write(b)
		if err != nil {
			return nil, err
		}

		bytes, err := ioutil.ReadAll(b)
		if err != nil {
			return nil, err
		}
		fmt.Fprintln(os.Stderr, string(bytes))

		r := req.Body
		if r != nil {
			var by []byte
			_, err = r.Read(by)
			fmt.Fprintln(os.Stderr, string(by))
		}
	}
	resp, err := client.Client.Do(req)
	if err != nil {
		return nil, err
	}

	if Debug {
		fmt.Fprintf(os.Stderr, "\nResponse HTTP Status Code: %s\n", resp.Status)
	}

	err = handleResponseStatus(resp, status)
	if err != nil {
		resp.Body.Close()
	}
	return resp, err
}
<commit_msg>credentials: make usable from CLI<commit_after>package phraseapp

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"os"

	"github.com/bgentry/speakeasy"
)

var Debug bool

func EnableDebug() {
	Debug = true
}

type DefaultParams Action
type Action map[string]map[string]interface{}

type Client struct {
	http.Client
	Credentials *Credentials
}

type Credentials struct {
	Username string `cli:"opt --username -u desc='username used for authentication'"`
	Token    string `cli:"opt --access-token -t desc='access token used for authentication'"`
	TFA      bool   `cli:"opt --tfa desc='use Two-Factor Authentication'"`
	Host     string `cli:"opt --host desc='Host to send Request to'"`
	Debug    bool   `cli:"opt --verbose -v desc='Verbose output'"`
}

func NewClient(credentials Credentials, defaultCredentials *Credentials) (*Client, error) {
	client := &Client{Credentials: &Credentials{}}

	envToken := os.Getenv("PHRASEAPP_ACCESS_TOKEN")

	if credentials.Token != "" && client.Credentials.Token == "" && client.Credentials.Username == "" {
		client.Credentials.Token = credentials.Token
	} else if credentials.Username != "" && client.Credentials.Username == "" {
		client.Credentials.Username = credentials.Username
	} else if envToken != "" && credentials.Token == "" && credentials.Username == "" && client.Credentials.Username == "" {
		client.Credentials.Token = envToken
	}

	if credentials.TFA {
		client.Credentials.TFA = credentials.TFA
	}

	if credentials.Debug == true || ((defaultCredentials != nil) && defaultCredentials.Debug == true) {
		EnableDebug()
	}

	if credentials.Host != "" {
		client.Credentials.Host = credentials.Host
	} else {
		if defaultCredentials != nil && defaultCredentials.Host != "" {
			client.Credentials.Host = defaultCredentials.Host
		}
	}

	if client.Credentials.Host == "" {
		client.Credentials.Host = "https://api.phraseapp.com"
	}

	notSet := client.Credentials.Token == "" && client.Credentials.Username == ""
	if notSet && defaultCredentials != nil && defaultCredentials.Token != "" {
		client.Credentials.Token = defaultCredentials.Token
	}
	if notSet && defaultCredentials != nil && defaultCredentials.Username != "" {
		client.Credentials.Username = defaultCredentials.Username
	}

	return client, nil
}

func (client *Client) authenticate(req *http.Request) error {
	if client.Credentials == nil {
		return fmt.Errorf("no auth handler registered")
	}

	if err := client.Credentials.validate(); err != nil {
		return err
	}

	req.Header.Set("User-Agent", GetUserAgent())
	switch {
	case client.Credentials.Token != "":
		req.Header.Set("Authorization", "token "+client.Credentials.Token)
	case client.Credentials.Username != "":
		pwd, err := speakeasy.Ask("Password: ")
		if err != nil {
			return err
		}
		req.SetBasicAuth(client.Credentials.Username, pwd)

		if client.Credentials.TFA { // TFA only required for username+password based login.
			token, err := speakeasy.Ask("TFA-Token: ")
			if err != nil {
				return err
			}
			req.Header.Set("X-PhraseApp-OTP", token)
		}
	}

	return nil
}

func (ah *Credentials) validate() error {
	switch {
	case ah.Username == "" && ah.Token == "":
		return fmt.Errorf("either username or token must be given")
	default:
		return nil
	}
}

func (client *Client) sendRequestPaginated(method, rawurl, ctype string, r io.Reader, status, page, perPage int) (io.ReadCloser, error) {
	endpointUrl := client.Credentials.Host + rawurl
	u, err := url.Parse(endpointUrl)
	if err != nil {
		return nil, err
	}

	query := u.Query()
	query.Add("page", strconv.Itoa(page))
	query.Add("per_page", strconv.Itoa(perPage))

	u.RawQuery = query.Encode()

	if Debug {
		fmt.Fprintln(os.Stderr, method, u)
		if r != nil {
			bytes, err := ioutil.ReadAll(r)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s", err.Error())
			}
			str := string(bytes)
			fmt.Fprintln(os.Stderr, str)
			r = strings.NewReader(str)
		}
	}

	req, err := http.NewRequest(method, u.String(), r)
	if err != nil {
		return nil, err
	}

	if ctype != "" {
		req.Header.Add("Content-Type", ctype)
	}

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

	return resp.Body, nil
}

func (client *Client) sendRequest(method, url, ctype string, r io.Reader, status int) (io.ReadCloser, error) {
	endpointUrl := client.Credentials.Host + url
	if Debug {
		fmt.Fprintln(os.Stderr, method, url)
		if r != nil {
			bytes, err := ioutil.ReadAll(r)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s", err.Error())
			}
			str := string(bytes)
			fmt.Fprintln(os.Stderr, str)
			r = strings.NewReader(str)
		}
	}
	req, err := http.NewRequest(method, endpointUrl, r)
	if err != nil {
		return nil, err
	}

	if ctype != "" {
		req.Header.Add("Content-Type", ctype)
	}

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

func (client *Client) send(req *http.Request, status int) (*http.Response, error) {
	err := client.authenticate(req)
	if err != nil {
		return nil, err
	}

	if Debug {
		b := new(bytes.Buffer)
		err = req.Header.Write(b)
		if err != nil {
			return nil, err
		}

		bytes, err := ioutil.ReadAll(b)
		if err != nil {
			return nil, err
		}
		fmt.Fprintln(os.Stderr, string(bytes))

		r := req.Body
		if r != nil {
			var by []byte
			_, err = r.Read(by)
			fmt.Fprintln(os.Stderr, string(by))
		}
	}
	resp, err := client.Client.Do(req)
	if err != nil {
		return nil, err
	}

	if Debug {
		fmt.Fprintf(os.Stderr, "\nResponse HTTP Status Code: %s\n", resp.Status)
	}

	err = handleResponseStatus(resp, status)
	if err != nil {
		resp.Body.Close()
	}
	return resp, err
}
<|endoftext|>


<commit_before>// Copyright 2015 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

package stats

import "runtime"

// BasicMemStats includes a few of the fields from runtime.MemStats suitable for
// general logging.
type BasicMemStats struct {
	// General statistics.
	Alloc      uint64 // bytes allocated and still in use
	TotalAlloc uint64 // bytes allocated (even if freed)
	Sys        uint64 // bytes obtained from system (sum of XxxSys in runtime)
	Lookups    uint64 // number of pointer lookups
	Mallocs    uint64 // number of mallocs
	Frees      uint64 // number of frees

	// Main allocation heap statistics.
	HeapAlloc    uint64 // bytes allocated and still in use
	HeapSys      uint64 // bytes obtained from system
	HeapIdle     uint64 // bytes in idle spans
	HeapInuse    uint64 // bytes in non-idle span
	HeapReleased uint64 // bytes released to the OS
	HeapObjects  uint64 // total number of allocated objects

	// Garbage collector statistics.
	PauseTotalNs uint64
	PauseNs      uint64
}

type memStatsPlaceholder interface{}

// MemStatsWrapper wraps runtime.MemStats with an optionally less verbose JSON
// representation. The JSON field names correspond exactly to the runtime field
// names to avoid reimplementing the entire struct.
type MemStatsWrapper struct {
	memStatsPlaceholder `json:"Memory"`

	basic *BasicMemStats
	cache *runtime.MemStats
}

func NewMemStatsWrapper(verbose bool) *MemStatsWrapper {
	stats := &MemStatsWrapper{cache: &runtime.MemStats{}}

	if verbose {
		stats.memStatsPlaceholder = stats.cache
	} else {
		stats.basic = &BasicMemStats{}
		stats.memStatsPlaceholder = stats.basic
	}
	return stats
}

// Update fetches the current memstats from runtime and resets the cache.
func (s *MemStatsWrapper) Update() {
	runtime.ReadMemStats(s.cache)

	if s.basic != nil {
		// Gross, but any decent editor can generate this in a couple commands.
		s.basic.Alloc = s.cache.Alloc
		s.basic.TotalAlloc = s.cache.TotalAlloc
		s.basic.Sys = s.cache.Sys
		s.basic.Lookups = s.cache.Lookups
		s.basic.Mallocs = s.cache.Mallocs
		s.basic.Frees = s.cache.Frees
		s.basic.HeapAlloc = s.cache.HeapAlloc
		s.basic.HeapSys = s.cache.HeapSys
		s.basic.HeapIdle = s.cache.HeapIdle
		s.basic.HeapInuse = s.cache.HeapInuse
		s.basic.HeapReleased = s.cache.HeapReleased
		s.basic.HeapObjects = s.cache.HeapObjects
		s.basic.PauseTotalNs = s.cache.PauseTotalNs
		s.basic.PauseNs = s.cache.PauseNs[(s.cache.NumGC+255)%256]
	}
}
<commit_msg>stats: rename PauseNS to LatestPauseNS<commit_after>// Copyright 2015 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

package stats

import "runtime"

// BasicMemStats includes a few of the fields from runtime.MemStats suitable for
// general logging.
type BasicMemStats struct {
	// General statistics.
	Alloc      uint64 // bytes allocated and still in use
	TotalAlloc uint64 // bytes allocated (even if freed)
	Sys        uint64 // bytes obtained from system (sum of XxxSys in runtime)
	Lookups    uint64 // number of pointer lookups
	Mallocs    uint64 // number of mallocs
	Frees      uint64 // number of frees

	// Main allocation heap statistics.
	HeapAlloc    uint64 // bytes allocated and still in use
	HeapSys      uint64 // bytes obtained from system
	HeapIdle     uint64 // bytes in idle spans
	HeapInuse    uint64 // bytes in non-idle span
	HeapReleased uint64 // bytes released to the OS
	HeapObjects  uint64 // total number of allocated objects

	// Garbage collector statistics.
	PauseTotalNs  uint64
	LatestPauseNs uint64
}

type memStatsPlaceholder interface{}

// MemStatsWrapper wraps runtime.MemStats with an optionally less verbose JSON
// representation. The JSON field names correspond exactly to the runtime field
// names to avoid reimplementing the entire struct.
type MemStatsWrapper struct {
	memStatsPlaceholder `json:"Memory"`

	basic *BasicMemStats
	cache *runtime.MemStats
}

func NewMemStatsWrapper(verbose bool) *MemStatsWrapper {
	stats := &MemStatsWrapper{cache: &runtime.MemStats{}}

	if verbose {
		stats.memStatsPlaceholder = stats.cache
	} else {
		stats.basic = &BasicMemStats{}
		stats.memStatsPlaceholder = stats.basic
	}
	return stats
}

// Update fetches the current memstats from runtime and resets the cache.
func (s *MemStatsWrapper) Update() {
	runtime.ReadMemStats(s.cache)

	if s.basic != nil {
		// Gross, but any decent editor can generate this in a couple commands.
		s.basic.Alloc = s.cache.Alloc
		s.basic.TotalAlloc = s.cache.TotalAlloc
		s.basic.Sys = s.cache.Sys
		s.basic.Lookups = s.cache.Lookups
		s.basic.Mallocs = s.cache.Mallocs
		s.basic.Frees = s.cache.Frees
		s.basic.HeapAlloc = s.cache.HeapAlloc
		s.basic.HeapSys = s.cache.HeapSys
		s.basic.HeapIdle = s.cache.HeapIdle
		s.basic.HeapInuse = s.cache.HeapInuse
		s.basic.HeapReleased = s.cache.HeapReleased
		s.basic.HeapObjects = s.cache.HeapObjects
		s.basic.PauseTotalNs = s.cache.PauseTotalNs
		s.basic.LatestPauseNs = s.cache.PauseNs[(s.cache.NumGC+255)%256]
	}
}
<|endoftext|>