<commit_before>// Ogo

package ogo

import (
	//"fmt"
	"github.com/zenazn/goji"
	"github.com/zenazn/goji/web"
	"net/http"
	"strings"
)

type Handler func(c *RESTContext)

type Route struct {
	Pattern string
	Method  string
	Handler Handler
}

type Controller struct {
	Endpoint string
	Routes   map[string]*Route
	ReqCount int //访问计数
}

type ControllerInterface interface {
	//Init(endpoint string, c ControllerInterface)
	Get(c *RESTContext)
	Post(c *RESTContext)
	Put(c *RESTContext)
	Delete(c *RESTContext)
	Patch(c *RESTContext)
	Head(c *RESTContext)
}

func NewRoute(p string, m string, h Handler) *Route {
	return &Route{
		Pattern: p,
		Method:  m,
		Handler: h,
	}
}

// 封装
func handlerWrap(f Handler) web.HandlerFunc { //这里封装了webC到本地的结构中
	return func(c web.C, w http.ResponseWriter, r *http.Request) {
		f(getContext(c, w, r))
	}
}

func (ctr *Controller) Init(endpoint string, c ControllerInterface) {
	ctr.Endpoint = endpoint
	//ctr.Routes = make(map[string]*Route)
	//默认路由
	ctr.DefaultRoutes(c)
	if len(ctr.Routes) > 0 { //如果有自定义路由, 则优先增加
		for _, rt := range ctr.Routes {
			switch strings.ToLower(rt.Method) {
			case "get":
				ctr.RouteGet(rt)
			case "post":
				ctr.RoutePost(rt)
			case "put":
				ctr.RoutePut(rt)
			case "delete":
				ctr.RouteDelete(rt)
			case "patch":
				ctr.RoutePatch(rt)
			case "head":
				ctr.RouteHead(rt)
			default:
				// unknow method
			}
		}
	}
}

func (ctr *Controller) Get(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Post(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Put(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Delete(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Patch(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Head(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}

func (ctr *Controller) AddRoute(m string, p string, h Handler) {
	key := strings.ToUpper(m) + " " + p
	if _, ok := ctr.Routes[key]; ok {
		//手动加路由, 以最后加的为准,overwrite
	}
	ctr.Routes[key] = NewRoute(p, m, h)
	//switch strings.ToLower(m) {
	//case "get":
	//	ctr.RouteGet(rt)
	//case "post":
	//	ctr.RoutePost(rt)
	//case "put":
	//	ctr.RoutePut(rt)
	//case "delete":
	//	ctr.RouteDelete(rt)
	//case "patch":
	//	ctr.RoutePatch(rt)
	//case "head":
	//	ctr.RouteHead(rt)
	//default:
	//	// unknow method
	//}
}

// controller default route
// 默认路由, 如果已经定义了则忽略，没有定义则加上
//func (ctr *Controller) DefaultRoutes() {
func (ctr *Controller) DefaultRoutes(c ControllerInterface) {
	var pattern, method, key string
	// GET /{endpoint}
	pattern = "/" + ctr.Endpoint
	method = "GET"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Get)
		ctr.Routes[key] = rt
	}

	// GET /{endpoint}/{id}
	pattern = "/" + ctr.Endpoint + "/:_id_"
	method = "GET"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Get)
		ctr.Routes[key] = rt
	}

	// POST /{endpoint}
	pattern = "/" + ctr.Endpoint
	method = "POST"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Post)
		ctr.Routes[key] = rt
	}

	// DELETE /{endpoint}/{id}
	pattern = "/" + ctr.Endpoint + "/:_id_"
	method = "DELETE"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Delete)
		ctr.Routes[key] = rt
	}

	// PATCH /{endpoint}/{id}
	pattern = "/" + ctr.Endpoint + "/:_id_"
	method = "PATCH"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Patch)
		ctr.Routes[key] = rt
	}

	// PUT /{endpoint}/{id}
	pattern = "/" + ctr.Endpoint + "/:_id_"
	method = "PUT"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Put)
		ctr.Routes[key] = rt
	}
}

func (ctr *Controller) RouteGet(rt *Route) {
	goji.Get(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RoutePost(rt *Route) {
	goji.Post(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RoutePut(rt *Route) {
	goji.Put(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RouteDelete(rt *Route) {
	goji.Delete(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RoutePatch(rt *Route) {
	goji.Patch(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RouteHead(rt *Route) {
	goji.Head(rt.Pattern, handlerWrap(rt.Handler))
}
<commit_msg>add not found handler<commit_after>// Ogo

package ogo

import (
	//"fmt"
	"github.com/zenazn/goji"
	"github.com/zenazn/goji/web"
	"net/http"
	"strings"
)

type Handler func(c *RESTContext)

type Route struct {
	Pattern string
	Method  string
	Handler Handler
}

type Controller struct {
	Endpoint string
	Routes   map[string]*Route
	ReqCount int //访问计数
}

type ControllerInterface interface {
	//Init(endpoint string, c ControllerInterface)
	Get(c *RESTContext)
	Post(c *RESTContext)
	Put(c *RESTContext)
	Delete(c *RESTContext)
	Patch(c *RESTContext)
	Head(c *RESTContext)
	Options(c *RESTContext)
	Trace(c *RESTContext)
	NotFound(c *RESTContext)
}

func NewRoute(p string, m string, h Handler) *Route {
	return &Route{
		Pattern: p,
		Method:  m,
		Handler: h,
	}
}

// 封装
func handlerWrap(f Handler) web.HandlerFunc { //这里封装了webC到本地的结构中
	return func(c web.C, w http.ResponseWriter, r *http.Request) {
		f(getContext(c, w, r))
	}
}

func (ctr *Controller) Init(endpoint string, c ControllerInterface) {
	ctr.Endpoint = endpoint
	//ctr.Routes = make(map[string]*Route)
	//默认路由
	ctr.DefaultRoutes(c)
	if len(ctr.Routes) > 0 {
		for _, rt := range ctr.Routes {
			switch strings.ToLower(rt.Method) {
			case "get":
				ctr.RouteGet(rt)
			case "post":
				ctr.RoutePost(rt)
			case "put":
				ctr.RoutePut(rt)
			case "delete":
				ctr.RouteDelete(rt)
			case "patch":
				ctr.RoutePatch(rt)
			case "head":
				ctr.RouteHead(rt)
			default:
				// unknow method
			}
		}
	}
}

func (ctr *Controller) Get(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Post(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Put(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Delete(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Patch(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Head(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) Options(c *RESTContext) {
	c.HTTPError(http.StatusMethodNotAllowed)
}
func (ctr *Controller) NotFound(c *RESTContext) {
	c.HTTPError(http.StatusNotFound)
}

func (ctr *Controller) AddRoute(m string, p string, h Handler) {
	key := strings.ToUpper(m) + " " + p
	if _, ok := ctr.Routes[key]; ok {
		//手动加路由, 以最后加的为准,overwrite
	}
	ctr.Routes[key] = NewRoute(p, m, h)
}

// controller default route
// 默认路由, 如果已经定义了则忽略，没有定义则加上
//func (ctr *Controller) DefaultRoutes() {
func (ctr *Controller) DefaultRoutes(c ControllerInterface) {
	var pattern, method, key string
	// GET /{endpoint}
	pattern = "/" + ctr.Endpoint
	method = "GET"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Get)
		ctr.Routes[key] = rt
	}

	// GET /{endpoint}/{id}
	pattern = "/" + ctr.Endpoint + "/:_id_"
	method = "GET"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Get)
		ctr.Routes[key] = rt
	}

	// POST /{endpoint}
	pattern = "/" + ctr.Endpoint
	method = "POST"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Post)
		ctr.Routes[key] = rt
	}

	// DELETE /{endpoint}/{id}
	pattern = "/" + ctr.Endpoint + "/:_id_"
	method = "DELETE"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Delete)
		ctr.Routes[key] = rt
	}

	// PATCH /{endpoint}/{id}
	pattern = "/" + ctr.Endpoint + "/:_id_"
	method = "PATCH"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Patch)
		ctr.Routes[key] = rt
	}

	// PUT /{endpoint}/{id}
	pattern = "/" + ctr.Endpoint + "/:_id_"
	method = "PUT"
	key = method + " " + pattern
	if _, ok := ctr.Routes[key]; ok {
		// exists, warning, 默认路由不能覆盖自定义路由
	} else {
		rt := NewRoute(pattern, method, c.Put)
		ctr.Routes[key] = rt
	}

	//Not Found
}

func (ctr *Controller) RouteGet(rt *Route) {
	goji.Get(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RoutePost(rt *Route) {
	goji.Post(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RoutePut(rt *Route) {
	goji.Put(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RouteDelete(rt *Route) {
	goji.Delete(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RoutePatch(rt *Route) {
	goji.Patch(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RouteHead(rt *Route) {
	goji.Head(rt.Pattern, handlerWrap(rt.Handler))
}

func (ctr *Controller) RouteNotFound(rt *Route) {
	goji.NotFound(handlerWrap(rt.Handler))
}
<|endoftext|>


<commit_before>package models

import "time"

type NotificationSettings struct {
	// unique idetifier of the notification setting
	Id int64 `json:"id,string"`

	// Id of the channel
	ChannelId int64 `json:"channelId,string"       sql:"NOT NULL"`

	// Creator of the notification settting
	AccountId int64 `json:"accountId,string"               sql:"NOT NULL"`

	// Holds dektop setting type
	DesktopSetting string `json:"desktopSetting"	sql:"NOT NULL"`

	// Holds mobile setting type
	MobileSetting string `json:"mobileSetting"			sql:"NOT NULL"`

	// Holds the data if channel is muted or not
	IsMuted bool `json:"isMuted"`

	// Holds data that getting notification when @channel is written
	// If user doesn't want to get notification
	// when written to channel @channel, @here or name of the user.
	// User uses 'suppress' feature and doesn't get notification
	IsSuppressed bool `json:"isSuppressed"`

	// Creation date of the notification settings
	CreatedAt time.Time `json:"createdAt"          sql:"NOT NULL"`

	// Modification date of the notification settings
	UpdatedAt time.Time `json:"updatedAt"          sql:"NOT NULL"`
}
<commit_msg>socialapi: constants are added for notification settings<commit_after>package models

import "time"

type NotificationSettings struct {
	// unique idetifier of the notification setting
	Id int64 `json:"id,string"`

	// Id of the channel
	ChannelId int64 `json:"channelId,string"       sql:"NOT NULL"`

	// Creator of the notification settting
	AccountId int64 `json:"accountId,string"               sql:"NOT NULL"`

	// Holds dektop setting type
	DesktopSetting string `json:"desktopSetting"	sql:"NOT NULL"`

	// Holds mobile setting type
	MobileSetting string `json:"mobileSetting"			sql:"NOT NULL"`

	// Holds the data if channel is muted or not
	IsMuted bool `json:"isMuted"`

	// Holds data that getting notification when @channel is written
	// If user doesn't want to get notification
	// when written to channel @channel, @here or name of the user.
	// User uses 'suppress' feature and doesn't get notification
	IsSuppressed bool `json:"isSuppressed"`

	// Creation date of the notification settings
	CreatedAt time.Time `json:"createdAt"          sql:"NOT NULL"`

	// Modification date of the notification settings
	UpdatedAt time.Time `json:"updatedAt"          sql:"NOT NULL"`
}

const (
	// Describes that user want to be notified for all notifications
	NotificationSettings_STATUS_ALL = "all"
	// Describes that user want to be notified
	// for user's own name or with highlighted words
	NotificationSettings_STATUS_PERSONAL = "personal"
	// Describes that user doesn't want to get any notification
	NotificationSettings_STATUS_NEVER = "never"
)
<|endoftext|>


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

// Package stream provides utility functions to consume Entry streams.
package stream // import "kythe.io/kythe/go/storage/stream"

import (
	"encoding/json"
	"fmt"
	"io"
	"log"

	"kythe.io/kythe/go/platform/delimited"
	"kythe.io/kythe/go/util/schema/facts"

	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"

	cpb "kythe.io/kythe/proto/common_go_proto"
	spb "kythe.io/kythe/proto/storage_go_proto"
)

// EntryReader functions read a stream of entries, passing each to a handler
// function.
type EntryReader func(func(*spb.Entry) error) error

// ReadEntries reads a stream of Entry protobufs from r.
func ReadEntries(r io.Reader) <-chan *spb.Entry {
	ch := make(chan *spb.Entry)
	go func() {
		defer close(ch)
		if err := NewReader(r)(func(e *spb.Entry) error {
			ch <- e
			return nil
		}); err != nil {
			log.Fatal(err)
		}
	}()
	return ch
}

// NewReader reads a stream of Entry protobufs from r.
func NewReader(r io.Reader) EntryReader {
	return func(f func(*spb.Entry) error) error {
		rd := delimited.NewReader(r)
		for {
			var entry spb.Entry
			if err := rd.NextProto(&entry); err == io.EOF {
				return nil
			} else if err != nil {
				return fmt.Errorf("error decoding Entry: %v", err)
			}
			if err := f((*spb.Entry)(&entry)); err != nil {
				return err
			}
		}
	}
}

// ReadJSONEntries reads a JSON stream of Entry protobufs from r.
func ReadJSONEntries(r io.Reader) <-chan *spb.Entry {
	ch := make(chan *spb.Entry)
	go func() {
		defer close(ch)
		if err := NewJSONReader(r)(func(e *spb.Entry) error {
			ch <- e
			return nil
		}); err != nil {
			log.Fatal(err)
		}
	}()
	return ch
}

// StructuredEntry has custom marshaling behavior to handle structured FactValues
type StructuredEntry spb.Entry

// Reset calls the implementation for Entry
func (r *StructuredEntry) Reset() {
	(*spb.Entry)(r).Reset()
}

// String calls the implementation for Entry
func (r *StructuredEntry) String() string {
	return (*spb.Entry)(r).String()
}

// ProtoMessage calls the implementation for Entry
func (r *StructuredEntry) ProtoMessage() {
	(*spb.Entry)(r).ProtoMessage()
}

// NewStructuredJSONReader reads a JSON stream of StructuredEntry protobufs from r.
func NewStructuredJSONReader(r io.Reader) EntryReader {
	return func(f func(*spb.Entry) error) error {
		de := json.NewDecoder(r)
		for {
			var entry StructuredEntry
			if err := de.Decode(&entry); err == io.EOF {
				return nil
			} else if err != nil {
				return fmt.Errorf("error decoding JSON Entry: %v", err)
			}
			if err := f((*spb.Entry)(&entry)); err != nil {
				return err
			}
		}
	}
}

// NewJSONReader reads a JSON stream of Entry protobufs from r.
func NewJSONReader(r io.Reader) EntryReader {
	return func(f func(*spb.Entry) error) error {
		de := json.NewDecoder(r)
		for {
			var raw json.RawMessage
			if err := de.Decode(&raw); err == io.EOF {
				return nil
			} else if err != nil {
				return fmt.Errorf("error decoding JSON Entry: %w", err)
			}
			var entry spb.Entry
			if err := protojson.Unmarshal(raw, &entry); err != nil {
				return fmt.Errorf("error decoding JSON Entry: %w", err)
			}
			if err := f(&entry); err != nil {
				return err
			}
		}
	}
}

var marshaler = protojson.MarshalOptions{UseProtoNames: true}

// richJSONEntry delays the unmarshaling of the fact_value field
type richJSONEntry struct {
	*spb.Entry
	FactValue json.RawMessage `json:"fact_value"`
}

// StructuredFactValueJSON creates a json object from e.FactValue
func StructuredFactValueJSON(e *spb.Entry) (json.RawMessage, error) {
	if e.FactName == facts.Code {
		var ms cpb.MarkedSource
		if err := proto.Unmarshal(e.FactValue, &ms); err != nil {
			return nil, err
		}
		rec, err := marshaler.Marshal(&ms)
		if err != nil {
			return nil, err
		}
		return rec, nil
	}
	return json.Marshal(e.FactValue)
}

// Structured creates an entry that serializes factValue to a full value
func Structured(e *spb.Entry) *StructuredEntry {
	return (*StructuredEntry)(e)
}

// UnmarshalJSON unmarshals r including an object representation of FactValue when appropriate
func (r *StructuredEntry) UnmarshalJSON(data []byte) error {
	var jsonEntry = richJSONEntry{(*spb.Entry)(r), nil}
	if err := json.Unmarshal(data, &jsonEntry); err != nil {
		return err
	}
	if jsonEntry.FactName == facts.Code {
		var ms cpb.MarkedSource
		if err := protojson.Unmarshal(jsonEntry.FactValue, &ms); err != nil {
			return err
		}
		pb, err := proto.Marshal(&ms)
		if err != nil {
			return err
		}
		jsonEntry.Entry.FactValue = pb
	} else if err := json.Unmarshal(jsonEntry.FactValue, &jsonEntry.Entry.FactValue); err != nil {
		return err
	}

	return nil
}

// MarshalJSON marshals r including an object representation of FactValue when appropriate
func (r *StructuredEntry) MarshalJSON() ([]byte, error) {
	var err error
	var jsonEntry = richJSONEntry{(*spb.Entry)(r), nil}
	if jsonEntry.FactValue, err = StructuredFactValueJSON((*spb.Entry)(r)); err != nil {
		return nil, err
	}

	return json.Marshal(jsonEntry)
}
<commit_msg>chore: use protojson to marshal StructuredEntry structs (#4809)<commit_after>/*
 * Copyright 2014 The Kythe Authors. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Package stream provides utility functions to consume Entry streams.
package stream // import "kythe.io/kythe/go/storage/stream"

import (
	"encoding/json"
	"fmt"
	"io"
	"log"

	"kythe.io/kythe/go/platform/delimited"
	"kythe.io/kythe/go/util/schema/facts"

	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"

	cpb "kythe.io/kythe/proto/common_go_proto"
	spb "kythe.io/kythe/proto/storage_go_proto"
)

// EntryReader functions read a stream of entries, passing each to a handler
// function.
type EntryReader func(func(*spb.Entry) error) error

// ReadEntries reads a stream of Entry protobufs from r.
func ReadEntries(r io.Reader) <-chan *spb.Entry {
	ch := make(chan *spb.Entry)
	go func() {
		defer close(ch)
		if err := NewReader(r)(func(e *spb.Entry) error {
			ch <- e
			return nil
		}); err != nil {
			log.Fatal(err)
		}
	}()
	return ch
}

// NewReader reads a stream of Entry protobufs from r.
func NewReader(r io.Reader) EntryReader {
	return func(f func(*spb.Entry) error) error {
		rd := delimited.NewReader(r)
		for {
			var entry spb.Entry
			if err := rd.NextProto(&entry); err == io.EOF {
				return nil
			} else if err != nil {
				return fmt.Errorf("error decoding Entry: %v", err)
			}
			if err := f((*spb.Entry)(&entry)); err != nil {
				return err
			}
		}
	}
}

// ReadJSONEntries reads a JSON stream of Entry protobufs from r.
func ReadJSONEntries(r io.Reader) <-chan *spb.Entry {
	ch := make(chan *spb.Entry)
	go func() {
		defer close(ch)
		if err := NewJSONReader(r)(func(e *spb.Entry) error {
			ch <- e
			return nil
		}); err != nil {
			log.Fatal(err)
		}
	}()
	return ch
}

// StructuredEntry has custom marshaling behavior to handle structured FactValues
type StructuredEntry spb.Entry

// Reset calls the implementation for Entry
func (r *StructuredEntry) Reset() {
	(*spb.Entry)(r).Reset()
}

// String calls the implementation for Entry
func (r *StructuredEntry) String() string {
	return (*spb.Entry)(r).String()
}

// ProtoMessage calls the implementation for Entry
func (r *StructuredEntry) ProtoMessage() {
	(*spb.Entry)(r).ProtoMessage()
}

// NewStructuredJSONReader reads a JSON stream of StructuredEntry protobufs from r.
func NewStructuredJSONReader(r io.Reader) EntryReader {
	return func(f func(*spb.Entry) error) error {
		de := json.NewDecoder(r)
		for {
			var entry StructuredEntry
			if err := de.Decode(&entry); err == io.EOF {
				return nil
			} else if err != nil {
				return fmt.Errorf("error decoding JSON Entry: %v", err)
			}
			if err := f((*spb.Entry)(&entry)); err != nil {
				return err
			}
		}
	}
}

// NewJSONReader reads a JSON stream of Entry protobufs from r.
func NewJSONReader(r io.Reader) EntryReader {
	return func(f func(*spb.Entry) error) error {
		de := json.NewDecoder(r)
		for {
			var raw json.RawMessage
			if err := de.Decode(&raw); err == io.EOF {
				return nil
			} else if err != nil {
				return fmt.Errorf("error decoding JSON Entry: %w", err)
			}
			var entry spb.Entry
			if err := protojson.Unmarshal(raw, &entry); err != nil {
				return fmt.Errorf("error decoding JSON Entry: %w", err)
			}
			if err := f(&entry); err != nil {
				return err
			}
		}
	}
}

var marshaler = protojson.MarshalOptions{UseProtoNames: true}

// richJSONEntry delays the unmarshaling of the fact_value field
type richJSONEntry struct {
	Source    json.RawMessage `json:"source,omitempty"`
	Target    json.RawMessage `json:"target,omitempty"`
	EdgeKind  string          `json:"edge_kind,omitempty"`
	FactName  string          `json:"fact_name,omitempty"`
	FactValue json.RawMessage `json:"fact_value,omitempty"`
}

// StructuredFactValueJSON creates a json object from e.FactValue
func StructuredFactValueJSON(e *spb.Entry) (json.RawMessage, error) {
	if e.FactName == facts.Code {
		var ms cpb.MarkedSource
		if err := proto.Unmarshal(e.FactValue, &ms); err != nil {
			return nil, err
		}
		rec, err := marshaler.Marshal(&ms)
		if err != nil {
			return nil, err
		}
		return rec, nil
	}
	return json.Marshal(e.FactValue)
}

// Structured creates an entry that serializes factValue to a full value
func Structured(e *spb.Entry) *StructuredEntry {
	return (*StructuredEntry)(e)
}

// UnmarshalJSON unmarshals r including an object representation of FactValue when appropriate
func (r *StructuredEntry) UnmarshalJSON(data []byte) error {
	var jsonEntry richJSONEntry
	if err := json.Unmarshal(data, &jsonEntry); err != nil {
		return err
	}
	if jsonEntry.FactName == facts.Code {
		var ms cpb.MarkedSource
		if err := protojson.Unmarshal(jsonEntry.FactValue, &ms); err != nil {
			return err
		}
		pb, err := proto.Marshal(&ms)
		if err != nil {
			return err
		}
		r.FactValue = pb
	} else if err := json.Unmarshal(jsonEntry.FactValue, &r.FactValue); err != nil {
		return err
	}

	r.EdgeKind = jsonEntry.EdgeKind
	r.FactName = jsonEntry.FactName

	var err error
	if r.Source, err = unmarshalVName(jsonEntry.Source); err != nil {
		return err
	} else if r.Target, err = unmarshalVName(jsonEntry.Target); err != nil {
		return err
	}

	return nil
}

func marshalVName(v *spb.VName) (json.RawMessage, error) {
	if v == nil {
		return nil, nil
	}
	return protojson.Marshal(v)
}

func unmarshalVName(msg json.RawMessage) (*spb.VName, error) {
	if len(msg) == 0 {
		return nil, nil
	}
	var v spb.VName
	return &v, protojson.Unmarshal(msg, &v)
}

// MarshalJSON marshals r including an object representation of FactValue when appropriate
func (r *StructuredEntry) MarshalJSON() ([]byte, error) {
	jsonEntry := richJSONEntry{
		EdgeKind: r.EdgeKind,
		FactName: r.FactName,
	}
	var err error
	if jsonEntry.Source, err = marshalVName(r.Source); err != nil {
		return nil, err
	} else if jsonEntry.Target, err = marshalVName(r.Target); err != nil {
		return nil, err
	} else if jsonEntry.FactValue, err = StructuredFactValueJSON((*spb.Entry)(r)); err != nil {
		return nil, err
	}

	return json.Marshal(jsonEntry)
}
<|endoftext|>


<commit_before>package core

import (
	"errors"
	"log"

	"github.com/nikhan/go-fetch"
)

// NewBlock creates a new block from a spec
func NewBlock(s Spec) *Block {
	var in []Route
	var out []Output

	for _, v := range s.Inputs {

		q, _ := fetch.Parse(".")
		in = append(in, Route{
			Name:  v.Name,
			Value: q,
			C:     make(chan Message),
		})
	}

	for _, v := range s.Outputs {
		out = append(out, Output{
			Name:        v.Name,
			Connections: make(map[Connection]struct{}),
		})
	}

	return &Block{
		state: BlockState{
			make(MessageMap),
			make(MessageMap),
			make(MessageMap),
			make(Manifest),
			false,
		},
		routing: BlockRouting{
			Inputs:        in,
			Outputs:       out,
			InterruptChan: make(chan Interrupt),
			Shared: SharedStore{
				Type: s.Shared,
			},
		},
		kernel: s.Kernel,
	}
}

// suture: the main routine the block runs
func (b *Block) Serve() {
	for {
		var interrupt Interrupt

		b.routing.RLock()
		for {
			interrupt = b.receive()
			if interrupt != nil {
				break
			}

			interrupt = b.process()
			if interrupt != nil {
				break
			}

			interrupt = b.broadcast()
			if interrupt != nil {
				break
			}

			b.crank()
		}
		b.routing.RUnlock()
		b.routing.Lock()
		if ok := interrupt(); !ok {
			return
		}
		b.routing.Unlock()
	}
}

func (b *Block) exportRoute(id RouteID) (*Route, error) {
	if int(id) >= len(b.routing.Inputs) || int(id) < 0 {
		return nil, errors.New("index out of range")
	}

	var v interface{}
	switch n := b.routing.Inputs[id].Value.(type) {
	case *fetch.Query:
		// yuck copy
		v, _ = fetch.Parse(n.String())
	default:
		v = Copy(n)
	}

	return &Route{
		Value: v,
		C:     b.routing.Inputs[id].C,
		Name:  b.routing.Inputs[id].Name,
	}, nil
}

// Input returns the specfied Route
func (b *Block) GetRoute(id RouteID) (*Route, error) {
	b.routing.RLock()
	r, err := b.exportRoute(id)
	b.routing.RUnlock()
	return r, err
}

func (b *Block) GetRoutes() []Route {
	b.routing.RLock()
	re := make([]Route, len(b.routing.Inputs), len(b.routing.Inputs))
	for i, _ := range b.routing.Inputs {
		r, _ := b.exportRoute(RouteID(i))
		re[i] = *r
	}
	b.routing.RUnlock()
	return re
}

// Outputs return a list of manifest pairs for the block
func (b *Block) GetOutputs() []Output {
	b.routing.RLock()
	m := make([]Output, len(b.routing.Outputs), len(b.routing.Outputs))
	for id, out := range b.routing.Outputs {
		m[id] = out
	}
	b.routing.RUnlock()
	return m
}

func (b *Block) GetStore() Store {
	b.routing.RLock()
	v := b.routing.Shared.Store
	b.routing.RUnlock()
	return v
}

// sets a store for the block. can be set to nil
func (b *Block) SetStore(s Store) {
	b.routing.InterruptChan <- func() bool {
		b.routing.Shared.Store = s
		return true
	}
}

// RouteValue sets the route to always be the specified value
func (b *Block) SetRoute(id RouteID, v interface{}) error {
	returnVal := make(chan error, 1)
	b.routing.InterruptChan <- func() bool {
		if int(id) < 0 || int(id) >= len(b.routing.Inputs) {
			returnVal <- errors.New("input out of range")
			return true
		}

		b.routing.Inputs[id].Value = v
		returnVal <- nil
		return true
	}
	return <-returnVal
}

// Connect connects a Route, specified by ID, to a connection
func (b *Block) Connect(id RouteID, c Connection) error {
	returnVal := make(chan error, 1)
	b.routing.InterruptChan <- func() bool {
		if int(id) < 0 || int(id) >= len(b.routing.Outputs) {
			returnVal <- errors.New("output out of range")
			return true
		}

		if _, ok := b.routing.Outputs[id].Connections[c]; ok {
			returnVal <- errors.New("this connection already exists on this output")
			return true
		}

		b.routing.Outputs[id].Connections[c] = struct{}{}
		returnVal <- nil
		return true
	}
	return <-returnVal
}

// Disconnect removes a connection from a Route
func (b *Block) Disconnect(id RouteID, c Connection) error {
	returnVal := make(chan error, 1)
	b.routing.InterruptChan <- func() bool {
		if int(id) < 0 || int(id) >= len(b.routing.Outputs) {
			returnVal <- errors.New("output out of range")
			return true
		}

		if _, ok := b.routing.Outputs[id].Connections[c]; !ok {
			returnVal <- errors.New("connection does not exist")
			return true
		}

		delete(b.routing.Outputs[id].Connections, c)
		returnVal <- nil
		return true
	}
	return <-returnVal
}

// suture: stop the block
func (b *Block) Stop() {
	b.routing.InterruptChan <- func() bool {
		return false
	}
}

// wait and listen for all kernel inputs to be filled.
func (b *Block) receive() Interrupt {
	var err error
	for id, input := range b.routing.Inputs {
		//if we have already received a value on this input, skip.
		if _, ok := b.state.inputValues[RouteID(id)]; ok {
			continue
		}

		// if there is a value set for this input, place value on
		// buffer and set it in map.
		query, ok := input.Value.(*fetch.Query)
		if !ok {
			b.state.inputValues[RouteID(id)] = Copy(input.Value)
			continue
		}

		select {
		case m := <-input.C:
			b.state.inputValues[RouteID(id)], err = fetch.Run(query, m)
			if err != nil {
				log.Fatal(err)
			}
		case f := <-b.routing.InterruptChan:
			return f
		}
	}
	return nil
}

// run kernel on inputs, produce outputs
func (b *Block) process() Interrupt {
	if b.state.Processed == true {
		return nil
	}

	// if this kernel relies on an external shared state then we need to
	// block until an interrupt connects us to a shared external state.
	if b.routing.Shared.Type != NONE && b.routing.Shared.Store == nil {
		select {
		case f := <-b.routing.InterruptChan:
			return f
		}
	}

	// we should only be able to get here if
	// - we don't need an shared state
	// - we have an external shared state and it has been attached
	if b.routing.Shared.Type != NONE {
		b.routing.Shared.Store.Lock()
	}

	// run the kernel
	interrupt := b.kernel(b.state.inputValues,
		b.state.outputValues,
		b.state.internalValues,
		b.routing.Shared.Store,
		b.routing.InterruptChan)

	if interrupt != nil {
		if b.routing.Shared.Type != NONE {
			b.routing.Shared.Store.Unlock()
		}
		return interrupt
	}

	if b.routing.Shared.Type != NONE {
		b.routing.Shared.Store.Unlock()
	}

	b.state.Processed = true

	return nil
}

// broadcast the kernel output to all connections on all outputs.
func (b *Block) broadcast() Interrupt {
	for id, out := range b.routing.Outputs {
		// if there no connection for this output then wait until there
		// is one. that means we have to wait for an interrupt.
		if len(out.Connections) == 0 {
			select {
			case f := <-b.routing.InterruptChan:
				return f
			}
		}
		for c, _ := range out.Connections {
			// check to see if we have delivered a message to this
			// connection for this block crank. if we have, then
			// skip this delivery.
			m := ManifestPair{id, c}
			if _, ok := b.state.manifest[m]; ok {
				continue
			}

			select {
			case c <- b.state.outputValues[RouteID(id)]:
				// set that we have delivered the message.
				b.state.manifest[m] = struct{}{}
			case f := <-b.routing.InterruptChan:
				return f
			}
		}

	}
	return nil
}

// cleanup all block state for this crank of the block
func (b *Block) crank() {
	for k, _ := range b.state.inputValues {
		delete(b.state.inputValues, k)
	}
	for k, _ := range b.state.outputValues {
		delete(b.state.outputValues, k)
	}
	for k, _ := range b.state.manifest {
		delete(b.state.manifest, k)
	}
	b.state.Processed = false
}
<commit_msg>fixing block outputs<commit_after>package core

import (
	"errors"
	"log"

	"github.com/nikhan/go-fetch"
)

// NewBlock creates a new block from a spec
func NewBlock(s Spec) *Block {
	var in []Route
	var out []Output

	for _, v := range s.Inputs {

		q, _ := fetch.Parse(".")
		in = append(in, Route{
			Name:  v.Name,
			Value: q,
			C:     make(chan Message),
		})
	}

	for _, v := range s.Outputs {
		out = append(out, Output{
			Name:        v.Name,
			Connections: make(map[Connection]struct{}),
		})
	}

	return &Block{
		state: BlockState{
			make(MessageMap),
			make(MessageMap),
			make(MessageMap),
			make(Manifest),
			false,
		},
		routing: BlockRouting{
			Inputs:        in,
			Outputs:       out,
			InterruptChan: make(chan Interrupt),
			Shared: SharedStore{
				Type: s.Shared,
			},
		},
		kernel: s.Kernel,
	}
}

// suture: the main routine the block runs
func (b *Block) Serve() {
	for {
		var interrupt Interrupt

		b.routing.RLock()
		for {
			interrupt = b.receive()
			if interrupt != nil {
				break
			}

			interrupt = b.process()
			if interrupt != nil {
				break
			}

			interrupt = b.broadcast()
			if interrupt != nil {
				break
			}

			b.crank()
		}
		b.routing.RUnlock()
		b.routing.Lock()
		if ok := interrupt(); !ok {
			return
		}
		b.routing.Unlock()
	}
}

func (b *Block) exportRoute(id RouteID) (*Route, error) {
	if int(id) >= len(b.routing.Inputs) || int(id) < 0 {
		return nil, errors.New("index out of range")
	}

	var v interface{}
	switch n := b.routing.Inputs[id].Value.(type) {
	case *fetch.Query:
		// yuck copy
		v, _ = fetch.Parse(n.String())
	default:
		v = Copy(n)
	}

	return &Route{
		Value: v,
		C:     b.routing.Inputs[id].C,
		Name:  b.routing.Inputs[id].Name,
	}, nil
}

// Input returns the specfied Route
func (b *Block) GetRoute(id RouteID) (*Route, error) {
	b.routing.RLock()
	r, err := b.exportRoute(id)
	b.routing.RUnlock()
	return r, err
}

func (b *Block) GetRoutes() []*Route {
	b.routing.RLock()
	re := make([]*Route, len(b.routing.Inputs), len(b.routing.Inputs))
	for i, _ := range b.routing.Inputs {
		re[i], _ = b.exportRoute(RouteID(i))
	}
	b.routing.RUnlock()
	return re
}

// Outputs return a list of manifest pairs for the block
func (b *Block) GetOutputs() []Output {
	b.routing.RLock()
	m := make([]Output, len(b.routing.Outputs), len(b.routing.Outputs))
	for id, out := range b.routing.Outputs {
		m[id] = out
	}
	b.routing.RUnlock()
	return m
}

func (b *Block) GetStore() Store {
	b.routing.RLock()
	v := b.routing.Shared.Store
	b.routing.RUnlock()
	return v
}

// sets a store for the block. can be set to nil
func (b *Block) SetStore(s Store) {
	b.routing.InterruptChan <- func() bool {
		b.routing.Shared.Store = s
		return true
	}
}

// RouteValue sets the route to always be the specified value
func (b *Block) SetRoute(id RouteID, v interface{}) error {
	returnVal := make(chan error, 1)
	b.routing.InterruptChan <- func() bool {
		if int(id) < 0 || int(id) >= len(b.routing.Inputs) {
			returnVal <- errors.New("input out of range")
			return true
		}

		b.routing.Inputs[id].Value = v
		returnVal <- nil
		return true
	}
	return <-returnVal
}

// Connect connects a Route, specified by ID, to a connection
func (b *Block) Connect(id RouteID, c Connection) error {
	returnVal := make(chan error, 1)
	b.routing.InterruptChan <- func() bool {
		if int(id) < 0 || int(id) >= len(b.routing.Outputs) {
			returnVal <- errors.New("output out of range")
			return true
		}

		if _, ok := b.routing.Outputs[id].Connections[c]; ok {
			returnVal <- errors.New("this connection already exists on this output")
			return true
		}

		b.routing.Outputs[id].Connections[c] = struct{}{}
		returnVal <- nil
		return true
	}
	return <-returnVal
}

// Disconnect removes a connection from a Route
func (b *Block) Disconnect(id RouteID, c Connection) error {
	returnVal := make(chan error, 1)
	b.routing.InterruptChan <- func() bool {
		if int(id) < 0 || int(id) >= len(b.routing.Outputs) {
			returnVal <- errors.New("output out of range")
			return true
		}

		if _, ok := b.routing.Outputs[id].Connections[c]; !ok {
			returnVal <- errors.New("connection does not exist")
			return true
		}

		delete(b.routing.Outputs[id].Connections, c)
		returnVal <- nil
		return true
	}
	return <-returnVal
}

// suture: stop the block
func (b *Block) Stop() {
	b.routing.InterruptChan <- func() bool {
		return false
	}
}

// wait and listen for all kernel inputs to be filled.
func (b *Block) receive() Interrupt {
	var err error
	for id, input := range b.routing.Inputs {
		//if we have already received a value on this input, skip.
		if _, ok := b.state.inputValues[RouteID(id)]; ok {
			continue
		}

		// if there is a value set for this input, place value on
		// buffer and set it in map.
		query, ok := input.Value.(*fetch.Query)
		if !ok {
			b.state.inputValues[RouteID(id)] = Copy(input.Value)
			continue
		}

		select {
		case m := <-input.C:
			b.state.inputValues[RouteID(id)], err = fetch.Run(query, m)
			if err != nil {
				log.Fatal(err)
			}
		case f := <-b.routing.InterruptChan:
			return f
		}
	}
	return nil
}

// run kernel on inputs, produce outputs
func (b *Block) process() Interrupt {
	if b.state.Processed == true {
		return nil
	}

	// if this kernel relies on an external shared state then we need to
	// block until an interrupt connects us to a shared external state.
	if b.routing.Shared.Type != NONE && b.routing.Shared.Store == nil {
		select {
		case f := <-b.routing.InterruptChan:
			return f
		}
	}

	// we should only be able to get here if
	// - we don't need an shared state
	// - we have an external shared state and it has been attached
	if b.routing.Shared.Type != NONE {
		b.routing.Shared.Store.Lock()
	}

	// run the kernel
	interrupt := b.kernel(b.state.inputValues,
		b.state.outputValues,
		b.state.internalValues,
		b.routing.Shared.Store,
		b.routing.InterruptChan)

	if interrupt != nil {
		if b.routing.Shared.Type != NONE {
			b.routing.Shared.Store.Unlock()
		}
		return interrupt
	}

	if b.routing.Shared.Type != NONE {
		b.routing.Shared.Store.Unlock()
	}

	b.state.Processed = true

	return nil
}

// broadcast the kernel output to all connections on all outputs.
func (b *Block) broadcast() Interrupt {
	for id, out := range b.routing.Outputs {
		// if there no connection for this output then wait until there
		// is one. that means we have to wait for an interrupt.
		if len(out.Connections) == 0 {
			select {
			case f := <-b.routing.InterruptChan:
				return f
			}
		}
		for c, _ := range out.Connections {
			// check to see if we have delivered a message to this
			// connection for this block crank. if we have, then
			// skip this delivery.
			m := ManifestPair{id, c}
			if _, ok := b.state.manifest[m]; ok {
				continue
			}

			select {
			case c <- b.state.outputValues[RouteID(id)]:
				// set that we have delivered the message.
				b.state.manifest[m] = struct{}{}
			case f := <-b.routing.InterruptChan:
				return f
			}
		}

	}
	return nil
}

// cleanup all block state for this crank of the block
func (b *Block) crank() {
	for k, _ := range b.state.inputValues {
		delete(b.state.inputValues, k)
	}
	for k, _ := range b.state.outputValues {
		delete(b.state.outputValues, k)
	}
	for k, _ := range b.state.manifest {
		delete(b.state.manifest, k)
	}
	b.state.Processed = false
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"time"

	"github.com/centrifugal/centrifugo/logger"

	"github.com/garyburd/redigo/redis"
)

// redisEngine uses Redis datastructures and PUB/SUB to manage Centrifuge logic.
// This engine allows to scale Centrifuge - you can run several Centrifuge instances
// connected to the same Redis and load balance clients between instances.
type redisEngine struct {
	app       *application
	pool      *redis.Pool
	psc       redis.PubSubConn
	connected bool
}

func newRedisEngine(app *application, host, port, password, db, url string, api bool) *redisEngine {
	server := host + ":" + port
	pool := newPool(server, password, db)
	return &redisEngine{
		app:  app,
		pool: pool,
	}
}

func newPool(server, password, db string) *redis.Pool {
	return &redis.Pool{
		MaxIdle:     3,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", server)
			if err != nil {
				logger.CRITICAL.Println(err)
				return nil, err
			}
			if password != "" {
				if _, err := c.Do("AUTH", password); err != nil {
					c.Close()
					logger.CRITICAL.Println(err)
					return nil, err
				}
			}
			if _, err := c.Do("SELECT", db); err != nil {
				c.Close()
				logger.CRITICAL.Println(err)
				return nil, err
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}

func (e *redisEngine) getName() string {
	return "Redis"
}

func (e *redisEngine) initialize() error {
	go e.initializePubSub()
	go e.checkConnectionStatus()
	return nil
}

func (e *redisEngine) checkConnectionStatus() {
	for {
		time.Sleep(time.Second)
		if e.connected {
			continue
		}
		go e.initializePubSub()
	}
}

func (e *redisEngine) initializePubSub() {
	e.connected = true
	e.psc = redis.PubSubConn{e.pool.Get()}
	defer e.psc.Close()
	err := e.psc.Subscribe(e.app.adminChannel)
	if err != nil {
		e.connected = false
		e.psc.Close()
		return
	}
	err = e.psc.Subscribe(e.app.controlChannel)
	if err != nil {
		e.connected = false
		e.psc.Close()
		return
	}
	for _, channel := range e.app.clientSubscriptionHub.getChannels() {
		err = e.psc.Subscribe(channel)
		if err != nil {
			e.connected = false
			e.psc.Close()
			return
		}
	}
	for {
		switch n := e.psc.Receive().(type) {
		case redis.Message:
			e.app.handleMessage(n.Channel, n.Data)
		case redis.Subscription:
		case error:
			logger.ERROR.Printf("error: %v\n", n)
			e.psc.Close()
			e.connected = false
			return
		}
	}
}

func (e *redisEngine) publish(channel string, message []byte) error {
	conn := e.pool.Get()
	defer conn.Close()
	_, err := conn.Do("PUBLISH", channel, message)
	return err
}

func (e *redisEngine) subscribe(channel string) error {
	return e.psc.Subscribe(channel)
}

func (e *redisEngine) unsubscribe(channel string) error {
	return e.psc.Unsubscribe(channel)
}

func (e *redisEngine) getHashKey(channel string) string {
	return e.app.channelPrefix + ".presence.hash." + channel
}

func (e *redisEngine) getSetKey(channel string) string {
	return e.app.channelPrefix + ".presence.set." + channel
}

func (e *redisEngine) addPresence(channel, uid string, info interface{}) error {
	conn := e.pool.Get()
	defer conn.Close()
	infoJson, err := json.Marshal(info)
	if err != nil {
		return err
	}
	expireAt := time.Now().Unix() + e.app.presenceExpireInterval
	hashKey := e.getHashKey(channel)
	setKey := e.getSetKey(channel)
	conn.Send("MULTI")
	conn.Send("ZADD", setKey, expireAt, uid)
	conn.Send("HSET", hashKey, uid, infoJson)
	_, err = conn.Do("EXEC")
	return err
}

func (e *redisEngine) removePresence(channel, uid string) error {
	conn := e.pool.Get()
	defer conn.Close()
	hashKey := e.getHashKey(channel)
	setKey := e.getSetKey(channel)
	conn.Send("MULTI")
	conn.Send("HDEL", hashKey, uid)
	conn.Send("ZREM", setKey, uid)
	_, err := conn.Do("EXEC")
	return err
}

func mapStringInterface(result interface{}, err error) (map[string]interface{}, error) {
	values, err := redis.Values(result, err)
	if err != nil {
		return nil, err
	}
	if len(values)%2 != 0 {
		return nil, errors.New("mapStringInterface expects even number of values result")
	}
	m := make(map[string]interface{}, len(values)/2)
	for i := 0; i < len(values); i += 2 {
		key, okKey := values[i].([]byte)
		value, okValue := values[i+1].([]byte)
		if !okKey || !okValue {
			return nil, errors.New("ScanMap key not a bulk string value")
		}
		var f interface{}
		err = json.Unmarshal(value, &f)
		if err != nil {
			return nil, errors.New("can not unmarshal value to interface")
		}
		m[string(key)] = f
	}
	return m, nil
}

func (e *redisEngine) getPresence(channel string) (map[string]interface{}, error) {
	conn := e.pool.Get()
	defer conn.Close()
	now := time.Now().Unix()
	hashKey := e.getHashKey(channel)
	setKey := e.getSetKey(channel)
	reply, err := conn.Do("ZRANGEBYSCORE", setKey, 0, now)
	if err != nil {
		return nil, err
	}
	expiredKeys, err := redis.Strings(reply, nil)
	if err != nil {
		return nil, err
	}
	if len(expiredKeys) > 0 {
		conn.Send("MULTI")
		conn.Send("ZREMRANGEBYSCORE", setKey, 0, now)
		for _, key := range expiredKeys {
			conn.Send("HDEL", hashKey, key)
		}
		_, err = conn.Do("EXEC")
		if err != nil {
			return nil, err
		}
	}
	reply, err = conn.Do("HGETALL", hashKey)
	if err != nil {
		return nil, err
	}
	presence, err := mapStringInterface(reply, nil)
	return presence, err
}

func (e *redisEngine) addHistoryMessage(channel string, message interface{}, size, lifetime int64) error {
	conn := e.pool.Get()
	defer conn.Close()
	return nil
}

func (e *redisEngine) getHistory(channel string) ([]interface{}, error) {
	conn := e.pool.Get()
	defer conn.Close()
	return []interface{}{}, nil
}
<commit_msg>redis engine history<commit_after>package main

import (
	"encoding/json"
	"errors"
	"time"

	"github.com/centrifugal/centrifugo/logger"

	"github.com/garyburd/redigo/redis"
)

// redisEngine uses Redis datastructures and PUB/SUB to manage Centrifuge logic.
// This engine allows to scale Centrifuge - you can run several Centrifuge instances
// connected to the same Redis and load balance clients between instances.
type redisEngine struct {
	app       *application
	pool      *redis.Pool
	psc       redis.PubSubConn
	connected bool
}

func newRedisEngine(app *application, host, port, password, db, url string, api bool) *redisEngine {
	server := host + ":" + port
	pool := newPool(server, password, db)
	return &redisEngine{
		app:  app,
		pool: pool,
	}
}

func newPool(server, password, db string) *redis.Pool {
	return &redis.Pool{
		MaxIdle:     3,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", server)
			if err != nil {
				logger.CRITICAL.Println(err)
				return nil, err
			}
			if password != "" {
				if _, err := c.Do("AUTH", password); err != nil {
					c.Close()
					logger.CRITICAL.Println(err)
					return nil, err
				}
			}
			if _, err := c.Do("SELECT", db); err != nil {
				c.Close()
				logger.CRITICAL.Println(err)
				return nil, err
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}

func (e *redisEngine) getName() string {
	return "Redis"
}

func (e *redisEngine) initialize() error {
	go e.initializePubSub()
	go e.checkConnectionStatus()
	return nil
}

func (e *redisEngine) checkConnectionStatus() {
	for {
		time.Sleep(time.Second)
		if e.connected {
			continue
		}
		go e.initializePubSub()
	}
}

func (e *redisEngine) initializePubSub() {
	e.connected = true
	e.psc = redis.PubSubConn{e.pool.Get()}
	defer e.psc.Close()
	err := e.psc.Subscribe(e.app.adminChannel)
	if err != nil {
		e.connected = false
		e.psc.Close()
		return
	}
	err = e.psc.Subscribe(e.app.controlChannel)
	if err != nil {
		e.connected = false
		e.psc.Close()
		return
	}
	for _, channel := range e.app.clientSubscriptionHub.getChannels() {
		err = e.psc.Subscribe(channel)
		if err != nil {
			e.connected = false
			e.psc.Close()
			return
		}
	}
	for {
		switch n := e.psc.Receive().(type) {
		case redis.Message:
			e.app.handleMessage(n.Channel, n.Data)
		case redis.Subscription:
		case error:
			logger.ERROR.Printf("error: %v\n", n)
			e.psc.Close()
			e.connected = false
			return
		}
	}
}

func (e *redisEngine) publish(channel string, message []byte) error {
	conn := e.pool.Get()
	defer conn.Close()
	_, err := conn.Do("PUBLISH", channel, message)
	return err
}

func (e *redisEngine) subscribe(channel string) error {
	return e.psc.Subscribe(channel)
}

func (e *redisEngine) unsubscribe(channel string) error {
	return e.psc.Unsubscribe(channel)
}

func (e *redisEngine) getHashKey(channel string) string {
	return e.app.channelPrefix + ".presence.hash." + channel
}

func (e *redisEngine) getSetKey(channel string) string {
	return e.app.channelPrefix + ".presence.set." + channel
}

func (e *redisEngine) getHistoryKey(channel string) string {
	return e.app.channelPrefix + ".history.list." + channel
}

func (e *redisEngine) addPresence(channel, uid string, info interface{}) error {
	conn := e.pool.Get()
	defer conn.Close()
	infoJson, err := json.Marshal(info)
	if err != nil {
		return err
	}
	expireAt := time.Now().Unix() + e.app.presenceExpireInterval
	hashKey := e.getHashKey(channel)
	setKey := e.getSetKey(channel)
	conn.Send("MULTI")
	conn.Send("ZADD", setKey, expireAt, uid)
	conn.Send("HSET", hashKey, uid, infoJson)
	_, err = conn.Do("EXEC")
	return err
}

func (e *redisEngine) removePresence(channel, uid string) error {
	conn := e.pool.Get()
	defer conn.Close()
	hashKey := e.getHashKey(channel)
	setKey := e.getSetKey(channel)
	conn.Send("MULTI")
	conn.Send("HDEL", hashKey, uid)
	conn.Send("ZREM", setKey, uid)
	_, err := conn.Do("EXEC")
	return err
}

func mapStringInterface(result interface{}, err error) (map[string]interface{}, error) {
	values, err := redis.Values(result, err)
	if err != nil {
		return nil, err
	}
	if len(values)%2 != 0 {
		return nil, errors.New("mapStringInterface expects even number of values result")
	}
	m := make(map[string]interface{}, len(values)/2)
	for i := 0; i < len(values); i += 2 {
		key, okKey := values[i].([]byte)
		value, okValue := values[i+1].([]byte)
		if !okKey || !okValue {
			return nil, errors.New("ScanMap key not a bulk string value")
		}
		var f interface{}
		err = json.Unmarshal(value, &f)
		if err != nil {
			return nil, errors.New("can not unmarshal value to interface")
		}
		m[string(key)] = f
	}
	return m, nil
}

func (e *redisEngine) getPresence(channel string) (map[string]interface{}, error) {
	conn := e.pool.Get()
	defer conn.Close()
	now := time.Now().Unix()
	hashKey := e.getHashKey(channel)
	setKey := e.getSetKey(channel)
	reply, err := conn.Do("ZRANGEBYSCORE", setKey, 0, now)
	if err != nil {
		return nil, err
	}
	expiredKeys, err := redis.Strings(reply, nil)
	if err != nil {
		return nil, err
	}
	if len(expiredKeys) > 0 {
		conn.Send("MULTI")
		conn.Send("ZREMRANGEBYSCORE", setKey, 0, now)
		for _, key := range expiredKeys {
			conn.Send("HDEL", hashKey, key)
		}
		_, err = conn.Do("EXEC")
		if err != nil {
			return nil, err
		}
	}
	reply, err = conn.Do("HGETALL", hashKey)
	if err != nil {
		return nil, err
	}
	presence, err := mapStringInterface(reply, nil)
	return presence, err
}

func (e *redisEngine) addHistoryMessage(channel string, message interface{}, size, lifetime int64) error {
	conn := e.pool.Get()
	defer conn.Close()
	if size <= 0 {
		return nil
	}
	historyKey := e.getHistoryKey(channel)
	messageJson, err := json.Marshal(message)
	if err != nil {
		return err
	}
	conn.Send("MULTI")
	conn.Send("LPUSH", historyKey, messageJson)
	conn.Send("LTRIM", historyKey, 0, size-1)
	if lifetime <= 0 {
		conn.Send("PERSIST", historyKey)
	} else {
		conn.Send("EXPIRE", historyKey, lifetime)
	}
	_, err = conn.Do("EXEC")
	return err
}

func (e *redisEngine) getHistory(channel string) ([]interface{}, error) {
	conn := e.pool.Get()
	defer conn.Close()
	historyKey := e.getHistoryKey(channel)
	values, err := redis.Values(conn.Do("LRANGE", historyKey, 0, -1))
	return values, err
}
<|endoftext|>


<commit_before>//	The MIT License (MIT)
//
//	Copyright (c) 2016, Cagatay Dogan
//
//	Permission is hereby granted, free of charge, to any person obtaining a copy
//	of this software and associated documentation files (the "Software"), to deal
//	in the Software without restriction, including without limitation the rights
//	to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//	copies of the Software, and to permit persons to whom the Software is
//	furnished to do so, subject to the following conditions:
//
//		The above copyright notice and this permission notice shall be included in
//		all copies or substantial portions of the Software.
//
//		THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//		IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//		FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//		AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//		LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//		OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//		THE SOFTWARE.

package lib

import (
	"sync"
)

type WorkGroup struct {
	workerCount int
	wg          sync.WaitGroup
}

func (w *WorkGroup) Add(delta int) {
	if delta != 0 {
		w.workerCount += delta
		w.wg.Add(delta)
	}
}

func (w *WorkGroup) Done() {
	if w.workerCount > 0 {
		w.workerCount--
		w.wg.Done()
	}
}

func (w *WorkGroup) Count() int {
	return w.workerCount
}

func (w *WorkGroup) Wait() {
	if w.workerCount > 0 {
		w.wg.Wait()
		w.workerCount = 0
	}
}

func (w *WorkGroup) Close() {
	if w.workerCount > 0 {
		w.wg.Add(-w.workerCount)
		w.workerCount = 0
	}
}
<commit_msg>Atomic update<commit_after>//	The MIT License (MIT)
//
//	Copyright (c) 2016, Cagatay Dogan
//
//	Permission is hereby granted, free of charge, to any person obtaining a copy
//	of this software and associated documentation files (the "Software"), to deal
//	in the Software without restriction, including without limitation the rights
//	to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//	copies of the Software, and to permit persons to whom the Software is
//	furnished to do so, subject to the following conditions:
//
//		The above copyright notice and this permission notice shall be included in
//		all copies or substantial portions of the Software.
//
//		THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//		IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//		FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//		AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//		LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//		OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//		THE SOFTWARE.

package lib

import (
	"sync"
	"sync/atomic"
)

type WorkGroup struct {
	workerCount int64
	wg          sync.WaitGroup
}

func (w *WorkGroup) Add(delta int) {
	if delta != 0 {
		atomic.AddInt64(&w.workerCount, int64(delta))
		w.wg.Add(delta)
	}
}

func (w *WorkGroup) Done() {
	if atomic.AddInt64(&w.workerCount, -1) < 0 {
		atomic.StoreInt64(&w.workerCount, 0)
	} else {
		w.wg.Done()
	}
}

func (w *WorkGroup) Count() int {
	return int(w.workerCount)
}

func (w *WorkGroup) Wait() {
	if atomic.LoadInt64(&w.workerCount) > 0 {
		w.wg.Wait()
		atomic.StoreInt64(&w.workerCount, 0)
	}
}

func (w *WorkGroup) Close() {
	old := atomic.SwapInt64(&w.workerCount, 0)
	if old > 0 {
		w.wg.Add(int(-old))
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 docker-cluster authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build: race

package cluster

import (
	"github.com/fsouza/go-dockerclient"
	"runtime"
	"sync"
	"testing"
)

func TestRoundRobinScheduleIsRaceFree(t *testing.T) {
	const tasks = 8
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(tasks))
	c, err := New(&roundRobin{}, &MapStorage{})
	if err != nil {
		t.Fatalf("unexpected error: %s", err.Error())
	}
	c.Register("url1", nil)
	c.Register("url2", nil)
	opts := docker.CreateContainerOptions{Config: &docker.Config{}}
	var wg sync.WaitGroup
	wg.Add(8)
	for i := 0; i < tasks; i++ {
		go func() {
			defer wg.Done()
			node, err := c.scheduler.Schedule(c, opts, nil)
			if err != nil {
				t.Fatal(err)
			}
			if node.Address != "url1" && node.Address != "url2" {
				t.Errorf("Wrong node. Wanted url1 or url2. Got %q.", node)
			}
		}()
	}
	wg.Wait()
}
<commit_msg>cluster: fix vet errors<commit_after>// Copyright 2014 docker-cluster authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build race

package cluster

import (
	"github.com/fsouza/go-dockerclient"
	"runtime"
	"sync"
	"testing"
)

func TestRoundRobinScheduleIsRaceFree(t *testing.T) {
	const tasks = 8
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(tasks))
	c, err := New(&roundRobin{}, &MapStorage{})
	if err != nil {
		t.Fatalf("unexpected error: %s", err.Error())
	}
	c.Register("url1", nil)
	c.Register("url2", nil)
	opts := docker.CreateContainerOptions{Config: &docker.Config{}}
	var wg sync.WaitGroup
	wg.Add(8)
	for i := 0; i < tasks; i++ {
		go func() {
			defer wg.Done()
			node, err := c.scheduler.Schedule(c, opts, nil)
			if err != nil {
				t.Fatal(err)
			}
			if node.Address != "url1" && node.Address != "url2" {
				t.Errorf("Wrong node. Wanted url1 or url2. Got %q.", node.Address)
			}
		}()
	}
	wg.Wait()
}
<|endoftext|>


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

// Main entry point for Kubernetes listener
package main

import (
	"flag"
	"fmt"
	"github.com/romana/core/common"
	"github.com/romana/core/romana/kubernetes"
)

func main() {
	rootURL := flag.String("rootURL", "", "Romana Root service URL")
	version := flag.Bool("version", false, "Build Information.")
	username := flag.String("username", "", "Username")
	password := flag.String("password", "", "Password")

	flag.Parse()

	if *version {
		fmt.Println(common.BuildInfo())
		return
	}
	cred := common.MakeCredentialFromCliArgs(*username, *password)
	kubernetes.Run(*rootURL, cred)
}
<commit_msg>Add service loop in kubernetes listener main.go<commit_after>// Copyright (c) 2016 Pani Networks
// All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

// Main entry point for Kubernetes listener
package main

import (
	"flag"
	"fmt"
	"github.com/romana/core/common"
	"github.com/romana/core/romana/kubernetes"
)

func main() {
	rootURL := flag.String("rootURL", "", "Romana Root service URL")
	version := flag.Bool("version", false, "Build Information.")
	username := flag.String("username", "", "Username")
	password := flag.String("password", "", "Password")

	flag.Parse()

	if *version {
		fmt.Println(common.BuildInfo())
		return
	}
	cred := common.MakeCredentialFromCliArgs(*username, *password)
	svcInfo, err := kubernetes.Run(*rootURL, cred)
	if err != nil {
		panic(err)
	}
	for {
		msg := <-svcInfo.Channel
		fmt.Println(msg)
	}
}
<|endoftext|>


<commit_before>package kinetic

import (
	"errors"
	"runtime"
	"sync"
	"syscall"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	awsKinesis "github.com/aws/aws-sdk-go/service/kinesis"
	awsKinesisIface "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
	. "github.com/smartystreets/goconvey/convey"
)

const testEndpoint = "http://127.0.0.1:4567"

func CreateAndWaitForStream(client awsKinesisIface.KinesisAPI, name string) {
	client.CreateStream(&awsKinesis.CreateStreamInput{
		StreamName: aws.String(name),
		ShardCount: aws.Int64(1),
	})
	stream := &awsKinesis.DescribeStreamInput{StreamName: aws.String(name), Limit: aws.Int64(1)}
	client.WaitUntilStreamExists(stream)
}

func TestListenerStop(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")
	CreateAndWaitForStream(listener.client, "stream-name")
	listener.ReInit()

	Convey("Given a running listener", t, func() {
		go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
			wg.Done()
		})

		Convey("It should stop listening if sent an interrupt signal", func() {
			listener.interrupts <- syscall.SIGINT
			runtime.Gosched()
			// Let it finish stopping
			time.Sleep(3 * time.Second)

			So(listener.IsListening(), ShouldEqual, false)
		})
	})

	listener.Close()
}

func TestListenerSyncStop(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")
	CreateAndWaitForStream(listener.client, "stream-name")
	listener.ReInit()

	Convey("Given a running listener", t, func() {
		go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
			wg.Done()
		})

		Convey("It should stop listening if sent an interrupt signal", func() {
			err := listener.CloseSync()
			So(err, ShouldBeNil)
			So(listener.IsListening(), ShouldEqual, false)
		})
	})

	listener.Close()
}

func TestListenerError(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")
	CreateAndWaitForStream(listener.client, "stream-name")
	listener.ReInit()

	Convey("Given a running listener", t, func() {
		go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
			wg.Done()
		})

		Convey("It should handle errors successfully", func() {
			listener.errors <- errors.New("All your base are belong to us")

			// Let the error propagate
			<-time.After(3 * time.Second)

			So(listener.getErrCount(), ShouldNotEqual, 0)
			So(listener.IsListening(), ShouldEqual, true)
		})
	})

	listener.Close()
}

func TestListenerMessage(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")
	CreateAndWaitForStream(listener.client, "stream-name")
	listener.ReInit()

	go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
		wg.Done()
	})

	<-time.After(3 * time.Second)

	for _, c := range cases {
		Convey("Given a running listener", t, func() {
			listener.messages <- new(Message).Init(c.message, "test")

			Convey("It should handle messages successfully", func() {
				So(listener.IsListening(), ShouldEqual, true)
				So(listener.Errors(), ShouldNotResemble, nil)
			})
		})
	}

	listener.Close()
}

func TestRetrieveMessage(t *testing.T) {
	listener, _ := new(Listener).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10)
	producer, _ := new(KinesisProducer).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10)

	listener.NewEndpoint(testEndpoint, "your-stream")
	producer.NewEndpoint(testEndpoint, "your-stream")
	CreateAndWaitForStream(listener.client, "your-stream")
	listener.ReInit()
	producer.ReInit()

	time.Sleep(10 * time.Millisecond)
	for _, c := range cases {
		Convey("Given a valid message", t, func() {
			producer.Send(new(Message).Init(c.message, "test"))
			time.Sleep(3 * time.Millisecond)

			Convey("It should be passed on the queue without error", func() {
				msg, err := listener.Retrieve()
				// if err != nil {
				// 	t.Fatalf(err.Error())
				// }
				So(err, ShouldBeNil)

				So(string(msg.Value()), ShouldResemble, string(c.message))
			})
		})
	}

	producer.Close()
	listener.Close()
}

var cases = []struct {
	message []byte
}{
	{
		message: []byte(`{"foo":"bar"}`),
	},
	{
		message: []byte(`{"bar":"baz"}`),
	},
	{
		message: []byte(`{"baz":"qux"}`),
	},
}
<commit_msg>Fix another data race in the tests<commit_after>package kinetic

import (
	"errors"
	"runtime"
	"sync"
	"syscall"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	awsKinesis "github.com/aws/aws-sdk-go/service/kinesis"
	awsKinesisIface "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
	. "github.com/smartystreets/goconvey/convey"
)

const testEndpoint = "http://127.0.0.1:4567"

func CreateAndWaitForStream(client awsKinesisIface.KinesisAPI, name string) {
	client.CreateStream(&awsKinesis.CreateStreamInput{
		StreamName: aws.String(name),
		ShardCount: aws.Int64(1),
	})
	stream := &awsKinesis.DescribeStreamInput{StreamName: aws.String(name), Limit: aws.Int64(1)}
	client.WaitUntilStreamExists(stream)
}

func TestListenerStop(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")
	CreateAndWaitForStream(listener.client, "stream-name")
	listener.ReInit()

	Convey("Given a running listener", t, func() {
		go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
			wg.Done()
		})

		Convey("It should stop listening if sent an interrupt signal", func() {
			listener.interrupts <- syscall.SIGINT
			runtime.Gosched()
			// Let it finish stopping
			time.Sleep(3 * time.Second)

			So(listener.IsListening(), ShouldEqual, false)
		})
	})

	listener.Close()
}

func TestListenerSyncStop(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")
	CreateAndWaitForStream(listener.client, "stream-name")
	listener.ReInit()

	Convey("Given a running listener", t, func() {
		go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
			wg.Done()
		})

		Convey("It should stop listening if sent an interrupt signal", func() {
			err := listener.CloseSync()
			So(err, ShouldBeNil)
			So(listener.IsListening(), ShouldEqual, false)
		})
	})

	listener.Close()
}

func TestListenerError(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")
	CreateAndWaitForStream(listener.client, "stream-name")
	listener.ReInit()

	Convey("Given a running listener", t, func() {
		go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
			wg.Done()
		})

		Convey("It should handle errors successfully", func() {
			listener.errors <- errors.New("All your base are belong to us")

			// Let the error propagate
			<-time.After(3 * time.Second)

			So(listener.getErrCount(), ShouldNotEqual, 0)
			So(listener.IsListening(), ShouldEqual, true)
		})
	})

	listener.Close()
}

func TestListenerMessage(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")
	CreateAndWaitForStream(listener.client, "stream-name")
	listener.ReInit()

	go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
		wg.Done()
	})

	time.Sleep(3 * time.Second)

	for _, c := range cases {
		Convey("Given a running listener", t, func() {
			listener.messages <- new(Message).Init(c.message, "test")

			Convey("It should handle messages successfully", func() {
				So(listener.IsListening(), ShouldEqual, true)
				So(listener.Errors(), ShouldNotResemble, nil)
			})
		})
	}
	time.Sleep(2 * time.Second)
	listener.CloseSync()
}

func TestRetrieveMessage(t *testing.T) {
	listener, _ := new(Listener).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10)
	producer, _ := new(KinesisProducer).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10)

	listener.NewEndpoint(testEndpoint, "your-stream")
	producer.NewEndpoint(testEndpoint, "your-stream")
	CreateAndWaitForStream(listener.client, "your-stream")
	listener.ReInit()
	producer.ReInit()

	time.Sleep(10 * time.Millisecond)
	for _, c := range cases {
		Convey("Given a valid message", t, func() {
			producer.Send(new(Message).Init(c.message, "test"))
			time.Sleep(3 * time.Millisecond)

			Convey("It should be passed on the queue without error", func() {
				msg, err := listener.Retrieve()
				So(err, ShouldBeNil)

				So(string(msg.Value()), ShouldResemble, string(c.message))
			})
		})
	}
	time.Sleep(1 * time.Second)
	producer.CloseSync()
	time.Sleep(1 * time.Second)
	listener.CloseSync()
}

var cases = []struct {
	message []byte
}{
	{
		message: []byte(`{"foo":"bar"}`),
	},
	{
		message: []byte(`{"bar":"baz"}`),
	},
	{
		message: []byte(`{"baz":"qux"}`),
	},
}
<|endoftext|>


<commit_before>package testblas

import (
	"math"
	"testing"

	"github.com/gonum/blas"
)

// throwPanic will throw unexpected panics if true, or will just report them as errors if false
const throwPanic = true

func dTolEqual(a, b float64) bool {
	m := math.Max(math.Abs(a), math.Abs(b))
	if m > 1 {
		a /= m
		b /= m
	}
	if math.Abs(a-b) < 1e-14 {
		return true
	}
	return false
}

func dSliceTolEqual(a, b []float64) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if !dTolEqual(a[i], b[i]) {
			return false
		}
	}
	return true
}

func dStridedSliceTolEqual(n int, a []float64, inca int, b []float64, incb int) bool {
	ia := 0
	ib := 0
	if inca <= 0 {
		ia = -(n - 1) * inca
	}
	if incb <= 0 {
		ib = -(n - 1) * incb
	}
	for i := 0; i < n; i++ {
		if !dTolEqual(a[ia], b[ib]) {
			return false
		}
		ia += inca
		ib += incb
	}
	return true
}

func dSliceEqual(a, b []float64) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if !(a[i] == b[i]) {
			return false
		}
	}
	return true
}

func dCopyTwoTmp(x, xTmp, y, yTmp []float64) {
	if len(x) != len(xTmp) {
		panic("x size mismatch")
	}
	if len(y) != len(yTmp) {
		panic("y size mismatch")
	}
	for i, val := range x {
		xTmp[i] = val
	}
	for i, val := range y {
		yTmp[i] = val
	}
}

// returns true if the function panics
func panics(f func()) (b bool) {
	defer func() {
		err := recover()
		if err != nil {
			b = true
		}
	}()
	f()
	return
}

func testpanics(f func(), name string, t *testing.T) {
	b := panics(f)
	if !b {
		t.Errorf("%v should panic and does not", name)
	}
}

func sliceOfSliceCopy(a [][]float64) [][]float64 {
	n := make([][]float64, len(a))
	for i := range a {
		n[i] = make([]float64, len(a[i]))
		copy(n[i], a[i])
	}
	return n
}

func sliceCopy(a []float64) []float64 {
	n := make([]float64, len(a))
	copy(n, a)
	return n
}

func flatten(a [][]float64) []float64 {
	if len(a) == 0 {
		return nil
	}
	m := len(a)
	n := len(a[0])
	s := make([]float64, m*n)
	for i := 0; i < m; i++ {
		for j := 0; j < n; j++ {
			s[i*n+j] = a[i][j]
		}
	}
	return s
}

func unflatten(a []float64, m, n int) [][]float64 {
	s := make([][]float64, m)
	for i := 0; i < m; i++ {
		s[i] = make([]float64, n)
		for j := 0; j < n; j++ {
			s[i][j] = a[i*n+j]
		}
	}
	return s
}

// flattenTriangular turns the upper or lower triangle of a dense slice of slice
// into a single slice with packed storage.
func flattenTriangular(a [][]float64, ul blas.Uplo) []float64 {
	m := len(a)
	n := len(a[0])
	if m != n {
		panic("must be square")
	}
	aFlat := make([]float64, n*(n+1)/2)
	var count int
	if ul == blas.Upper {
		for i := 0; i < m; i++ {
			for j := i; j < n; j++ {
				aFlat[count] = a[i][j]
				count++
			}
		}
		return aFlat
	}
	for i := 0; i < m; i++ {
		for j := 0; j <= i; j++ {
			aFlat[count] = a[i][j]
			count++
		}
	}
	return aFlat
}

// flattenBanded turns a dense banded slice of slice into the compact banded matrix format
func flattenBanded(a [][]float64, ku, kl int) []float64 {
	m := len(a)
	n := len(a[0])
	if ku < 0 || kl < 0 {
		panic("testblas: negative band length")
	}
	nRows := m
	nCols := (ku + kl + 1)
	aflat := make([]float64, nRows*nCols)
	for i := range aflat {
		aflat[i] = math.NaN()
	}
	// loop over the rows, and then the bands
	// elements in the ith row stay in the ith row
	// order in bands is kept
	for i := 0; i < nRows; i++ {
		min := -kl
		if i-kl < 0 {
			min = -i
		}
		max := ku
		if i+ku >= n {
			max = n - i - 1
		}
		for j := min; j <= max; j++ {
			col := kl + j
			aflat[i*nCols+col] = a[i][i+j]
		}
	}
	return aflat
}

// makeIncremented takes a slice with inc == 1 and makes an incremented version
// and adds extra values on the end
func makeIncremented(x []float64, inc int, extra int) []float64 {
	if inc == 0 {
		panic("zero inc")
	}
	absinc := inc
	if absinc < 0 {
		absinc = -inc
	}
	xcopy := make([]float64, len(x))
	if inc > 0 {
		copy(xcopy, x)
	} else {
		for i := 0; i < len(x); i++ {
			xcopy[i] = x[len(x)-i-1]
		}
	}

	// don't use NaN because it makes comparison hard
	// Do use a weird unique value for easier debugging
	counter := 100.0
	var xnew []float64
	for i, v := range xcopy {
		xnew = append(xnew, v)
		if i != len(x)-1 {
			for j := 0; j < absinc-1; j++ {
				xnew = append(xnew, counter)
				counter++
			}
		}
	}
	for i := 0; i < extra; i++ {
		xnew = append(xnew, counter)
		counter++
	}
	return xnew
}
<commit_msg>remove flatten to pull from master<commit_after>package testblas

import (
	"math"
	"testing"
)

// throwPanic will throw unexpected panics if true, or will just report them as errors if false
const throwPanic = true

func dTolEqual(a, b float64) bool {
	m := math.Max(math.Abs(a), math.Abs(b))
	if m > 1 {
		a /= m
		b /= m
	}
	if math.Abs(a-b) < 1e-14 {
		return true
	}
	return false
}

func dSliceTolEqual(a, b []float64) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if !dTolEqual(a[i], b[i]) {
			return false
		}
	}
	return true
}

func dStridedSliceTolEqual(n int, a []float64, inca int, b []float64, incb int) bool {
	ia := 0
	ib := 0
	if inca <= 0 {
		ia = -(n - 1) * inca
	}
	if incb <= 0 {
		ib = -(n - 1) * incb
	}
	for i := 0; i < n; i++ {
		if !dTolEqual(a[ia], b[ib]) {
			return false
		}
		ia += inca
		ib += incb
	}
	return true
}

func dSliceEqual(a, b []float64) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if !(a[i] == b[i]) {
			return false
		}
	}
	return true
}

func dCopyTwoTmp(x, xTmp, y, yTmp []float64) {
	if len(x) != len(xTmp) {
		panic("x size mismatch")
	}
	if len(y) != len(yTmp) {
		panic("y size mismatch")
	}
	for i, val := range x {
		xTmp[i] = val
	}
	for i, val := range y {
		yTmp[i] = val
	}
}

// returns true if the function panics
func panics(f func()) (b bool) {
	defer func() {
		err := recover()
		if err != nil {
			b = true
		}
	}()
	f()
	return
}

func testpanics(f func(), name string, t *testing.T) {
	b := panics(f)
	if !b {
		t.Errorf("%v should panic and does not", name)
	}
}

func sliceOfSliceCopy(a [][]float64) [][]float64 {
	n := make([][]float64, len(a))
	for i := range a {
		n[i] = make([]float64, len(a[i]))
		copy(n[i], a[i])
	}
	return n
}

func sliceCopy(a []float64) []float64 {
	n := make([]float64, len(a))
	copy(n, a)
	return n
}

func flatten(a [][]float64) []float64 {
	if len(a) == 0 {
		return nil
	}
	m := len(a)
	n := len(a[0])
	s := make([]float64, m*n)
	for i := 0; i < m; i++ {
		for j := 0; j < n; j++ {
			s[i*n+j] = a[i][j]
		}
	}
	return s
}

func unflatten(a []float64, m, n int) [][]float64 {
	s := make([][]float64, m)
	for i := 0; i < m; i++ {
		s[i] = make([]float64, n)
		for j := 0; j < n; j++ {
			s[i][j] = a[i*n+j]
		}
	}
	return s
}

// flattenBanded turns a dense banded slice of slice into the compact banded matrix format
func flattenBanded(a [][]float64, ku, kl int) []float64 {
	m := len(a)
	n := len(a[0])
	if ku < 0 || kl < 0 {
		panic("testblas: negative band length")
	}
	nRows := m
	nCols := (ku + kl + 1)
	aflat := make([]float64, nRows*nCols)
	for i := range aflat {
		aflat[i] = math.NaN()
	}
	// loop over the rows, and then the bands
	// elements in the ith row stay in the ith row
	// order in bands is kept
	for i := 0; i < nRows; i++ {
		min := -kl
		if i-kl < 0 {
			min = -i
		}
		max := ku
		if i+ku >= n {
			max = n - i - 1
		}
		for j := min; j <= max; j++ {
			col := kl + j
			aflat[i*nCols+col] = a[i][i+j]
		}
	}
	return aflat
}

// makeIncremented takes a slice with inc == 1 and makes an incremented version
// and adds extra values on the end
func makeIncremented(x []float64, inc int, extra int) []float64 {
	if inc == 0 {
		panic("zero inc")
	}
	absinc := inc
	if absinc < 0 {
		absinc = -inc
	}
	xcopy := make([]float64, len(x))
	if inc > 0 {
		copy(xcopy, x)
	} else {
		for i := 0; i < len(x); i++ {
			xcopy[i] = x[len(x)-i-1]
		}
	}

	// don't use NaN because it makes comparison hard
	// Do use a weird unique value for easier debugging
	counter := 100.0
	var xnew []float64
	for i, v := range xcopy {
		xnew = append(xnew, v)
		if i != len(x)-1 {
			for j := 0; j < absinc-1; j++ {
				xnew = append(xnew, counter)
				counter++
			}
		}
	}
	for i := 0; i < extra; i++ {
		xnew = append(xnew, counter)
		counter++
	}
	return xnew
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/base64"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"time"

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

type sshChannel struct {
	ssh.Channel
	name string
}

func sendGlobalRequset(conn ssh.Conn, name string, wantReply bool, payload []byte) {
	time.Sleep(500 * time.Millisecond)
	accepted, response, err := conn.SendRequest(name, wantReply, payload)
	fmt.Printf(">global request %v\n  %#v\n  %#v\n", name, accepted, response)
	if err != nil {
		panic(err)
	}
}

func openChannel(conn ssh.Conn, name string, data []byte, success bool) *sshChannel {
	time.Sleep(500 * time.Millisecond)
	channel, requests, err := conn.OpenChannel(name, data)
	fmt.Printf(">channel %v\n", name)
	if success {
		if err != nil {
			panic(err)
		}
	} else {
		if err == nil {
			panic(err)
		}
		return nil
	}
	if err != nil {
		if success {
			panic(err)
		} else {
			return nil
		}
	}
	go func() {
		for request := range requests {
			fmt.Printf("<channel %v request\n  %#v\n  %#v\n  %#v\n", name, request.Type, request.WantReply, request.Payload)
			if request.WantReply {
				panic("WantReply")
			}
		}
		fmt.Printf("<close channel %v requests\n", name)
	}()
	go func() {
		scanner := bufio.NewScanner(channel)
		for scanner.Scan() {
			fmt.Printf("<channel %v stdout\n  %#v\n", name, scanner.Text())
		}
		fmt.Printf("<close channel %v stdout\n", name)
		if err := scanner.Err(); err != nil {
			panic(err)
		}
	}()
	go func() {
		scanner := bufio.NewScanner(channel.Stderr())
		for scanner.Scan() {
			fmt.Printf("<channel %v stderr\n  %#v\n", name, scanner.Text())
		}
		fmt.Printf("<close channel %v stderr\n", name)
		if err := scanner.Err(); err != nil {
			panic(err)
		}
	}()
	return &sshChannel{channel, name}
}

func (channel *sshChannel) sendRequset(name string, wantReply bool, payload []byte) {
	time.Sleep(500 * time.Millisecond)
	accepted, err := channel.SendRequest(name, wantReply, payload)
	fmt.Printf(">channel %v request %v\n  %#v\n", channel.name, name, accepted)
	if err != nil {
		panic(err)
	}
}

func (channel *sshChannel) write(data string, close bool) {
	time.Sleep(500 * time.Millisecond)
	_, err := channel.Write([]byte(data))
	fmt.Printf(">channel %v data\n  %#v\n", channel.name, data)
	if err != nil {
		panic(err)
	}
	if close {
		err = channel.CloseWrite()
		fmt.Printf(">closewrite channel %v\n", channel.name)
		if err != nil {
			panic(err)
		}
	}
}

func (channel *sshChannel) close() {
	time.Sleep(500 * time.Millisecond)
	err := channel.Close()
	fmt.Printf(">close channel %v\n", channel.name)
	if err != nil {
		panic(err)
	}
}

func main() {
	addr := flag.String("addr", "127.0.0.1:22", "")
	clientVersion := flag.String("client_version", "SSH-2.0-sshesame", "")
	user := flag.String("user", "root", "")
	password := flag.String("password", "", "")
	key := flag.String("key", "", "")
	flag.Parse()

	config := &ssh.ClientConfig{
		User: *user,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			fmt.Printf("<host key\n  %#v\n", key)
			return nil
		},
		BannerCallback: func(message string) error {
			fmt.Printf("<banner\n  %#v\n", message)
			return nil
		},
		ClientVersion: *clientVersion,
	}
	if *password != "" {
		config.Auth = append(config.Auth, ssh.Password(*password))
	}
	if *key != "" {
		keyBytes, err := ioutil.ReadFile(*key)
		if err != nil {
			panic(err)
		}
		signer, err := ssh.ParsePrivateKey(keyBytes)
		if err != nil {
			panic(err)
		}
		config.Auth = append(config.Auth, ssh.PublicKeys(signer))
	}

	conn, err := net.Dial("tcp", *addr)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	sshClientConn, channels, requests, err := ssh.NewClientConn(conn, *addr, config)
	if err != nil {
		panic(err)
	}
	defer sshClientConn.Close()
	fmt.Printf(">connect\n  %#v\n", string(sshClientConn.ServerVersion()))

	go func() {
		for request := range requests {
			fmt.Printf("<global request\n  %#v\n  %#v\n  %#v\n", request.Type, request.WantReply, request.Payload)
			if request.WantReply {
				panic("WantReply")
			}
		}
		fmt.Printf("<close global requests\n")
	}()

	go func() {
		for channel := range channels {
			panic(channel)
		}
		fmt.Printf("<close channels\n")
	}()

	sendGlobalRequset(sshClientConn, "nope", true, []byte("nope"))

	// Causes the connection to close
	// sendGlobalRequset(sshClientConn, "tcpip-forward", true, []byte("nope"))

	sendGlobalRequset(sshClientConn, "tcpip-forward", true, ssh.Marshal(struct {
		string
		uint32
	}{"127.0.0.1", 1234}))

	sendGlobalRequset(sshClientConn, "tcpip-forward", true, ssh.Marshal(struct {
		string
		uint32
	}{"127.0.0.1", 0}))

	sendGlobalRequset(sshClientConn, "cancel-tcpip-forward", true, ssh.Marshal(struct {
		string
		uint32
	}{"127.0.0.1", 1234}))

	openChannel(sshClientConn, "nope", []byte("nope"), false)

	// Causes the connection to close
	// openChannel(sshClientConn, "direct-tcpip", []byte("nope"), false)

	tcpipChannel := openChannel(sshClientConn, "direct-tcpip", ssh.Marshal(struct {
		address           string
		port              uint32
		originatorAddress string
		originatorPort    uint32
	}{"github.com", 80, "127.0.0.1", 8080}), true)

	tcpipChannel.sendRequset("shell", true, nil)

	tcpipChannel.write("GET / HTTP/1.1\r\nHost: github.com\r\n\r\n", false)

	tcpipChannel.close()

	sessionChannel := openChannel(sshClientConn, "session", nil, true)

	// Blocks indefinitely
	// sessionChannel.write("foo")

	sessionChannel.sendRequset("nope", true, []byte("nope"))

	// Causes the connection to close
	// sessionChannel.sendRequset("exec", true, []byte("nope"))

	sessionChannel.sendRequset("exec", true, ssh.Marshal(struct {
		command string
	}{"true"}))

	sessionChannel = openChannel(sshClientConn, "session", nil, true)

	sessionChannel.sendRequset("shell", true, nil)

	sessionChannel.write("true\nfalse\nuname", true)

	sessionChannel = openChannel(sshClientConn, "session", nil, true)

	terminalModes, err := base64.RawStdEncoding.DecodeString("gQAAJYCAAAAlgAEAAAADAgAAABwDAAAAfwQAAAAVBQAAAAQGAAAA/wcAAAD/CAAAABEJAAAAEwoAAAAaCwAAABkMAAAAEg0AAAAXDgAAABYRAAAAFBIAAAAPHgAAAAAfAAAAACAAAAAAIQAAAAAiAAAAACMAAAAAJAAAAAEmAAAAACcAAAABKAAAAAApAAAAASoAAAABMgAAAAEzAAAAATUAAAABNgAAAAE3AAAAADgAAAAAOQAAAAA6AAAAADsAAAABPAAAAAE9AAAAAT4AAAAARgAAAAFIAAAAAUkAAAAASgAAAABLAAAAAFoAAAABWwAAAAFcAAAAAF0AAAAAAA")
	if err != nil {
		panic(err)
	}
	sessionChannel.sendRequset("pty-req", true, ssh.Marshal(struct {
		Term                                   string
		Width, Height, PixelWidth, PixelHeight uint32
		Modes                                  string
	}{"xterm-256color", 120, 80, 0, 0, string(terminalModes)}))

	sessionChannel.sendRequset("shell", true, nil)

	sessionChannel.sendRequset("pty-req", true, ssh.Marshal(struct {
		Term                                   string
		Width, Height, PixelWidth, PixelHeight uint32
		Modes                                  string
	}{"xterm-256color", 120, 80, 0, 0, string(terminalModes)}))

	sessionChannel.sendRequset("exec", true, ssh.Marshal(struct {
		command string
	}{"true"}))

	sessionChannel.write("true\nfalse\nuname\n", true)

	sendGlobalRequset(sshClientConn, "no-more-sessions@openssh.com", false, nil)

	openChannel(sshClientConn, "session", nil, false)

	time.Sleep(5 * time.Second)
}
<commit_msg>Test client improvements<commit_after>package main

import (
	"bufio"
	"encoding/base64"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"time"

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

type sshChannel struct {
	ssh.Channel
	name string
}

func sendGlobalRequset(conn ssh.Conn, name string, wantReply bool, payload []byte) {
	time.Sleep(500 * time.Millisecond)
	accepted, response, err := conn.SendRequest(name, wantReply, payload)
	fmt.Printf(">global request %v\n  %#v\n  %#v\n", name, accepted, response)
	if err != nil {
		panic(err)
	}
}

func openChannel(conn ssh.Conn, name string, data []byte, success bool) *sshChannel {
	time.Sleep(500 * time.Millisecond)
	channel, requests, err := conn.OpenChannel(name, data)
	fmt.Printf(">channel %v\n  %#v\n", name, err)
	if success {
		if err != nil {
			panic(err)
		}
	} else {
		if _, ok := err.(*ssh.OpenChannelError); !ok {
			panic(err)
		}
		return nil
	}
	if err != nil {
		if success {
			panic(err)
		} else {
			return nil
		}
	}
	go func() {
		for request := range requests {
			fmt.Printf("<channel %v request\n  %#v\n  %#v\n  %#v\n", name, request.Type, request.WantReply, request.Payload)
			if request.WantReply {
				panic("WantReply")
			}
		}
		fmt.Printf("<close channel %v requests\n", name)
	}()
	go func() {
		scanner := bufio.NewScanner(channel)
		for scanner.Scan() {
			fmt.Printf("<channel %v stdout\n  %#v\n", name, scanner.Text())
		}
		fmt.Printf("<close channel %v stdout\n", name)
		if err := scanner.Err(); err != nil {
			panic(err)
		}
	}()
	go func() {
		scanner := bufio.NewScanner(channel.Stderr())
		for scanner.Scan() {
			fmt.Printf("<channel %v stderr\n  %#v\n", name, scanner.Text())
		}
		fmt.Printf("<close channel %v stderr\n", name)
		if err := scanner.Err(); err != nil {
			panic(err)
		}
	}()
	return &sshChannel{channel, name}
}

func (channel *sshChannel) sendRequset(name string, wantReply bool, payload []byte) {
	time.Sleep(500 * time.Millisecond)
	accepted, err := channel.SendRequest(name, wantReply, payload)
	fmt.Printf(">channel %v request %v\n  %#v\n", channel.name, name, accepted)
	if err != nil {
		panic(err)
	}
}

func (channel *sshChannel) write(data string, close bool) {
	time.Sleep(500 * time.Millisecond)
	_, err := channel.Write([]byte(data))
	fmt.Printf(">channel %v data\n  %#v\n", channel.name, data)
	if err != nil {
		panic(err)
	}
	if close {
		err = channel.CloseWrite()
		fmt.Printf(">closewrite channel %v\n", channel.name)
		if err != nil {
			panic(err)
		}
	}
}

func main() {
	addr := flag.String("addr", "127.0.0.1:22", "")
	clientVersion := flag.String("client_version", "SSH-2.0-sshesame", "")
	user := flag.String("user", "root", "")
	password := flag.String("password", "", "")
	key := flag.String("key", "", "")
	flag.Parse()

	config := &ssh.ClientConfig{
		User: *user,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			fmt.Printf("<host key\n  %#v\n", key)
			return nil
		},
		BannerCallback: func(message string) error {
			fmt.Printf("<banner\n  %#v\n", message)
			return nil
		},
		ClientVersion: *clientVersion,
	}
	if *password != "" {
		config.Auth = append(config.Auth, ssh.Password(*password))
	}
	if *key != "" {
		keyBytes, err := ioutil.ReadFile(*key)
		if err != nil {
			panic(err)
		}
		signer, err := ssh.ParsePrivateKey(keyBytes)
		if err != nil {
			panic(err)
		}
		config.Auth = append(config.Auth, ssh.PublicKeys(signer))
	}

	conn, err := net.Dial("tcp", *addr)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	sshClientConn, channels, requests, err := ssh.NewClientConn(conn, *addr, config)
	if err != nil {
		panic(err)
	}
	defer sshClientConn.Close()
	fmt.Printf(">connect\n  %#v\n", string(sshClientConn.ServerVersion()))

	go func() {
		for request := range requests {
			fmt.Printf("<global request\n  %#v\n  %#v\n  %#v\n", request.Type, request.WantReply, request.Payload)
			if request.WantReply {
				panic("WantReply")
			}
		}
		fmt.Printf("<close global requests\n")
	}()

	go func() {
		for channel := range channels {
			panic(channel)
		}
		fmt.Printf("<close channels\n")
	}()

	sendGlobalRequset(sshClientConn, "nope", true, []byte("nope"))

	// Causes the connection to close (data expected, nil sent)
	// sendGlobalRequset(sshClientConn, "tcpip-forward", true, nil)

	// Causes the connection to close (data expected, invalid data sent)
	// sendGlobalRequset(sshClientConn, "tcpip-forward", true, []byte("nope"))

	sendGlobalRequset(sshClientConn, "tcpip-forward", true, ssh.Marshal(struct {
		string
		uint32
	}{"127.0.0.1", 1234}))

	sendGlobalRequset(sshClientConn, "tcpip-forward", true, ssh.Marshal(struct {
		string
		uint32
	}{"127.0.0.1", 0}))

	sendGlobalRequset(sshClientConn, "cancel-tcpip-forward", true, ssh.Marshal(struct {
		string
		uint32
	}{"127.0.0.1", 1234}))

	openChannel(sshClientConn, "nope", []byte("nope"), false)

	// Causes the connection to close (data expected, nil sent)
	// openChannel(sshClientConn, "direct-tcpip", nil, false)

	// Causes the connection to close (data expected, invalid data sent)
	// openChannel(sshClientConn, "direct-tcpip", []byte("nope"), false)

	tcpipChannel := openChannel(sshClientConn, "direct-tcpip", ssh.Marshal(struct {
		address           string
		port              uint32
		originatorAddress string
		originatorPort    uint32
	}{"github.com", 80, "127.0.0.1", 8080}), true)
	tcpipChannel.sendRequset("shell", true, nil)
	tcpipChannel.sendRequset("shell", true, []byte("nope"))
	tcpipChannel.sendRequset("exec", true, ssh.Marshal(struct {
		command string
	}{"true"}))
	tcpipChannel.sendRequset("exec", true, nil)
	tcpipChannel.sendRequset("exec", true, []byte("nope"))
	tcpipChannel.write("GET / HTTP/1.1\r\nHost: github.com\r\n\r\n", true)

	// Causes the connection to close (nil expected, data sent)
	// openChannel(sshClientConn, "session", []byte("nope"), false)

	sessionChannel := openChannel(sshClientConn, "session", nil, true)
	// Blocks indefinitely
	// sessionChannel.write("foo")
	sessionChannel.sendRequset("nope", true, []byte("nope"))
	// Causes the connection to close (data expected, nil sent)
	// sessionChannel.sendRequset("exec", true, nil)
	// Causes the connection to close (data expected, invalid data sent)
	// sessionChannel.sendRequset("exec", true, []byte("nope"))
	// Causes the connection to close (nil expected, data sent)
	// sessionChannel.sendRequset("shell", true, []byte("nope"))
	sessionChannel.sendRequset("exec", true, ssh.Marshal(struct {
		command string
	}{"true"}))

	sessionChannel = openChannel(sshClientConn, "session", nil, true)
	sessionChannel.sendRequset("shell", true, nil)
	sessionChannel.write("true\nfalse\nuname\n\x04", false)

	sessionChannel = openChannel(sshClientConn, "session", nil, true)
	sessionChannel.sendRequset("shell", true, nil)
	sessionChannel.write("true\nfalse\nuname\n", true)

	sessionChannel = openChannel(sshClientConn, "session", nil, true)
	terminalModes, err := base64.RawStdEncoding.DecodeString("gQAAJYCAAAAlgAEAAAADAgAAABwDAAAAfwQAAAAVBQAAAAQGAAAA/wcAAAD/CAAAABEJAAAAEwoAAAAaCwAAABkMAAAAEg0AAAAXDgAAABYRAAAAFBIAAAAPHgAAAAAfAAAAACAAAAAAIQAAAAAiAAAAACMAAAAAJAAAAAEmAAAAACcAAAABKAAAAAApAAAAASoAAAABMgAAAAEzAAAAATUAAAABNgAAAAE3AAAAADgAAAAAOQAAAAA6AAAAADsAAAABPAAAAAE9AAAAAT4AAAAARgAAAAFIAAAAAUkAAAAASgAAAABLAAAAAFoAAAABWwAAAAFcAAAAAF0AAAAAAA")
	if err != nil {
		panic(err)
	}
	sessionChannel.sendRequset("pty-req", true, ssh.Marshal(struct {
		Term                                   string
		Width, Height, PixelWidth, PixelHeight uint32
		Modes                                  string
	}{"xterm-256color", 120, 80, 0, 0, string(terminalModes)}))
	sessionChannel.sendRequset("shell", true, nil)
	sessionChannel.sendRequset("pty-req", true, ssh.Marshal(struct {
		Term                                   string
		Width, Height, PixelWidth, PixelHeight uint32
		Modes                                  string
	}{"xterm-256color", 120, 80, 0, 0, string(terminalModes)}))
	sessionChannel.sendRequset("exec", true, ssh.Marshal(struct {
		command string
	}{"true"}))
	sessionChannel.write("true\rfalse\runame\r\x04", false)

	sessionChannel = openChannel(sshClientConn, "session", nil, true)
	sessionChannel.sendRequset("pty-req", true, ssh.Marshal(struct {
		Term                                   string
		Width, Height, PixelWidth, PixelHeight uint32
		Modes                                  string
	}{"xterm-256color", 120, 80, 0, 0, string(terminalModes)}))
	sessionChannel.sendRequset("shell", true, nil)
	sessionChannel.write("true\rfalse\runame\r", true)

	sendGlobalRequset(sshClientConn, "no-more-sessions@openssh.com", false, nil)

	// Causes the connection to close
	// openChannel(sshClientConn, "session", nil, false)

	time.Sleep(5 * time.Second)
}
<|endoftext|>