<commit_before>// SPDX-License-Identifier: MIT

package server

import (
	"fmt"
	"net/http"

	"github.com/issue9/logs/v3"
)

const exited status = 0

type (
	// HandlerFunc 路由项处理函数原型
	//
	// 如果返回 nil，表示未出现任何错误，可以继续后续操作，
	// 非 nil，表示需要中断执行并向用户输出返回的对象。
	HandlerFunc func(*Context) Responser

	// Responser 表示向客户端输出的对象
	Responser interface {
		// Status 状态码
		Status() int

		// Headers 输出的报头
		Headers() map[string]string

		// Body 输出到 body 部分的对象
		//
		// 该对象最终经由 serialization.MarshalFunc 转换成文本输出。
		Body() interface{}
	}

	status int

	object struct {
		status  int
		headers map[string]string
		body    interface{}
	}
)

// FileServer 返回以当前模块作为文件系统的静态文件服务
func (m *Module) FileServer(name, index string) HandlerFunc {
	return m.Server().FileServer(http.FS(m), name, index)
}

func (ctx *Context) renderResponser(resp Responser) {
	if resp == nil || resp == exited {
		return
	}

	if err := ctx.Marshal(resp.Status(), resp.Body(), resp.Headers()); err != nil {
		ctx.Server().Logs().Error(err)
	}
}

func (s status) Status() int { return int(s) }

func (s status) Headers() map[string]string { return nil }

func (s status) Body() interface{} { return nil }

func (o *object) Status() int { return o.status }

func (o *object) Headers() map[string]string { return o.headers }

func (o *object) Body() interface{} { return o.body }

// Error 输出日志到 ERROR 通道并向用户输出指定状态码的页面
//
// NOTE:应该在出错的地方直接调用 Error，而不是将 Error 嵌套在另外的函数里，
// 否则出错信息的位置信息将不准确。
func (ctx *Context) Error(status int, v ...interface{}) Responser {
	ctx.Log(logs.LevelError, 2, v...)
	return Status(status)
}

// Errorf 输出日志到 ERROR 通道并向用户输出指定状态码的页面
func (ctx *Context) Errorf(status int, format string, v ...interface{}) Responser {
	ctx.Logf(logs.LevelError, 2, format, v...)
	return Status(status)
}

// Critical 输出日志到 CRITICAL 通道并向用户输出指定状态码的页面
func (ctx *Context) Critical(status int, v ...interface{}) Responser {
	ctx.Log(logs.LevelCritical, 2, v...)
	return Status(status)
}

// Criticalf 输出日志到 CRITICAL 通道并向用户输出指定状态码的页面
func (ctx *Context) Criticalf(status int, format string, v ...interface{}) Responser {
	ctx.Logf(logs.LevelCritical, 2, format, v...)
	return Status(status)
}

func Object(status int, body interface{}, headers map[string]string) Responser {
	return &object{
		status:  status,
		headers: headers,
		body:    body,
	}
}

// Status 仅包含状态码的 Responser
func Status(statusCode int) Responser {
	if statusCode < 100 || statusCode >= 600 {
		panic(fmt.Sprintf("无效的状态码 %d", statusCode))
	}
	return status(statusCode)
}

// Exit 不再执行后续操作退出当前请求
//
// 与其它返回的区别在于，Exit 表示已经向客户端输出相关内容，仅作退出。
func Exit() Responser { return exited }

// Result 返回 Result 实例
//
// 如果找不到 code 对应的错误信息，则会直接 panic。
func (ctx *Context) Result(code string, fields ResultFields) Responser {
	rslt := ctx.Server().Result(ctx.LocalePrinter, code, fields)
	return Object(rslt.Status(), rslt, nil)
}

// Redirect 重定向至新的 URL
func (ctx *Context) Redirect(status int, url string) Responser {
	http.Redirect(ctx.Response, ctx.Request, url, status)
	return Exit()
}
<commit_msg>refactor(server): 修正无效的类型转换<commit_after>// SPDX-License-Identifier: MIT

package server

import (
	"fmt"
	"net/http"

	"github.com/issue9/logs/v3"
)

const exited status = 0

type (
	// HandlerFunc 路由项处理函数原型
	//
	// 如果返回 nil，表示未出现任何错误，可以继续后续操作，
	// 非 nil，表示需要中断执行并向用户输出返回的对象。
	HandlerFunc func(*Context) Responser

	// Responser 表示向客户端输出的对象
	Responser interface {
		// Status 状态码
		Status() int

		// Headers 输出的报头
		Headers() map[string]string

		// Body 输出到 body 部分的对象
		//
		// 该对象最终经由 serialization.MarshalFunc 转换成文本输出。
		Body() interface{}
	}

	status int

	object struct {
		status  int
		headers map[string]string
		body    interface{}
	}
)

// FileServer 返回以当前模块作为文件系统的静态文件服务
func (m *Module) FileServer(name, index string) HandlerFunc {
	return m.Server().FileServer(m, name, index)
}

func (ctx *Context) renderResponser(resp Responser) {
	if resp == nil || resp == exited {
		return
	}

	if err := ctx.Marshal(resp.Status(), resp.Body(), resp.Headers()); err != nil {
		ctx.Server().Logs().Error(err)
	}
}

func (s status) Status() int { return int(s) }

func (s status) Headers() map[string]string { return nil }

func (s status) Body() interface{} { return nil }

func (o *object) Status() int { return o.status }

func (o *object) Headers() map[string]string { return o.headers }

func (o *object) Body() interface{} { return o.body }

// Error 输出日志到 ERROR 通道并向用户输出指定状态码的页面
//
// NOTE:应该在出错的地方直接调用 Error，而不是将 Error 嵌套在另外的函数里，
// 否则出错信息的位置信息将不准确。
func (ctx *Context) Error(status int, v ...interface{}) Responser {
	ctx.Log(logs.LevelError, 2, v...)
	return Status(status)
}

// Errorf 输出日志到 ERROR 通道并向用户输出指定状态码的页面
func (ctx *Context) Errorf(status int, format string, v ...interface{}) Responser {
	ctx.Logf(logs.LevelError, 2, format, v...)
	return Status(status)
}

// Critical 输出日志到 CRITICAL 通道并向用户输出指定状态码的页面
func (ctx *Context) Critical(status int, v ...interface{}) Responser {
	ctx.Log(logs.LevelCritical, 2, v...)
	return Status(status)
}

// Criticalf 输出日志到 CRITICAL 通道并向用户输出指定状态码的页面
func (ctx *Context) Criticalf(status int, format string, v ...interface{}) Responser {
	ctx.Logf(logs.LevelCritical, 2, format, v...)
	return Status(status)
}

func Object(status int, body interface{}, headers map[string]string) Responser {
	return &object{
		status:  status,
		headers: headers,
		body:    body,
	}
}

// Status 仅包含状态码的 Responser
func Status(statusCode int) Responser {
	if statusCode < 100 || statusCode >= 600 {
		panic(fmt.Sprintf("无效的状态码 %d", statusCode))
	}
	return status(statusCode)
}

// Exit 不再执行后续操作退出当前请求
//
// 与其它返回的区别在于，Exit 表示已经向客户端输出相关内容，仅作退出。
func Exit() Responser { return exited }

// Result 返回 Result 实例
//
// 如果找不到 code 对应的错误信息，则会直接 panic。
func (ctx *Context) Result(code string, fields ResultFields) Responser {
	rslt := ctx.Server().Result(ctx.LocalePrinter, code, fields)
	return Object(rslt.Status(), rslt, nil)
}

// Redirect 重定向至新的 URL
func (ctx *Context) Redirect(status int, url string) Responser {
	http.Redirect(ctx.Response, ctx.Request, url, status)
	return Exit()
}
<|endoftext|>


<commit_before>package service

import (
	"github.com/m110/cort/rpc"
)

type Service struct {
	name   string
	server *rpc.Server
}

func NewService(name string, uri string) *Service {
	service := &Service{
		name:   name,
		server: rpc.NewServer(uri),
	}

	return service
}

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

	err = s.server.Start()
	if err != nil {
		return err
	}

	err = s.Register()
	if err != nil {
		return err
	}

	return nil
}

func (s *Service) Stop() error {
	var err error

	err = s.Deregister()
	if err != nil {
		return err
	}

	err = s.server.Stop()
	if err != nil {
		return err
	}

	return nil
}

func (s *Service) Register() error {
	// TODO Register the service to be discovered by other services
	return nil
}

func (s *Service) Deregister() error {
	// TODO Deregister the service
	return nil
}
<commit_msg>Service: Register before the server loop.<commit_after>package service

import (
	"github.com/m110/cort/rpc"
)

type Service struct {
	name   string
	server *rpc.Server
}

func NewService(name string, uri string) *Service {
	service := &Service{
		name:   name,
		server: rpc.NewServer(uri),
	}

	return service
}

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

	err = s.Register()
	if err != nil {
		return err
	}

	err = s.server.Start()
	if err != nil {
		return err
	}

	return nil
}

func (s *Service) Stop() error {
	var err error

	err = s.Deregister()
	if err != nil {
		return err
	}

	err = s.server.Stop()
	if err != nil {
		return err
	}

	return nil
}

func (s *Service) Register() error {
	// TODO Register the service to be discovered by other services
	return nil
}

func (s *Service) Deregister() error {
	// TODO Deregister the service
	return nil
}
<|endoftext|>


<commit_before>package endly

import (
	"fmt"
	"github.com/viant/toolbox"
	"strings"
	"time"
)

const ProcessServiceId = "process"

type ProcessStartRequest struct {
	Name          string
	Target        *Resource
	Options       *ExecutionOptions
	SystemService bool
	Directory     string
	Command       string
	Arguments     []string
}

type ProcessInfo struct {
	Name      string
	Pid       int
	Command   string
	Arguments []string
	Stdin     string
	Stdout    string
}

type ProcessStatusRequest struct {
	Target  *Resource
	Command string
}

type ProcessStopRequest struct {
	Target *Resource
	Pid    int
}

type processService struct {
	*AbstractService
}

func (s *processService) Run(context *Context, request interface{}) *ServiceResponse {
	var response = &ServiceResponse{Status: "ok"}
	var err error
	switch actualRequest := request.(type) {
	case *ProcessStartRequest:
		response.Response, err = s.startProcess(context, actualRequest)
		if err != nil {
			response.Error = fmt.Sprintf("Failed to start process: %v, %v", actualRequest.Name, err)
		}
	}
	if response.Error != "" {
		response.Status = "err"
	}
	return response
}

func (s *processService) checkProcess(context *Context, request *ProcessStatusRequest) ([]*ProcessInfo, error) {

	commandResponse, err := context.Execute(request.Target, &ManagedCommand{
		Executions: []*Execution{
			{
				Command: "ps -ev | grep " + request.Command,
			},
		},
	})

	if err != nil {
		return nil, err
	}
	var result = make([]*ProcessInfo, 0)

	for _, line := range strings.Split(commandResponse.Stdout(), "\r\n") {
		if strings.Contains(line, "grep") {
			continue
		}
		pid, ok := ExtractColumn(line, 0)
		if !ok {
			continue
		}

		argumentsIndex := strings.Index(line, request.Command)
		var arguments []string
		if argumentsIndex != -1 {
			args := strings.Trim(line[argumentsIndex+len(request.Command)+1:], " &\t")
			arguments = strings.Split(args, " ")
		}
		info := &ProcessInfo{
			Pid:       toolbox.AsInt(pid),
			Command:   request.Command,
			Arguments: arguments,
			Stdout:    line,
		}
		result = append(result, info)
	}
	return result, nil
}

func (s *processService) stopProcess(context *Context, request *ProcessStopRequest) (*CommandInfo, error) {
	commandResult, err := context.Execute(request.Target, &ManagedCommand{
		Executions: []*Execution{
			{
				Command: fmt.Sprintf("kill -9 %v", request.Pid),
			},
		},
	})
	if err != nil {
		return nil, err
	}
	return commandResult, err
}

func indexProcesses(processes ...*ProcessInfo) map[int]*ProcessInfo {
	var result = make(map[int]*ProcessInfo)
	for _, process := range processes {
		result[process.Pid] = process
	}
	return result
}

func (s *processService) startProcess(context *Context, request *ProcessStartRequest) (*ProcessInfo, error) {
	origProcesses, err := s.checkProcess(context, &ProcessStatusRequest{
		Target:  request.Target,
		Command: request.Command,
	})
	if err != nil {
		return nil, err
	}
	for _, process := range origProcesses {

		if strings.Join(process.Arguments, " ") == strings.Join(request.Arguments, " ") {
			_, err := s.stopProcess(context, &ProcessStopRequest{
				Pid:    process.Pid,
				Target: request.Target,
			})
			if err != nil {
				return nil, err
			}
		}
	}
	changeDirCommand := fmt.Sprintf("cd %v ", request.Directory)
	startCommand := fmt.Sprintf("nohup %v %v &", request.Command, strings.Join(request.Arguments, " "))
	commandInfo, err := context.Execute(request.Target, &ManagedCommand{
		Options: request.Options,
		Executions: []*Execution{
			{
				Command: changeDirCommand,
			},
			{
				Command: startCommand,
			},
		},
	})
	if err != nil {
		return nil, err
	}
	time.Sleep(time.Second)
	newProcesses, err := s.checkProcess(context, &ProcessStatusRequest{
		Target:  request.Target,
		Command: request.Command,
	})
	if err != nil {
		return nil, err
	}

	if len(newProcesses) == 1 {
		return newProcesses[0], nil
	}
	existingProcesses := indexProcesses(origProcesses...)
	var result *ProcessInfo
	for _, candidate := range newProcesses {

		if _, has := existingProcesses[candidate.Pid]; !has {
			result = candidate
			break
		}
	}
	if result == nil {
		return nil, fmt.Errorf("Failed to get info about prorcess %v", request.Command)
	}
	result.Stdout = commandInfo.Stdout()
	result.Stdin = fmt.Sprintf("%v && %v", changeDirCommand, startCommand)
	return result, nil
}

func (s *processService) NewRequest(name string) (interface{}, error) {
	switch name {
	case "start":
		return &ProcessStartRequest{}, nil
	case "check":
		return &ProcessStatusRequest{}, nil
	case "stop":
		return &ProcessStopRequest{}, nil

	}
	return nil, fmt.Errorf("Unsupported name: %v", name)
}

func NewProcessService() Service {
	var result = &processService{
		AbstractService: NewAbstractService(ProcessServiceId),
	}
	result.AbstractService.Service = result
	return result
}
<commit_msg>added stop process<commit_after>package endly

import (
	"fmt"
	"github.com/viant/toolbox"
	"strings"
	"time"
	"github.com/viant/endly/common"
)

const ProcessServiceId = "process"
const processesKey = "pid"
type ProcessStartRequest struct {
	Name          string
	Target        *Resource
	Options       *ExecutionOptions
	SystemService bool
	Directory     string
	Command       string
	Arguments     []string
}

type ProcessInfo struct {
	Name      string
	Pid       int
	Command   string
	Arguments []string
	Stdin     string
	Stdout    string
}

type ProcessStatusRequest struct {
	Target  *Resource
	Command string
}

type ProcessStopRequest struct {
	Target *Resource
	Pid    int
}

type processService struct {
	*AbstractService
}

func (s *processService) Run(context *Context, request interface{}) *ServiceResponse {
	var response = &ServiceResponse{Status: "ok"}
	var err error
	switch actualRequest := request.(type) {
	case *ProcessStartRequest:
		response.Response, err = s.startProcess(context, actualRequest)
		if err != nil {
			response.Error = fmt.Sprintf("Failed to start process: %v, %v", actualRequest.Name, err)
		}
	case *ProcessStopRequest:
		response.Response, err = s.stopProcess(context, actualRequest)
		if err != nil {
			response.Error = fmt.Sprintf("Failed to stop process: %v, %v", actualRequest.Pid, err)
		}

	}
	if response.Error != "" {
		response.Status = "err"
	}
	return response
}

func (s *processService) checkProcess(context *Context, request *ProcessStatusRequest) ([]*ProcessInfo, error) {
	commandResponse, err := context.Execute(request.Target, &ManagedCommand{
		Executions: []*Execution{
			{
				Command: "ps -ev | grep " + request.Command,
			},
		},
	})

	if err != nil {
		return nil, err
	}

	var state = context.State()
	if ! state.Has(processesKey) {
		state.Put(processesKey, common.NewMap())
	}
	var processes = state.GetMap(processesKey)
	processes.Put(request.Command, 0)
	var result = make([]*ProcessInfo, 0)

	for _, line := range strings.Split(commandResponse.Stdout(), "\r\n") {
		if strings.Contains(line, "grep") {
			continue
		}
		pid, ok := ExtractColumn(line, 0)
		if !ok {
			continue
		}
		argumentsIndex := strings.Index(line, request.Command)
		var arguments []string
		if argumentsIndex != -1 {
			args := strings.Trim(line[argumentsIndex+len(request.Command)+1:], " &\t")
			arguments = strings.Split(args, " ")
		}
		info := &ProcessInfo{
			Pid:       toolbox.AsInt(pid),
			Command:   request.Command,
			Arguments: arguments,
			Stdout:    line,
		}
		result = append(result, info)
		processes.Put(request.Command, info.Pid)
	}
	return result, nil
}

func (s *processService) stopProcess(context *Context, request *ProcessStopRequest) (*CommandInfo, error) {
	commandResult, err := context.Execute(request.Target, &ManagedCommand{
		Executions: []*Execution{
			{
				Command: fmt.Sprintf("kill -9 %v", request.Pid),
			},
		},
	})
	if err != nil {
		return nil, err
	}
	state := context.State()
	if ! state.Has(processesKey) {
		state.Put(processesKey, common.NewMap())
	}
	var processes = state.GetMap(processesKey)
	for k, pid := range processes {
		if toolbox.AsInt(pid) == request.Pid {
			state[k] = 0
			break;
		}
	}
	return commandResult, err
}

func indexProcesses(processes ...*ProcessInfo) map[int]*ProcessInfo {
	var result = make(map[int]*ProcessInfo)
	for _, process := range processes {
		result[process.Pid] = process
	}
	return result
}

func (s *processService) startProcess(context *Context, request *ProcessStartRequest) (*ProcessInfo, error) {
	origProcesses, err := s.checkProcess(context, &ProcessStatusRequest{
		Target:  request.Target,
		Command: request.Command,
	})
	if err != nil {
		return nil, err
	}
	for _, process := range origProcesses {
		if strings.Join(process.Arguments, " ") == strings.Join(request.Arguments, " ") {
			_, err := s.stopProcess(context, &ProcessStopRequest{
				Pid:    process.Pid,
				Target: request.Target,
			})
			if err != nil {
				return nil, err
			}
		}
	}
	changeDirCommand := fmt.Sprintf("cd %v ", request.Directory)
	startCommand := fmt.Sprintf("nohup %v %v &", request.Command, strings.Join(request.Arguments, " "))
	commandInfo, err := context.Execute(request.Target, &ManagedCommand{
		Options: request.Options,
		Executions: []*Execution{
			{
				Command: changeDirCommand,
			},
			{
				Command: startCommand,
			},
		},
	})
	if err != nil {
		return nil, err
	}
	time.Sleep(time.Second)
	newProcesses, err := s.checkProcess(context, &ProcessStatusRequest{
		Target:  request.Target,
		Command: request.Command,
	})
	if err != nil {
		return nil, err
	}

	if len(newProcesses) == 1 {
		return newProcesses[0], nil
	}
	existingProcesses := indexProcesses(origProcesses...)
	var result *ProcessInfo
	for _, candidate := range newProcesses {
			if _, has := existingProcesses[candidate.Pid]; !has {
			result = candidate
			break
		}
	}
	if result == nil {
		return nil, fmt.Errorf("Failed to get info about prorcess %v", request.Command)
	}
	result.Stdout = commandInfo.Stdout()
	result.Stdin = fmt.Sprintf("%v && %v", changeDirCommand, startCommand)
	return result, nil
}

func (s *processService) NewRequest(name string) (interface{}, error) {
	switch name {
	case "start":
		return &ProcessStartRequest{}, nil
	case "check":
		return &ProcessStatusRequest{}, nil
	case "stop":
		return &ProcessStopRequest{}, nil

	}
	return nil, fmt.Errorf("Unsupported name: %v", name)
}

func NewProcessService() Service {
	var result = &processService{
		AbstractService: NewAbstractService(ProcessServiceId),
	}
	result.AbstractService.Service = result
	return result
}
<|endoftext|>


<commit_before>package session

import (
	"bytes"
	"log"
	"net/http"
)

// Deprecated: Middleware previously defined the signature for the session management
// middleware returned by Manage. Manage now returns a func(h http.Handler) http.Handler
// directly instead, so it's easier to use with middleware chaining packages like Alice.
type Middleware func(h http.Handler) http.Handler

/*
Manage returns a new session manager middleware instance. The first parameter
should be a valid storage engine, followed by zero or more functional options.

For example:

	session.Manage(memstore.New(0))

	session.Manage(memstore.New(0), session.Lifetime(14*24*time.Hour))

	session.Manage(memstore.New(0),
		session.Secure(true),
		session.Persist(true),
		session.Lifetime(14*24*time.Hour),
	)

The returned session manager can be used to wrap any http.Handler. It automatically
loads sessions based on the HTTP request and saves session data as and when necessary.
*/
func Manage(engine Engine, opts ...Option) func(h http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		do := *defaultOptions

		m := &manager{
			h:      h,
			engine: engine,
			opts:   &do,
		}

		for _, option := range opts {
			option(m.opts)
		}

		return m
	}
}

type manager struct {
	h      http.Handler
	engine Engine
	opts   *options
}

func (m *manager) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	sr, err := load(r, m.engine, m.opts)
	if err != nil {
		m.opts.errorFunc(w, r, err)
		return
	}
	bw := &bufferedResponseWriter{ResponseWriter: w}
	m.h.ServeHTTP(bw, sr)

	err = write(w, sr)
	if err != nil {
		m.opts.errorFunc(w, r, err)
		return
	}

	if bw.code != 0 {
		w.WriteHeader(bw.code)
	}
	w.Write(bw.buf.Bytes())
}

type bufferedResponseWriter struct {
	http.ResponseWriter
	buf  bytes.Buffer
	code int
}

func (bw *bufferedResponseWriter) Write(b []byte) (int, error) {
	return bw.buf.Write(b)
}

func (bw *bufferedResponseWriter) WriteHeader(code int) {
	bw.code = code
}

func (bw *bufferedResponseWriter) Flush() {
	f, ok := bw.ResponseWriter.(http.Flusher)
	if ok == true {
		bw.ResponseWriter.Write(bw.buf.Bytes())
		f.Flush()
		bw.buf.Reset()
	}
}

func defaultErrorFunc(w http.ResponseWriter, r *http.Request, err error) {
	log.Output(2, err.Error())
	http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
}
<commit_msg>[bugfix] Surface error from ResponseWriter.Write<commit_after>package session

import (
	"bytes"
	"log"
	"net/http"
)

// Deprecated: Middleware previously defined the signature for the session management
// middleware returned by Manage. Manage now returns a func(h http.Handler) http.Handler
// directly instead, so it's easier to use with middleware chaining packages like Alice.
type Middleware func(h http.Handler) http.Handler

/*
Manage returns a new session manager middleware instance. The first parameter
should be a valid storage engine, followed by zero or more functional options.

For example:

	session.Manage(memstore.New(0))

	session.Manage(memstore.New(0), session.Lifetime(14*24*time.Hour))

	session.Manage(memstore.New(0),
		session.Secure(true),
		session.Persist(true),
		session.Lifetime(14*24*time.Hour),
	)

The returned session manager can be used to wrap any http.Handler. It automatically
loads sessions based on the HTTP request and saves session data as and when necessary.
*/
func Manage(engine Engine, opts ...Option) func(h http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		do := *defaultOptions

		m := &manager{
			h:      h,
			engine: engine,
			opts:   &do,
		}

		for _, option := range opts {
			option(m.opts)
		}

		return m
	}
}

type manager struct {
	h      http.Handler
	engine Engine
	opts   *options
}

func (m *manager) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	sr, err := load(r, m.engine, m.opts)
	if err != nil {
		m.opts.errorFunc(w, r, err)
		return
	}
	bw := &bufferedResponseWriter{ResponseWriter: w}
	m.h.ServeHTTP(bw, sr)

	err = write(w, sr)
	if err != nil {
		m.opts.errorFunc(w, r, err)
		return
	}

	if bw.code != 0 {
		w.WriteHeader(bw.code)
	}

	_, err = w.Write(bw.buf.Bytes())
	if err != nil {
		m.opts.errorFunc(w, r, err)
	}
}

type bufferedResponseWriter struct {
	http.ResponseWriter
	buf  bytes.Buffer
	code int
}

func (bw *bufferedResponseWriter) Write(b []byte) (int, error) {
	return bw.buf.Write(b)
}

func (bw *bufferedResponseWriter) WriteHeader(code int) {
	bw.code = code
}

func (bw *bufferedResponseWriter) Flush() {
	f, ok := bw.ResponseWriter.(http.Flusher)
	if ok == true {
		bw.ResponseWriter.Write(bw.buf.Bytes())
		f.Flush()
		bw.buf.Reset()
	}
}

func defaultErrorFunc(w http.ResponseWriter, r *http.Request, err error) {
	log.Output(2, err.Error())
	http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
}
<|endoftext|>


<commit_before>// TODO: better name? network topology?
package graph

import (
	"container/ring"
	"crypto/md5"
	"encoding/hex"
	"io"

	"github.com/Sirupsen/logrus"
)

// TODO: maintain some maps for easier lookup
type NetworkGraph struct {
	// nodeName -> Node
	NodesMap map[string]*NetworkNode

	// nodeName,nodeName -> NetworkLink
	LinksMap map[string]*NetworkLink

	RoutesMap map[string]*NetworkRoute
}

func Create() *NetworkGraph {
	return &NetworkGraph{
		NodesMap:  make(map[string]*NetworkNode),
		LinksMap:  make(map[string]*NetworkLink),
		RoutesMap: make(map[string]*NetworkRoute),
	}
}

func (g *NetworkGraph) IncrNode(name string) *NetworkNode {
	n, ok := g.NodesMap[name]
	// if this one doesn't exist, lets add it
	if !ok {
		n = NewNetworkNode(name)
		g.NodesMap[name] = n
	}
	n.refCount++
	return n
}

func (g *NetworkGraph) GetNode(name string) *NetworkNode {
	n, _ := g.NodesMap[name]
	return n
}

func (g *NetworkGraph) GetNodeCount() int {
	return len(g.NodesMap)
}

func (g *NetworkGraph) DecrNode(name string) {
	n, ok := g.NodesMap[name]

	if !ok {
		logrus.Warningf("Attempted to remove node with ip %v which wasn't in the graph", name)
		return
	}

	n.refCount--
	if n.refCount == 0 {
		delete(g.NodesMap, name)
	}
}

func (g *NetworkGraph) IncrLink(src, dst string) *NetworkLink {
	key := src + "," + dst
	l, ok := g.LinksMap[key]
	if !ok {
		l = &NetworkLink{
			Src: g.IncrNode(src),
			Dst: g.IncrNode(dst),
		}
		g.LinksMap[key] = l
	}
	l.refCount++
	return l
}

func (g *NetworkGraph) GetLink(src, dst string) *NetworkLink {
	key := src + "," + dst
	l, _ := g.LinksMap[key]
	return l
}

func (g *NetworkGraph) GetLinkCount() int {
	return len(g.LinksMap)
}

func (g *NetworkGraph) DecrLink(src, dst string) {
	key := src + "," + dst
	l, ok := g.LinksMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove link %v which wasn't in the graph", key)
		return
	}
	// decrement ourselves
	l.refCount--
	if l.refCount == 0 {
		// Decrement our children
		g.DecrNode(src)
		g.DecrNode(dst)
		delete(g.LinksMap, key)
	}
}

func (g *NetworkGraph) pathKey(hops []string) string {
	h := md5.New()
	for _, hop := range hops {
		io.WriteString(h, hop)
	}
	return hex.EncodeToString(h.Sum(nil))
}

func (g *NetworkGraph) IncrRoute(hops []string) *NetworkRoute {
	key := g.pathKey(hops)

	// check if we have a route for this already
	route, ok := g.RoutesMap[key]
	// if we don't have it, lets make it
	if !ok {
		logrus.Infof("New Route: key=%s %v", key, hops)
		path := make([]*NetworkNode, 0, len(hops))
		for i, hop := range hops {
			path = append(path, g.IncrNode(hop))
			// If there was something prior-- lets add the link as well
			if i-1 > 0 {
				g.IncrLink(hops[i-1], hop)
			}
		}
		route = &NetworkRoute{
			Path:       path,
			State:      Up,
			metricRing: ring.New(10), // TODO: config
		}
		g.RoutesMap[key] = route
	}

	// increment route's refcount
	route.refCount++

	return route
}

func (g *NetworkGraph) GetRoute(hops []string) *NetworkRoute {
	r, _ := g.RoutesMap[g.pathKey(hops)]
	return r
}

func (g *NetworkGraph) GetRouteCount() int {
	return len(g.RoutesMap)
}

func (g *NetworkGraph) DecrRoute(hops []string) {
	key := g.pathKey(hops)
	r, ok := g.RoutesMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove route %v which wasn't in the graph", key)
		return
	}

	r.refCount--
	if r.refCount == 0 {
		// decrement all the links/nodes as well
		for i, node := range r.Path {
			g.DecrNode(node.Name)
			if i-1 > 0 {
				g.DecrLink(r.Path[i-1].Name, node.Name)
			}
		}

		delete(g.RoutesMap, key)
	}
}
<commit_msg>Some tag cleanup<commit_after>// TODO: better name? network topology?
package graph

import (
	"container/ring"
	"crypto/md5"
	"encoding/hex"
	"io"

	"github.com/Sirupsen/logrus"
)

// TODO: maintain some maps for easier lookup
type NetworkGraph struct {
	// nodeName -> Node
	NodesMap map[string]*NetworkNode `json:"nodes"`

	// nodeName,nodeName -> NetworkLink
	LinksMap map[string]*NetworkLink `json:"edges"`

	RoutesMap map[string]*NetworkRoute `json:"routes"`
}

func Create() *NetworkGraph {
	return &NetworkGraph{
		NodesMap:  make(map[string]*NetworkNode),
		LinksMap:  make(map[string]*NetworkLink),
		RoutesMap: make(map[string]*NetworkRoute),
	}
}

func (g *NetworkGraph) IncrNode(name string) *NetworkNode {
	n, ok := g.NodesMap[name]
	// if this one doesn't exist, lets add it
	if !ok {
		n = NewNetworkNode(name)
		g.NodesMap[name] = n
	}
	n.refCount++
	return n
}

func (g *NetworkGraph) GetNode(name string) *NetworkNode {
	n, _ := g.NodesMap[name]
	return n
}

func (g *NetworkGraph) GetNodeCount() int {
	return len(g.NodesMap)
}

func (g *NetworkGraph) DecrNode(name string) {
	n, ok := g.NodesMap[name]

	if !ok {
		logrus.Warningf("Attempted to remove node with ip %v which wasn't in the graph", name)
		return
	}

	n.refCount--
	if n.refCount == 0 {
		delete(g.NodesMap, name)
	}
}

func (g *NetworkGraph) IncrLink(src, dst string) *NetworkLink {
	key := src + "," + dst
	l, ok := g.LinksMap[key]
	if !ok {
		l = &NetworkLink{
			Src: g.IncrNode(src),
			Dst: g.IncrNode(dst),
		}
		g.LinksMap[key] = l
	}
	l.refCount++
	return l
}

func (g *NetworkGraph) GetLink(src, dst string) *NetworkLink {
	key := src + "," + dst
	l, _ := g.LinksMap[key]
	return l
}

func (g *NetworkGraph) GetLinkCount() int {
	return len(g.LinksMap)
}

func (g *NetworkGraph) DecrLink(src, dst string) {
	key := src + "," + dst
	l, ok := g.LinksMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove link %v which wasn't in the graph", key)
		return
	}
	// decrement ourselves
	l.refCount--
	if l.refCount == 0 {
		// Decrement our children
		g.DecrNode(src)
		g.DecrNode(dst)
		delete(g.LinksMap, key)
	}
}

func (g *NetworkGraph) pathKey(hops []string) string {
	h := md5.New()
	for _, hop := range hops {
		io.WriteString(h, hop)
	}
	return hex.EncodeToString(h.Sum(nil))
}

func (g *NetworkGraph) IncrRoute(hops []string) *NetworkRoute {
	key := g.pathKey(hops)

	// check if we have a route for this already
	route, ok := g.RoutesMap[key]
	// if we don't have it, lets make it
	if !ok {
		logrus.Infof("New Route: key=%s %v", key, hops)
		path := make([]*NetworkNode, 0, len(hops))
		for i, hop := range hops {
			path = append(path, g.IncrNode(hop))
			// If there was something prior-- lets add the link as well
			if i-1 > 0 {
				g.IncrLink(hops[i-1], hop)
			}
		}
		route = &NetworkRoute{
			Path:       path,
			State:      Up,
			metricRing: ring.New(10), // TODO: config
		}
		g.RoutesMap[key] = route
	}

	// increment route's refcount
	route.refCount++

	return route
}

func (g *NetworkGraph) GetRoute(hops []string) *NetworkRoute {
	r, _ := g.RoutesMap[g.pathKey(hops)]
	return r
}

func (g *NetworkGraph) GetRouteCount() int {
	return len(g.RoutesMap)
}

func (g *NetworkGraph) DecrRoute(hops []string) {
	key := g.pathKey(hops)
	r, ok := g.RoutesMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove route %v which wasn't in the graph", key)
		return
	}

	r.refCount--
	if r.refCount == 0 {
		// decrement all the links/nodes as well
		for i, node := range r.Path {
			g.DecrNode(node.Name)
			if i-1 > 0 {
				g.DecrLink(r.Path[i-1].Name, node.Name)
			}
		}

		delete(g.RoutesMap, key)
	}
}
<|endoftext|>


<commit_before>package graph

// Vertex ...
type Vertex interface{}

// Edge ...
type Edge interface {
	Vertices() (u, v Vertex)
}

// DirectedEdge ...
type DirectedEdge interface {
	Edge
	Source() Vertex
	Target() Vertex
}

// LabeledEdge ...
type LabeledEdge interface {
	Edge
	Label() string
}

// WeightedEdge ...
type WeightedEdge interface {
	Edge
	Weight() float64
}

// DirectedWeightedEdge ...
type DirectedWeightedEdge interface {
	Edge
	Source() Vertex
	Target() Vertex
	Weight() float64
}

type edge struct {
	u, v Vertex
}

type directedEdge struct {
	Edge
}

type labeledEdge struct {
	Edge
	label string
}

type weightedEdge struct {
	Edge
	weight float64
}

// NewEdge ...
func NewEdge(u, v Vertex) Edge {
	e := edge{u: u, v: v}
	return e
}

func (e edge) Vertices() (u, v Vertex) {
	return e.u, e.v
}

// NewDirectedEdge ...
func NewDirectedEdge(u, v Vertex) DirectedEdge {
	e := directedEdge{
		Edge: edge{u: u, v: v},
	}
	return e
}

func (e directedEdge) Vertices() (u, v Vertex) {
	return e.Vertices()
}

func (e directedEdge) Source() (u Vertex) {
	u, _ = e.Vertices()
	return u
}

func (e directedEdge) Target() (v Vertex) {
	_, v = e.Vertices()
	return v
}

// NewLabeledEdge ...
func NewLabeledEdge(u, v Vertex, s string) LabeledEdge {
	e := labeledEdge{
		Edge:  edge{u, v},
		label: s,
	}
	return e
}

func (e labeledEdge) Vertices() (u, v Vertex) {
	return e.Vertices()
}

func (e labeledEdge) Label() string {
	return e.label
}

// NewWeightedEdge ...
func NewWeightedEdge(u, v Vertex, w float64) WeightedEdge {
	e := weightedEdge{
		Edge:   edge{u, v},
		weight: w,
	}
	return e
}

func (e weightedEdge) Vertices() (u, v Vertex) {
	return e.Vertices()
}

func (e weightedEdge) Weight() float64 {
	return e.weight
}

// Graph ...
type Graph interface {
	Has(Vertex) bool
	HasE(Edge) bool
	Connected(a, b Vertex) bool
	EdgesOf(Vertex) []Edge

	Edges() []Edge
	Vertices() []Vertex

	AddEdge(Edge)
	DelEdge(Edge)
}

// DirectedGraph ...
type DirectedGraph interface {
	Graph
	InDegree(Vertex) int
	OutDegree(Vertex) int
}

// WeightedGraph ...
type WeightedGraph interface {
	Graph
	Sum(Vertex) float64
}
<commit_msg>graph: Type impl<commit_after>package graph

// Vertex ...
type Vertex interface{}

// Edge ...
type Edge interface {
	Vertices() (u, v Vertex)
}

// DirectedEdge ...
type DirectedEdge interface {
	Edge
	Source() Vertex
	Target() Vertex
}

// LabeledEdge ...
type LabeledEdge interface {
	Edge
	Label() string
}

// WeightedEdge ...
type WeightedEdge interface {
	Edge
	Weight() float64
}

// DirectedWeightedEdge ...
type DirectedWeightedEdge interface {
	Edge
	Source() Vertex
	Target() Vertex
	Weight() float64
}

type edge struct {
	u, v Vertex
}

type directedEdge struct {
	Edge
}

type labeledEdge struct {
	Edge
	label string
}

type weightedEdge struct {
	Edge
	weight float64
}

// NewEdge ...
func NewEdge(u, v Vertex) Edge {
	e := edge{u: u, v: v}
	return e
}

func (e edge) Vertices() (u, v Vertex) {
	return e.u, e.v
}

// NewDirectedEdge ...
func NewDirectedEdge(u, v Vertex) DirectedEdge {
	e := directedEdge{
		Edge: edge{u: u, v: v},
	}
	return e
}

func (e directedEdge) Vertices() (u, v Vertex) {
	return e.Vertices()
}

func (e directedEdge) Source() (u Vertex) {
	u, _ = e.Vertices()
	return u
}

func (e directedEdge) Target() (v Vertex) {
	_, v = e.Vertices()
	return v
}

// NewLabeledEdge ...
func NewLabeledEdge(u, v Vertex, s string) LabeledEdge {
	e := labeledEdge{
		Edge:  edge{u, v},
		label: s,
	}
	return e
}

func (e labeledEdge) Vertices() (u, v Vertex) {
	return e.Vertices()
}

func (e labeledEdge) Label() string {
	return e.label
}

// NewWeightedEdge ...
func NewWeightedEdge(u, v Vertex, w float64) WeightedEdge {
	e := weightedEdge{
		Edge:   edge{u, v},
		weight: w,
	}
	return e
}

func (e weightedEdge) Vertices() (u, v Vertex) {
	return e.Vertices()
}

func (e weightedEdge) Weight() float64 {
	return e.weight
}

// Graph ...
type Graph interface {
	Has(Vertex) bool
	HasE(Edge) bool
	Connected(a, b Vertex) bool
	EdgesOf(Vertex) []Edge

	Edges() []Edge
	Vertices() []Vertex

	AddEdge(Edge)
	DelEdge(Edge)
}

// DirectedGraph ...
type DirectedGraph interface {
	Graph
	InDegree(Vertex) int
	OutDegree(Vertex) int
}

// WeightedGraph ...
type WeightedGraph interface {
	Graph
	Sum(Vertex) float64
}

// Type ...
type Type int

const (
	// Simple ...
	Simple Type = 1 << iota
	// Loop ...
	Loop
	// Parallel ...
	Parallel

	// Directed ...
	Directed
	// Labeled ...
	Labeled
	// Weighted ...
	Weighted
)

// Build ...
func Build(t Type) Graph {
	return nil
}

var _ Graph = Build(Simple | Loop | Parallel | Directed | Labeled | Weighted)
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc.
//
// 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 graph manages programs stored as graphs.
package graph

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// Channel models a channel. It can be marshalled and unmarshalled to JSON sensibly.
type Channel struct {
	Name string `json:"name"`
	Type string `json:"type"`
	Cap  int    `json:"cap"`
}

// Graph describes a Go program as a graph. It can be marshalled and unmarshalled to JSON sensibly.
type Graph struct {
	SourcePath  string              `json:"-"` // path to the JSON source.
	Name        string              `json:"name"`
	PackagePath string              `json:"package_path"`
	Imports     []string            `json:"imports"`
	Nodes       map[string]*Node    `json:"nodes"`
	Channels    map[string]*Channel `json:"channels"`
}

// PackageName extracts the name of the package from the package path ("full" package name).
func (g *Graph) PackageName() string {
	i := strings.LastIndex(g.PackagePath, "/")
	if i < 0 {
		return g.PackagePath
	}
	return g.PackagePath[i:]
}

// LoadJSON loads a JSON-encoded Graph from an io.Reader.
func LoadJSON(r io.Reader, sourcePath string) (*Graph, error) {
	dec := json.NewDecoder(r)
	var g Graph
	if err := dec.Decode(&g); err != nil {
		return nil, err
	}
	g.SourcePath = sourcePath
	return &g, nil
}

// LoadJSONFile loads a JSON-encoded Graph from a file at a given path.
func LoadJSONFile(path string) (*Graph, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return LoadJSON(f, path)
}

// WriteJSONTo writes nicely-formatted JSON to the given Writer.
func (g *Graph) WriteJSONTo(w io.Writer) error {
	enc := json.NewEncoder(w)
	enc.SetIndent("", "\t") // For diffability!
	return enc.Encode(g)
}

// SaveJSONFile saves the JSON-encoded Graph to the SourcePath.
func (g *Graph) SaveJSONFile() error {
	f, err := ioutil.TempFile(filepath.Dir(g.SourcePath), filepath.Base(g.SourcePath))
	if err != nil {
		return err
	}
	defer f.Close()
	if err := g.WriteJSONTo(f); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return os.Rename(f.Name(), g.SourcePath)
}

// WriteDotTo writes the Dot language view of the graph to the io.Writer.
func (g *Graph) WriteDotTo(dst io.Writer) error { return dotTemplate.Execute(dst, g) }

// WriteGoTo writes the Go language view of the graph to the io.Writer.
func (g *Graph) WriteGoTo(w io.Writer) error {
	buf := &bytes.Buffer{}
	if err := goTemplate.Execute(buf, g); err != nil {
		return err
	}
	return goimports(w, buf)
}

// WriteGoRunnerTo outputs a simple main package and function for calling Run on
// a generated Go package.
func (g *Graph) WriteGoRunnerTo(w io.Writer) error {
	return goRunnerTemplate.Execute(w, g)
}

// Build saves the graph as Go source code and tries to build it.
func (g *Graph) Build() error {
	gopath, ok := os.LookupEnv("GOPATH")
	if !ok || gopath == "" {
		return errors.New("cannot use $GOPATH; empty or undefined")
	}
	pp := filepath.Join(gopath, "src", g.PackagePath)
	if err := os.Mkdir(pp, os.FileMode(0755)); err != nil {
		log.Printf("Could not make path %q, continuing: %v", pp, err)
	}
	mp := filepath.Join(pp, "generated.go")
	f, err := os.Create(mp)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := g.WriteGoTo(f); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return exec.Command(`go`, `build`, g.PackagePath).Run()
}

func (g *Graph) run() error {
	fn := filepath.Join(os.TempDir(), fmt.Sprintf("shenzhen-go-runner.%s.go", g.PackageName()))
	f, err := os.Create(fn)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := g.WriteGoRunnerTo(f); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return exec.Command(`go`, `run`, fn).Run()
}

// BuildAndRun saves the project as Go source code, builds it, and runs it.
func (g *Graph) BuildAndRun() error {
	if err := g.Build(); err != nil {
		return err
	}
	return g.run()
}

// DeclaredChannels returns the given channels which exist in g.Channels.
func (g *Graph) DeclaredChannels(chans []string) []string {
	r := make([]string, 0, len(chans))
	for _, d := range chans {
		if _, found := g.Channels[d]; found {
			r = append(r, d)
		}
	}
	return r
}
<commit_msg>Improve error output<commit_after>// Copyright 2016 Google Inc.
//
// 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 graph manages programs stored as graphs.
package graph

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// Channel models a channel. It can be marshalled and unmarshalled to JSON sensibly.
type Channel struct {
	Name string `json:"name"`
	Type string `json:"type"`
	Cap  int    `json:"cap"`
}

// Graph describes a Go program as a graph. It can be marshalled and unmarshalled to JSON sensibly.
type Graph struct {
	SourcePath  string              `json:"-"` // path to the JSON source.
	Name        string              `json:"name"`
	PackagePath string              `json:"package_path"`
	Imports     []string            `json:"imports"`
	Nodes       map[string]*Node    `json:"nodes"`
	Channels    map[string]*Channel `json:"channels"`
}

// PackageName extracts the name of the package from the package path ("full" package name).
func (g *Graph) PackageName() string {
	i := strings.LastIndex(g.PackagePath, "/")
	if i < 0 {
		return g.PackagePath
	}
	return g.PackagePath[i:]
}

// LoadJSON loads a JSON-encoded Graph from an io.Reader.
func LoadJSON(r io.Reader, sourcePath string) (*Graph, error) {
	dec := json.NewDecoder(r)
	var g Graph
	if err := dec.Decode(&g); err != nil {
		return nil, err
	}
	g.SourcePath = sourcePath
	return &g, nil
}

// LoadJSONFile loads a JSON-encoded Graph from a file at a given path.
func LoadJSONFile(path string) (*Graph, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return LoadJSON(f, path)
}

// WriteJSONTo writes nicely-formatted JSON to the given Writer.
func (g *Graph) WriteJSONTo(w io.Writer) error {
	enc := json.NewEncoder(w)
	enc.SetIndent("", "\t") // For diffability!
	return enc.Encode(g)
}

// SaveJSONFile saves the JSON-encoded Graph to the SourcePath.
func (g *Graph) SaveJSONFile() error {
	f, err := ioutil.TempFile(filepath.Dir(g.SourcePath), filepath.Base(g.SourcePath))
	if err != nil {
		return err
	}
	defer f.Close()
	if err := g.WriteJSONTo(f); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return os.Rename(f.Name(), g.SourcePath)
}

// WriteDotTo writes the Dot language view of the graph to the io.Writer.
func (g *Graph) WriteDotTo(dst io.Writer) error { return dotTemplate.Execute(dst, g) }

// WriteGoTo writes the Go language view of the graph to the io.Writer.
func (g *Graph) WriteGoTo(w io.Writer) error {
	buf := &bytes.Buffer{}
	if err := goTemplate.Execute(buf, g); err != nil {
		return err
	}
	return goimports(w, buf)
}

// WriteGoRunnerTo outputs a simple main package and function for calling Run on
// a generated Go package.
func (g *Graph) WriteGoRunnerTo(w io.Writer) error {
	return goRunnerTemplate.Execute(w, g)
}

// Build saves the graph as Go source code and tries to build it.
func (g *Graph) Build() error {
	gopath, ok := os.LookupEnv("GOPATH")
	if !ok || gopath == "" {
		return errors.New("cannot use $GOPATH; empty or undefined")
	}
	pp := filepath.Join(gopath, "src", g.PackagePath)
	if err := os.Mkdir(pp, os.FileMode(0755)); err != nil {
		log.Printf("Could not make path %q, continuing: %v", pp, err)
	}
	mp := filepath.Join(pp, "generated.go")
	f, err := os.Create(mp)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := g.WriteGoTo(f); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	o, err := exec.Command(`go`, `build`, g.PackagePath).CombinedOutput()
	if err != nil {
		// TODO: better error type
		return fmt.Errorf("go build: %v:\n%s", err, o)
	}
	return nil
}

func (g *Graph) run() error {
	fn := filepath.Join(os.TempDir(), fmt.Sprintf("shenzhen-go-runner.%s.go", g.PackageName()))
	f, err := os.Create(fn)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := g.WriteGoRunnerTo(f); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return exec.Command(`go`, `run`, fn).Run()
}

// BuildAndRun saves the project as Go source code, builds it, and runs it.
func (g *Graph) BuildAndRun() error {
	if err := g.Build(); err != nil {
		return err
	}
	return g.run()
}

// DeclaredChannels returns the given channels which exist in g.Channels.
func (g *Graph) DeclaredChannels(chans []string) []string {
	r := make([]string, 0, len(chans))
	for _, d := range chans {
		if _, found := g.Channels[d]; found {
			r = append(r, d)
		}
	}
	return r
}
<|endoftext|>


<commit_before>package mallory

import (
	"bufio"
	"bytes"
	"crypto/sha1"
	"crypto/tls"
	"encoding/pem"
	"fmt"
	"io"
	"math/big"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"sync"
	"time"
)

// all write on this should be sync between threads
type EngineGAE struct {
	// Global config
	Env *Env
	// work space for this engine
	Work string
	// place store certificates
	CertsDir string
	// Loaded certificate, contains the root certificate and private key
	RootCA *tls.Certificate
	// Pool of auto generated fake certificates signed by RootCert
	Certs *CertPool
}

func NewEngineGAE(e *Env) *EngineGAE {
	self := &EngineGAE{Env: e}
	self.Work = path.Join(self.Env.Work, "gae")
	self.CertsDir = path.Join(self.Work, "certs")
	self.Certs = NewCertPool()
	return self
}

func (self *EngineGAE) Init() error {
	err := os.MkdirAll(self.CertsDir, 0755)
	if err != nil && !os.IsExist(err) {
		return err
	}

	rcert, err := tls.LoadX509KeyPair(self.Env.Cert, self.Env.Key)
	if err != nil {
		return err
	}
	self.RootCA = &rcert
	return nil
}

// Data flow:
//  1. Receive client request R1
//  2. Write R1 as the body of a new request R2
//  3. Post request R2 to remote GAE
//  4. Receive response P1 from GAE
//  5. Read remote server(which the client want to connect with) resonse P2 from the body of P1
//  6. Send P2 as the response to client
func (self *EngineGAE) Serve(s *Session) {
	w, r := s.ResponseWriter, s.Request
	if r.Method == "CONNECT" {
		s.Error("this function can not handle CONNECT method")
		return
	}
	start := time.Now()

	// write the client request and post to remote
	// Note: WriteProxy keeps the full request URI
	var buf bytes.Buffer
	if err := r.WriteProxy(&buf); err != nil {
		s.Error("WriteProxy: %s", err.Error())
		return
	}

	// use httpS to keep all things secure,
	// the second phase of CONNECT also uses this.
	url := fmt.Sprintf("https://%s.appspot.com/http", self.Env.AppSpot)
	// for debug
	if self.Env.AppSpot == "debug" {
		url = "http://localhost:8080/http"
	}

	// post client request as body data
	resp, err := http.Post(url, "application/data", &buf)
	if err != nil {
		s.Error("Post: %s", err.Error())
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusInternalServerError {
		s.Error("Post: %s", resp.Status)
		return
	}

	// the response for the requst of client
	cres, err := http.ReadResponse(bufio.NewReader(resp.Body), r)
	if err != nil {
		s.Error("ReadResponse: %s", err.Error())
		return
	}
	defer cres.Body.Close()

	// copy headers
	CopyResponseHeader(w, cres)

	// please prepare header first and write them
	w.WriteHeader(cres.StatusCode)

	_, err = io.Copy(w, cres.Body)
	if err != nil {
		s.Error("Copy: %s", err.Error())
		return
	}

	s.Info("RESPONSE %s %s %s", r.URL.Host, resp.Status, time.Since(start).String())
}

// FIXME:
//  Impossible to connect gae and handle it as a normal TCP connection?
//  GAE only provide http handlers? At least I don't know how to handle to TCP connection on GAE server.
//  NOTE: GAE socket service can only be available for billing users. So free users is unable to use the
//  long term connection. And do what we did in EngineDirect.
//  So we can only use urlfetch.Client.Transport.RoundTrip to do http or https method.
//  Generally, the CONNECT method can be used for any purpose for the advantage of TCP connection.
//  The proxy doesn't need to know what the real underlying protocol or what it is, just need to copy
//  data from client to server, and copy the response from the server to client without any interpret.
//  Now what we can do and had been done by some GAE proxies is that, extract the underlying protocol!!!
//  GAE can only handle limited protocols with urlfetch module, such as http and https.
//  Use Hijacker to get the underlying connection
//
// Data flow:
//  1. Detect host and port
//  2. Hijack the client connection
//  3. Dial self
//  4. Return 200 OK if is successfully
//  5. Get cached or create new signed certificate
//  6. Wrap client connection with TLS and make handshake
//  7. Receive http request
//  8. Write request as a proxy request to self, HTTP handler
//  9. Copy response to client...
func (self *EngineGAE) Connect(s *Session) {
	w, r := s.ResponseWriter, s.Request
	if r.Method != "CONNECT" {
		s.Error("this function can only handle CONNECT method")
		return
	}
	start := time.Now()

	// Only support HTTPS protocol, which is connected with port 443
	host, port, err := net.SplitHostPort(r.URL.Host)
	if err != nil {
		s.Error("SplitHostPort: %s", err.Error())
		return
	}

	if port != "443" {
		s.Error("unsupported CONNECT port: %s", port)
		return
	}

	// hijack the connection to make SSL handshake
	hij, ok := w.(http.Hijacker)
	if !ok {
		s.Error("Server does not support Hijacker")
		return
	}

	conn, _, err := hij.Hijack()
	if err != nil {
		s.Error("Hijack: %s", err.Error())
		return
	}
	defer conn.Close()

	// dial self to transport application data, http request
	rconn, err := net.Dial("tcp", self.Env.Addr)
	if err != nil {
		s.Error("Dial: %s", err.Error())
		return
	}
	defer rconn.Close()

	// Once connected successfully, return OK
	conn.Write([]byte("HTTP/1.1 200 OK\r\n\r\n"))

	// get the fake cert, every host should have its own cert
	cert, err := self.GetCert(s, host)
	if err != nil {
		s.Error("GetCert: %s", err.Error())
		return
	}

	// assume the protocol of client connection is HTTPS
	// wrap it with TSL server
	config := &tls.Config{
		Certificates: []tls.Certificate{*cert},
		ServerName:   host,
	}
	sconn := tls.Server(conn, config)
	defer sconn.Close()

	// The TLS connection goes here
	if err := sconn.Handshake(); err != nil {
		// re-open browser to recover the handshake error:
		//    remote error: bad certificate
		s.Error("Handshake: %s", err.Error())
		return
	}

	// finally, we are at application layer, http request comes
	var wg sync.WaitGroup
	wg.Add(2)

	// read all requests, tls connection reues?
	go func() {
		for {
			req, err := http.ReadRequest(bufio.NewReader(sconn))
			if err != nil {
				if err != io.EOF {
					s.Error("ReadRequest: %s", err.Error())
				}
				break
			}

			// should re-wrap the URL with scheme "https://"
			req.URL, err = url.Parse("https://" + host + req.URL.String())
			req.Header.Set("Mallory-Session", strconv.FormatInt(s.ID, 10))

			// Now re-write the client request to self, HTTP handler
			err = req.WriteProxy(rconn)
			if err != nil {
				s.Error("WriteProxy: %s", err.Error())
				break
			}

			// close the persistent connection after reply the requset
			if req.Close {
				break
			}
		}
		wg.Done()
	}()

	// write back all responses
	go func() {
		for {
			// FIXME: how to get the right request instead of nil?
			cresp, err := http.ReadResponse(bufio.NewReader(rconn), nil)
			if err != nil {
				if err != io.EOF {
					s.Error("ReadResponse: %s", err.Error())
				}
				break
			}
			defer cresp.Body.Close()

			err = cresp.Write(sconn)
			if err != nil {
				// FIXME: write EOF? what does this mean?
				if err != io.EOF {
					s.Error("Write: %s", err.Error())
				}
				break
			}

			// close the persistent connection after reply the requset
			if cresp.Close {
				break
			}
		}
		wg.Done()
	}()

	// Keep connection until no data received from both client and server
	wg.Wait()

	s.Info("CLOSE %s %s", r.URL.Host, time.Since(start).String())
}

func (self *EngineGAE) GetCert(s *Session, host string) (cert *tls.Certificate, err error) {
	// firstly, try to find in memory
	cert = self.Certs.GetSafe(host)
	if cert != nil {
		return
	}

	// secondly, try to find on disk
	crtnam := path.Join(self.CertsDir, host+".crt")
	// we use the same key with CA
	crt, err := tls.LoadX509KeyPair(crtnam, self.Env.Key)
	if err == nil {
		cert = &crt
		self.Certs.AddSafe(host, cert)
		return
	} else if !os.IsNotExist(err) {
		s.Warn("LoadX509KeyPair: %s", err.Error())
	}

	// finally, try to create a new cert
	sn := sha1.Sum([]byte(host))
	config := &CertConfig{
		SerialNumber: new(big.Int).SetBytes(sn[:]),
		CommonName:   host, // FIXME: common name mismatch
	}
	cert, err = CreateSignedCert(self.RootCA, config)
	if err == nil {
		// add to memory
		self.Certs.AddSafe(host, cert)
		// save cert, fail is accepted
		fcrt, _err := os.Create(crtnam)
		if _err == nil {
			for _, c := range cert.Certificate {
				_err = pem.Encode(fcrt, &pem.Block{Type: "CERTIFICATE", Bytes: c})
				if _err != nil {
					break
				}
			}
			fcrt.Close()
			if _err != nil {
				s.Warn("Encode: %s", _err.Error())
				os.Remove(crtnam)
			}
		} else {
			s.Warn("Create %s", _err.Error())
		}
	}
	return
}
<commit_msg>simplify copy response<commit_after>package mallory

import (
	"bufio"
	"bytes"
	"crypto/sha1"
	"crypto/tls"
	"encoding/pem"
	"fmt"
	"io"
	"math/big"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"sync"
	"time"
)

// all write on this should be sync between threads
type EngineGAE struct {
	// Global config
	Env *Env
	// work space for this engine
	Work string
	// place store certificates
	CertsDir string
	// Loaded certificate, contains the root certificate and private key
	RootCA *tls.Certificate
	// Pool of auto generated fake certificates signed by RootCert
	Certs *CertPool
}

func NewEngineGAE(e *Env) *EngineGAE {
	self := &EngineGAE{Env: e}
	self.Work = path.Join(self.Env.Work, "gae")
	self.CertsDir = path.Join(self.Work, "certs")
	self.Certs = NewCertPool()
	return self
}

func (self *EngineGAE) Init() error {
	err := os.MkdirAll(self.CertsDir, 0755)
	if err != nil && !os.IsExist(err) {
		return err
	}

	rcert, err := tls.LoadX509KeyPair(self.Env.Cert, self.Env.Key)
	if err != nil {
		return err
	}
	self.RootCA = &rcert
	return nil
}

// Data flow:
//  1. Receive client request R1
//  2. Write R1 as the body of a new request R2
//  3. Post request R2 to remote GAE
//  4. Receive response P1 from GAE
//  5. Read remote server(which the client want to connect with) resonse P2 from the body of P1
//  6. Send P2 as the response to client
func (self *EngineGAE) Serve(s *Session) {
	w, r := s.ResponseWriter, s.Request
	if r.Method == "CONNECT" {
		s.Error("this function can not handle CONNECT method")
		return
	}
	start := time.Now()

	// write the client request and post to remote
	// Note: WriteProxy keeps the full request URI
	var buf bytes.Buffer
	if err := r.WriteProxy(&buf); err != nil {
		s.Error("WriteProxy: %s", err.Error())
		return
	}

	// use httpS to keep all things secure,
	// the second phase of CONNECT also uses this.
	url := fmt.Sprintf("https://%s.appspot.com/http", self.Env.AppSpot)
	// for debug
	if self.Env.AppSpot == "debug" {
		url = "http://localhost:8080/http"
	}

	// post client request as body data
	resp, err := http.Post(url, "application/data", &buf)
	if err != nil {
		s.Error("Post: %s", err.Error())
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusInternalServerError {
		s.Error("Post: %s", resp.Status)
		return
	}

	// the response for the requst of client
	cres, err := http.ReadResponse(bufio.NewReader(resp.Body), r)
	if err != nil {
		s.Error("ReadResponse: %s", err.Error())
		return
	}
	defer cres.Body.Close()

	// copy headers
	CopyResponseHeader(w, cres)

	// please prepare header first and write them
	w.WriteHeader(cres.StatusCode)

	_, err = io.Copy(w, cres.Body)
	if err != nil {
		s.Error("Copy: %s", err.Error())
		return
	}

	s.Info("RESPONSE %s %s %s", r.URL.Host, resp.Status, time.Since(start).String())
}

//  Impossible to connect gae and handle it as a normal TCP connection?
//  GAE only provide http handlers? At least I don't know how to handle to TCP connection on GAE server.
//  NOTE: GAE socket service can only be available for billing users. So free users is unable to use the
//  long term connection. And do what we did in EngineDirect.
//  So we can only use urlfetch.Client.Transport.RoundTrip to do http or https method.
//  Generally, the CONNECT method can be used for any purpose for the advantage of TCP connection.
//  The proxy doesn't need to know what the real underlying protocol or what it is, just need to copy
//  data from client to server, and copy the response from the server to client without any interpret.
//  Now what we can do and had been done by some GAE proxies is that, extract the underlying protocol!!!
//  GAE can only handle limited protocols with urlfetch module, such as http and https.
//  Use Hijacker to get the underlying connection
//
// Data flow:
//  1. Detect host and port
//  2. Hijack the client connection
//  3. Dial self
//  4. Return 200 OK if is successfully
//  5. Get cached or create new signed certificate
//  6. Wrap client connection with TLS and make handshake
//  7. Receive http request
//  8. Write request as a proxy request to self, HTTP handler
//  9. Copy response to client...
func (self *EngineGAE) Connect(s *Session) {
	w, r := s.ResponseWriter, s.Request
	if r.Method != "CONNECT" {
		s.Error("this function can only handle CONNECT method")
		return
	}
	start := time.Now()

	// Only support HTTPS protocol, which is connected with port 443
	host, port, err := net.SplitHostPort(r.URL.Host)
	if err != nil {
		s.Error("SplitHostPort: %s", err.Error())
		return
	}

	if port != "443" {
		s.Error("unsupported CONNECT port: %s", port)
		return
	}

	// hijack the connection to make SSL handshake
	hij, ok := w.(http.Hijacker)
	if !ok {
		s.Error("Server does not support Hijacker")
		return
	}

	conn, _, err := hij.Hijack()
	if err != nil {
		s.Error("Hijack: %s", err.Error())
		return
	}
	defer conn.Close()

	// dial self to transport application data, http request
	rconn, err := net.Dial("tcp", self.Env.Addr)
	if err != nil {
		s.Error("Dial: %s", err.Error())
		return
	}
	defer rconn.Close()

	// Once connected successfully, return OK
	conn.Write([]byte("HTTP/1.1 200 OK\r\n\r\n"))

	// get the fake cert, every host should have its own cert
	cert, err := self.GetCert(s, host)
	if err != nil {
		s.Error("GetCert: %s", err.Error())
		return
	}

	// assume the protocol of client connection is HTTPS
	// wrap it with TSL server
	config := &tls.Config{
		Certificates: []tls.Certificate{*cert},
		ServerName:   host,
	}
	sconn := tls.Server(conn, config)
	defer sconn.Close()

	// The TLS connection goes here
	if err := sconn.Handshake(); err != nil {
		// re-open browser to recover the handshake error:
		//    remote error: bad certificate
		s.Error("Handshake: %s", err.Error())
		return
	}

	// finally, we are at application layer, http request comes
	var wg sync.WaitGroup
	wg.Add(2)

	// read all requests, tls connection reues?
	go func() {
		for {
			req, err := http.ReadRequest(bufio.NewReader(sconn))
			if err != nil {
				if err != io.EOF {
					s.Error("ReadRequest: %s", err.Error())
				}
				break
			}

			// should re-wrap the URL with scheme "https://"
			req.URL, err = url.Parse("https://" + host + req.URL.String())
			req.Header.Set("Mallory-Session", strconv.FormatInt(s.ID, 10))

			// Now re-write the client request to self, HTTP handler
			err = req.WriteProxy(rconn)
			if err != nil {
				s.Error("WriteProxy: %s", err.Error())
				break
			}

			// close the persistent connection after reply the requset
			if req.Close {
				break
			}
		}
		wg.Done()
	}()

	// write back all responses
	go func() {
		_, err := io.Copy(sconn, rconn)
		if err != nil {
			s.Error("Copy: %s", err.Error())
		}
		wg.Done()
	}()

	// Keep connection until no data received from both client and server
	wg.Wait()

	s.Info("CLOSE %s %s", r.URL.Host, time.Since(start).String())
}

func (self *EngineGAE) GetCert(s *Session, host string) (cert *tls.Certificate, err error) {
	// firstly, try to find in memory
	cert = self.Certs.GetSafe(host)
	if cert != nil {
		return
	}

	// secondly, try to find on disk
	crtnam := path.Join(self.CertsDir, host+".crt")
	// we use the same key with CA
	crt, err := tls.LoadX509KeyPair(crtnam, self.Env.Key)
	if err == nil {
		cert = &crt
		self.Certs.AddSafe(host, cert)
		return
	} else if !os.IsNotExist(err) {
		s.Warn("LoadX509KeyPair: %s", err.Error())
	}

	// finally, try to create a new cert
	sn := sha1.Sum([]byte(host))
	config := &CertConfig{
		SerialNumber: new(big.Int).SetBytes(sn[:]),
		CommonName:   host, // FIXME: common name mismatch
	}
	cert, err = CreateSignedCert(self.RootCA, config)
	if err == nil {
		// add to memory
		self.Certs.AddSafe(host, cert)
		// save cert, fail is accepted
		fcrt, _err := os.Create(crtnam)
		if _err == nil {
			for _, c := range cert.Certificate {
				_err = pem.Encode(fcrt, &pem.Block{Type: "CERTIFICATE", Bytes: c})
				if _err != nil {
					break
				}
			}
			fcrt.Close()
			if _err != nil {
				s.Warn("Encode: %s", _err.Error())
				os.Remove(crtnam)
			}
		} else {
			s.Warn("Create %s", _err.Error())
		}
	}
	return
}
<|endoftext|>


<commit_before>/**
随机数处理
create by gloomy 2017-03-29 22:11:23
*/
package common

import "math"

/**
四舍五入取舍
create by gloomy 2017-03-29 22:11:18
需要取舍的浮点数 取舍几位
*/
func Rounding(f float64, n int) float64 {
	pow10_n := math.Pow10(n)
	return math.Trunc((f+0.5/pow10_n)*pow10_n) / pow10_n
}
<commit_msg>增加整数相除取舍 增加相除取获得百分比取小数<commit_after>/**
随机数处理
create by gloomy 2017-03-29 22:11:23
*/
package common

import "math"

/**
四舍五入取舍
create by gloomy 2017-03-29 22:11:18
需要取舍的浮点数 取舍几位
*/
func Rounding(f float64, n int) float64 {
	pow10_n := math.Pow10(n)
	return math.Trunc((f+0.5/pow10_n)*pow10_n) / pow10_n
}

/**
四舍五入取舍
create by gloomy 2017-03-29 22:18:46
除数 被除数 取舍几位
*/
func RoundingByInt(number, subNumber, n int) float64 {
	f := float64(number) / float64(subNumber)
	pow10_n := math.Pow10(n)
	return math.Trunc((f+0.5/pow10_n)*pow10_n) / pow10_n
}

/**
四舍五入取舍 百分比
create by gloomy 2017-03-29 22:18:46
除数 被除数 取舍几位
*/
func RoundingPercentageByInt(number, subNumber, n int) float64 {
	f := float64(number) / float64(subNumber)
	n += 2
	pow10_n := math.Pow10(n)
	percentage := math.Trunc((f+0.5/pow10_n)*pow10_n) / pow10_n
	return percentage * 100
}
<|endoftext|>


<commit_before>package sparse

import (
	"testing"

	"gonum.org/v1/gonum/mat"
)

// A:
// 1, 0, 2, 0,
// 0, 0, 0, 0,
// 0, 3, 4, 5,
var matrixPermutationsForA = []struct {
	name   string
	matrix BlasCompatibleSparser
}{
	{
		name: "CSR",
		matrix: NewCSR(
			3, 4,
			[]int{0, 2, 2, 5},
			[]int{0, 2, 1, 2, 3},
			[]float64{1, 2, 3, 4, 5}),
	},
	{
		name: "CSC",
		matrix: NewCSC(
			3, 4,
			[]int{0, 1, 2, 4, 5},
			[]int{0, 2, 0, 2, 2},
			[]float64{1, 3, 2, 4, 5}),
	},
	{
		name: "COO",
		matrix: NewCOO(
			3, 4,
			[]int{0, 2, 0, 2, 2},
			[]int{0, 3, 2, 1, 2},
			[]float64{1, 5, 2, 3, 4}),
	},
}

func TestMulMatVec(t *testing.T) {
	permsB := []struct {
		name   string
		vector mat.Vector
	}{
		{
			name:   "Dense",
			vector: mat.NewVecDense(4, []float64{1, 2, 0, 3}),
		},
		{
			name:   "Sparse",
			vector: NewVector(4, []int{0, 1, 3}, []float64{1, 2, 3}),
		},
		{
			name:   "Transposed Dense",
			vector: mat.NewVecDense(4, []float64{1, 2, 0, 3}).TVec(),
		},
	}

	tests := []struct {
		alpha  float64
		y      *mat.VecDense
		er, ec int
		eData  []float64
	}{
		{ // y = 0 * A * x
			alpha: 0,
			y:     mat.NewVecDense(3, []float64{0, 0, 0}),
			er:    3, ec: 1,
			eData: []float64{0, 0, 0},
		},
		{ // y = 1 * A * x
			alpha: 1,
			y:     mat.NewVecDense(3, []float64{0, 0, 0}),
			er:    3, ec: 1,
			eData: []float64{1, 0, 21},
		},
		{ // y = 2 * A * x
			alpha: 2,
			y:     mat.NewVecDense(3, []float64{0, 0, 0}),
			er:    3, ec: 1,
			eData: []float64{2, 0, 42},
		},
		{ // y = 1 * A * x + y
			alpha: 1,
			y:     mat.NewVecDense(3, []float64{1, 2, 0}),
			er:    3, ec: 1,
			eData: []float64{2, 2, 21},
		},
		{ // y = 1 * A * x + y
			alpha: 1,
			y:     nil,
			er:    3, ec: 1,
			eData: []float64{1, 0, 21},
		},
	}

	for ti, test := range tests {
		for _, transA := range []bool{false, true} {
			for _, b := range permsB {
				for _, a := range matrixPermutationsForA {
					amat := a.matrix
					var transInd string
					if transA {
						amat = amat.T().(BlasCompatibleSparser)
						transInd = "^T"
					}
					var ycopy *mat.VecDense
					if test.y != nil {
						ycopy = mat.VecDenseCopyOf(test.y)
					}
					y := MulMatVec(transA, test.alpha, amat, b.vector, ycopy)

					yr, yc := y.Dims()
					if yr != test.er || yc != test.ec {
						t.Errorf("Test %d (%s%s x %s): Incorrect dimensions: expected %d x %d but received %d x %d", ti+1, a.name, transInd, b.name, test.er, test.ec, yr, yc)
					}

					yraw := y.RawVector()

					for i, v := range test.eData {
						if v != yraw.Data[i] {
							e := mat.NewDense(test.er, test.ec, test.eData)
							t.Errorf("Test %d (%s%s x %s): Failed, expected\n%v\n but received \n%v", ti+1, a.name, transInd, b.name, mat.Formatted(e), mat.Formatted(y))
							break
						}
					}
				}
			}
		}
	}
}

func TestMulMatMat(t *testing.T) {
	permsB := []struct {
		name   string
		matrix mat.Matrix
	}{
		{
			name: "Dense",
			matrix: mat.NewDense(4, 5, []float64{
				1, 2, 3, 4, 5,
				6, 7, 8, 9, 1,
				2, 3, 4, 5, 6,
				7, 8, 9, 0, 1,
			}),
		},
		{
			name: "COO",
			matrix: NewCOO(4, 5,
				[]int{
					0, 0, 0, 0, 0,
					1, 1, 1, 1, 1,
					2, 2, 2, 2, 2,
					3, 3, 3, 3,
				},
				[]int{
					0, 1, 2, 3, 4,
					0, 1, 2, 3, 4,
					0, 1, 2, 3, 4,
					0, 1, 2, 4,
				},
				[]float64{
					1, 2, 3, 4, 5,
					6, 7, 8, 9, 1,
					2, 3, 4, 5, 6,
					7, 8, 9, 1,
				}),
		},
		{
			name: "CSR",
			matrix: NewCSR(4, 5,
				[]int{0, 5, 10, 15, 19},
				[]int{
					0, 1, 2, 3, 4,
					0, 1, 2, 3, 4,
					0, 1, 2, 3, 4,
					0, 1, 2, 4,
				},
				[]float64{
					1, 2, 3, 4, 5,
					6, 7, 8, 9, 1,
					2, 3, 4, 5, 6,
					7, 8, 9, 1,
				}),
		},
		{
			name: "CSC",
			matrix: NewCSC(
				4, 5,
				[]int{0, 4, 8, 12, 15, 19},
				[]int{
					0, 1, 2, 3,
					0, 1, 2, 3,
					0, 1, 2, 3,
					0, 1, 2,
					0, 1, 2, 3,
				},
				[]float64{
					1, 6, 2, 7,
					2, 7, 3, 8,
					3, 8, 4, 9,
					4, 9, 5,
					5, 1, 6, 1,
				},
			),
		},
	}

	tests := []struct {
		alpha  float64
		c      *mat.Dense
		er, ec int
		eData  []float64
	}{
		{ // C = 0 * A * B
			alpha: 0,
			c: mat.NewDense(3, 5, []float64{
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
			}),
			er: 3, ec: 5,
			eData: []float64{
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
			},
		},
		{ // C = 1 * A * B
			alpha: 1,
			c: mat.NewDense(3, 5, []float64{
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
			}),
			er: 3, ec: 5,
			eData: []float64{
				5, 8, 11, 14, 17,
				0, 0, 0, 0, 0,
				61, 73, 85, 47, 32,
			},
		},
		{ // C = 2 * A * B
			alpha: 2,
			c: mat.NewDense(3, 5, []float64{
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
			}),
			er: 3, ec: 5,
			eData: []float64{
				10, 16, 22, 28, 34,
				0, 0, 0, 0, 0,
				122, 146, 170, 94, 64,
			},
		},
		{ // C = 1 * A * B + C
			alpha: 1,
			c: mat.NewDense(3, 5, []float64{
				1, 0, 5, 0, 8,
				2, 0, 3, 0, 4,
				0, 0, 0, 7, 0,
			}),
			er: 3, ec: 5,
			eData: []float64{
				6, 8, 16, 14, 25,
				2, 0, 3, 0, 4,
				61, 73, 85, 54, 32,
			},
		},
		{ // C = 1 * A * B + C
			alpha: 1,
			c:     nil,
			er:    3, ec: 5,
			eData: []float64{
				5, 8, 11, 14, 17,
				0, 0, 0, 0, 0,
				61, 73, 85, 47, 32,
			},
		},
	}

	for ti, test := range tests {
		for _, transA := range []bool{false, true} {
			for _, b := range permsB {
				for _, a := range matrixPermutationsForA {
					amat := a.matrix
					var transInd string
					if transA {
						amat = amat.T().(BlasCompatibleSparser)
						transInd = "^T"
					}
					var ccopy *mat.Dense
					if test.c != nil {
						ccopy = mat.DenseCopyOf(test.c)
					}
					c := MulMatMat(transA, test.alpha, amat, b.matrix, ccopy)

					cr, cc := c.Dims()
					if cr != test.er || cc != test.ec {
						t.Errorf("Test %d (%s%s x %s): Incorrect dimensions: expected %d x %d but received %d x %d", ti+1, a.name, transInd, b.name, test.er, test.ec, cr, cc)
					}

					craw := c.RawMatrix()
					for i, v := range test.eData {
						if v != craw.Data[i] {
							e := mat.NewDense(test.er, test.ec, test.eData)
							t.Errorf("Test %d (%s%s x %s): Failed, expected\n%v\n but received \n%v", ti+1, a.name, transInd, b.name, mat.Formatted(e), mat.Formatted(c))
							break
						}
					}
				}
			}
		}
	}
}
<commit_msg>more testing for increased code coverage<commit_after>package sparse

import (
	"testing"

	"gonum.org/v1/gonum/mat"
)

// A:
// 1, 0, 2, 0,
// 0, 0, 0, 0,
// 0, 3, 4, 5,
var matrixPermutationsForA = []struct {
	name   string
	matrix BlasCompatibleSparser
}{
	{
		name: "CSR",
		matrix: NewCSR(
			3, 4,
			[]int{0, 2, 2, 5},
			[]int{0, 2, 1, 2, 3},
			[]float64{1, 2, 3, 4, 5}),
	},
	{
		name: "CSC",
		matrix: NewCSC(
			3, 4,
			[]int{0, 1, 2, 4, 5},
			[]int{0, 2, 0, 2, 2},
			[]float64{1, 3, 2, 4, 5}),
	},
	{
		name: "COO",
		matrix: NewCOO(
			3, 4,
			[]int{0, 2, 0, 2, 2},
			[]int{0, 3, 2, 1, 2},
			[]float64{1, 5, 2, 3, 4}),
	},
}

func TestMulMatVec(t *testing.T) {
	permsB := []struct {
		name   string
		vector mat.Vector
	}{
		{
			name:   "Dense",
			vector: mat.NewVecDense(4, []float64{1, 2, 0, 3}),
		},
		{
			name:   "Sparse",
			vector: NewVector(4, []int{0, 1, 3}, []float64{1, 2, 3}),
		},
		{
			name:   "Transposed Dense",
			vector: mat.NewVecDense(4, []float64{1, 2, 0, 3}).TVec(),
		},
	}

	tests := []struct {
		alpha  float64
		y      *mat.VecDense
		er, ec int
		eData  []float64
	}{
		{ // y = 0 * A * x
			alpha: 0,
			y:     mat.NewVecDense(3, []float64{0, 0, 0}),
			er:    3, ec: 1,
			eData: []float64{0, 0, 0},
		},
		{ // y = 1 * A * x
			alpha: 1,
			y:     mat.NewVecDense(3, []float64{0, 0, 0}),
			er:    3, ec: 1,
			eData: []float64{1, 0, 21},
		},
		{ // y = 2 * A * x
			alpha: 2,
			y:     mat.NewVecDense(3, []float64{0, 0, 0}),
			er:    3, ec: 1,
			eData: []float64{2, 0, 42},
		},
		{ // y = 1 * A * x + y
			alpha: 1,
			y:     mat.NewVecDense(3, []float64{1, 2, 0}),
			er:    3, ec: 1,
			eData: []float64{2, 2, 21},
		},
		{ // y = 1 * A * x + y
			alpha: 1,
			y:     nil,
			er:    3, ec: 1,
			eData: []float64{1, 0, 21},
		},
	}

	for ti, test := range tests {
		for _, transA := range []bool{false, true} {
			for _, b := range permsB {
				for _, a := range matrixPermutationsForA {
					amat := a.matrix
					var transInd string
					if transA {
						amat = amat.T().(BlasCompatibleSparser)
						transInd = "^T"
					}
					var ycopy *mat.VecDense
					if test.y != nil {
						ycopy = mat.VecDenseCopyOf(test.y)
					}
					y := MulMatVec(transA, test.alpha, amat, b.vector, ycopy)

					yr, yc := y.Dims()
					if yr != test.er || yc != test.ec {
						t.Errorf("Test %d (%s%s x %s): Incorrect dimensions: expected %d x %d but received %d x %d", ti+1, a.name, transInd, b.name, test.er, test.ec, yr, yc)
					}

					yraw := y.RawVector()

					for i, v := range test.eData {
						if v != yraw.Data[i] {
							e := mat.NewDense(test.er, test.ec, test.eData)
							t.Errorf("Test %d (%s%s x %s): Failed, expected\n%v\n but received \n%v", ti+1, a.name, transInd, b.name, mat.Formatted(e), mat.Formatted(y))
							break
						}
					}
				}
			}
		}
	}
}

func TestMulMatMat(t *testing.T) {
	permsB := []struct {
		name   string
		matrix mat.Matrix
	}{
		{
			name: "Dense",
			matrix: mat.NewDense(4, 5, []float64{
				1, 2, 3, 4, 5,
				6, 7, 8, 9, 1,
				2, 3, 4, 5, 6,
				7, 8, 9, 0, 1,
			}),
		},
		{
			name: "COO",
			matrix: NewCOO(4, 5,
				[]int{
					0, 0, 0, 0, 0,
					1, 1, 1, 1, 1,
					2, 2, 2, 2, 2,
					3, 3, 3, 3,
				},
				[]int{
					0, 1, 2, 3, 4,
					0, 1, 2, 3, 4,
					0, 1, 2, 3, 4,
					0, 1, 2, 4,
				},
				[]float64{
					1, 2, 3, 4, 5,
					6, 7, 8, 9, 1,
					2, 3, 4, 5, 6,
					7, 8, 9, 1,
				}),
		},
		{
			name: "CSR",
			matrix: NewCSR(4, 5,
				[]int{0, 5, 10, 15, 19},
				[]int{
					0, 1, 2, 3, 4,
					0, 1, 2, 3, 4,
					0, 1, 2, 3, 4,
					0, 1, 2, 4,
				},
				[]float64{
					1, 2, 3, 4, 5,
					6, 7, 8, 9, 1,
					2, 3, 4, 5, 6,
					7, 8, 9, 1,
				}),
		},
		{
			name: "CSC",
			matrix: NewCSC(
				4, 5,
				[]int{0, 4, 8, 12, 15, 19},
				[]int{
					0, 1, 2, 3,
					0, 1, 2, 3,
					0, 1, 2, 3,
					0, 1, 2,
					0, 1, 2, 3,
				},
				[]float64{
					1, 6, 2, 7,
					2, 7, 3, 8,
					3, 8, 4, 9,
					4, 9, 5,
					5, 1, 6, 1,
				},
			),
		},
	}

	tests := []struct {
		alpha  float64
		c      *mat.Dense
		er, ec int
		eData  []float64
	}{
		{ // C = 0 * A * B
			alpha: 0,
			c: mat.NewDense(3, 5, []float64{
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
			}),
			er: 3, ec: 5,
			eData: []float64{
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
			},
		},
		{ // C = 1 * A * B
			alpha: 1,
			c: mat.NewDense(3, 5, []float64{
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
			}),
			er: 3, ec: 5,
			eData: []float64{
				5, 8, 11, 14, 17,
				0, 0, 0, 0, 0,
				61, 73, 85, 47, 32,
			},
		},
		{ // C = 2 * A * B
			alpha: 2,
			c: mat.NewDense(3, 5, []float64{
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
				0, 0, 0, 0, 0,
			}),
			er: 3, ec: 5,
			eData: []float64{
				10, 16, 22, 28, 34,
				0, 0, 0, 0, 0,
				122, 146, 170, 94, 64,
			},
		},
		{ // C = 1 * A * B + C
			alpha: 1,
			c: mat.NewDense(3, 5, []float64{
				1, 0, 5, 0, 8,
				2, 0, 3, 0, 4,
				0, 0, 0, 7, 0,
			}),
			er: 3, ec: 5,
			eData: []float64{
				6, 8, 16, 14, 25,
				2, 0, 3, 0, 4,
				61, 73, 85, 54, 32,
			},
		},
		{ // C = 1 * A * B + C
			alpha: 1,
			c:     nil,
			er:    3, ec: 5,
			eData: []float64{
				5, 8, 11, 14, 17,
				0, 0, 0, 0, 0,
				61, 73, 85, 47, 32,
			},
		},
	}

	for ti, test := range tests {
		for _, transA := range []bool{false, true} {
			for _, b := range permsB {
				for _, a := range matrixPermutationsForA {
					amat := a.matrix
					var transInd string
					if transA {
						amat = amat.T().(BlasCompatibleSparser)
						transInd = "^T"
					}
					var ccopy *mat.Dense
					var cPrimecopy *mat.Dense
					if test.c != nil {
						ccopy = mat.DenseCopyOf(test.c)
						cPrimecopy = mat.DenseCopyOf(test.c)
					}
					c := MulMatMat(transA, test.alpha, amat, b.matrix, ccopy)
					cPrime := GenericMulMatMat(transA, test.alpha, amat, b.matrix, cPrimecopy)

					cr, cc := c.Dims()
					if cr != test.er || cc != test.ec {
						t.Errorf("Test %d (%s%s x %s): Incorrect dimensions: expected %d x %d but received %d x %d", ti+1, a.name, transInd, b.name, test.er, test.ec, cr, cc)
					}

					craw := c.RawMatrix()
					for i, v := range test.eData {
						if v != craw.Data[i] {
							e := mat.NewDense(test.er, test.ec, test.eData)
							t.Errorf("Test %d (%s%s x %s): Failed, expected\n%v\n but received \n%v", ti+1, a.name, transInd, b.name, mat.Formatted(e), mat.Formatted(c))
							break
						}
					}
					for row := 0; row < cr; row++ {
						for col := 0; col < cc; col++ {
							if cPrime.At(row, col) != c.At(row, col) {
								e := mat.NewDense(test.er, test.ec, test.eData)
								t.Errorf("Test %d (%s%s x %s): Failed, expected\n%v\n but received \n%v", ti+1, a.name, transInd, b.name, mat.Formatted(e), mat.Formatted(c))
								break
							}
						}
					}
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package md2man

import (
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/russross/blackfriday"
)

// roffRenderer implements the blackfriday.Renderer interface for creating
// roff format (manpages) from markdown text
type roffRenderer struct {
	extensions   blackfriday.Extensions
	listCounters []int
	firstHeader  bool
	defineTerm   bool
	listDepth    int
}

const (
	titleHeader      = ".TH "
	topLevelHeader   = "\n\n.SH "
	secondLevelHdr   = "\n.SH "
	otherHeader      = "\n.SS "
	crTag            = "\n"
	emphTag          = "\\fI"
	emphCloseTag     = "\\fP"
	strongTag        = "\\fB"
	strongCloseTag   = "\\fP"
	breakTag         = "\n.br\n"
	paraTag          = "\n.PP\n"
	hruleTag         = "\n.ti 0\n\\l'\\n(.lu'\n"
	linkTag          = "\n\\[la]"
	linkCloseTag     = "\\[ra]"
	codespanTag      = "\\fB\\fC"
	codespanCloseTag = "\\fR"
	codeTag          = "\n.PP\n.RS\n\n.nf\n"
	codeCloseTag     = "\n.fi\n.RE\n"
	quoteTag         = "\n.PP\n.RS\n"
	quoteCloseTag    = "\n.RE\n"
	listTag          = "\n.RS\n"
	listCloseTag     = "\n.RE\n"
	arglistTag       = "\n.TP\n"
	tableStart       = "\n.TS\nallbox;\n"
	tableEnd         = ".TE\n"
	tableCellStart   = "T{\n"
	tableCellEnd     = "\nT}\n"
)

// NewRoffRenderer creates a new blackfriday Renderer for generating roff documents
// from markdown
func NewRoffRenderer() *roffRenderer { // nolint: golint
	var extensions blackfriday.Extensions

	extensions |= blackfriday.NoIntraEmphasis
	extensions |= blackfriday.Tables
	extensions |= blackfriday.FencedCode
	extensions |= blackfriday.SpaceHeadings
	extensions |= blackfriday.Footnotes
	extensions |= blackfriday.Titleblock
	extensions |= blackfriday.DefinitionLists
	return &roffRenderer{
		extensions: extensions,
	}
}

// GetExtensions returns the list of extensions used by this renderer implementation
func (r *roffRenderer) GetExtensions() blackfriday.Extensions {
	return r.extensions
}

// RenderHeader handles outputting the header at document start
func (r *roffRenderer) RenderHeader(w io.Writer, ast *blackfriday.Node) {
	// disable hyphenation
	out(w, ".nh\n")
}

// RenderFooter handles outputting the footer at the document end; the roff
// renderer has no footer information
func (r *roffRenderer) RenderFooter(w io.Writer, ast *blackfriday.Node) {
}

// RenderNode is called for each node in a markdown document; based on the node
// type the equivalent roff output is sent to the writer
func (r *roffRenderer) RenderNode(w io.Writer, node *blackfriday.Node, entering bool) blackfriday.WalkStatus {

	var walkAction = blackfriday.GoToNext

	switch node.Type {
	case blackfriday.Text:
		r.handleText(w, node, entering)
	case blackfriday.Softbreak:
		out(w, crTag)
	case blackfriday.Hardbreak:
		out(w, breakTag)
	case blackfriday.Emph:
		if entering {
			out(w, emphTag)
		} else {
			out(w, emphCloseTag)
		}
	case blackfriday.Strong:
		if entering {
			out(w, strongTag)
		} else {
			out(w, strongCloseTag)
		}
	case blackfriday.Link:
		if entering {
			out(w, linkTag+string(node.LinkData.Destination))
		} else {
			out(w, linkCloseTag)
		}
	case blackfriday.Image:
		// ignore images
		walkAction = blackfriday.SkipChildren
	case blackfriday.Code:
		out(w, codespanTag)
		escapeSpecialChars(w, node.Literal)
		out(w, codespanCloseTag)
	case blackfriday.Document:
		break
	case blackfriday.Paragraph:
		// roff .PP markers break lists
		if r.listDepth > 0 {
			return blackfriday.GoToNext
		}
		if entering {
			out(w, paraTag)
		} else {
			out(w, crTag)
		}
	case blackfriday.BlockQuote:
		if entering {
			out(w, quoteTag)
		} else {
			out(w, quoteCloseTag)
		}
	case blackfriday.Heading:
		r.handleHeading(w, node, entering)
	case blackfriday.HorizontalRule:
		out(w, hruleTag)
	case blackfriday.List:
		r.handleList(w, node, entering)
	case blackfriday.Item:
		r.handleItem(w, node, entering)
	case blackfriday.CodeBlock:
		out(w, codeTag)
		escapeSpecialChars(w, node.Literal)
		out(w, codeCloseTag)
	case blackfriday.Table:
		r.handleTable(w, node, entering)
	case blackfriday.TableCell:
		r.handleTableCell(w, node, entering)
	case blackfriday.TableHead:
	case blackfriday.TableBody:
	case blackfriday.TableRow:
		// no action as cell entries do all the nroff formatting
		return blackfriday.GoToNext
	default:
		fmt.Fprintf(os.Stderr, "WARNING: go-md2man does not handle node type "+node.Type.String())
	}
	return walkAction
}

func (r *roffRenderer) handleText(w io.Writer, node *blackfriday.Node, entering bool) {
	var (
		start, end string
	)
	// handle special roff table cell text encapsulation
	if node.Parent.Type == blackfriday.TableCell {
		if len(node.Literal) > 30 {
			start = tableCellStart
			end = tableCellEnd
		} else {
			// end rows that aren't terminated by "tableCellEnd" with a cr if end of row
			if node.Parent.Next == nil && !node.Parent.IsHeader {
				end = crTag
			}
		}
	}
	out(w, start)
	escapeSpecialChars(w, node.Literal)
	out(w, end)
}

func (r *roffRenderer) handleHeading(w io.Writer, node *blackfriday.Node, entering bool) {
	if entering {
		switch node.Level {
		case 1:
			if !r.firstHeader {
				out(w, titleHeader)
				r.firstHeader = true
				break
			}
			out(w, topLevelHeader)
		case 2:
			out(w, secondLevelHdr)
		default:
			out(w, otherHeader)
		}
	}
}

func (r *roffRenderer) handleList(w io.Writer, node *blackfriday.Node, entering bool) {
	openTag := listTag
	closeTag := listCloseTag
	if node.ListFlags&blackfriday.ListTypeDefinition != 0 {
		// tags for definition lists handled within Item node
		openTag = ""
		closeTag = ""
	}
	if entering {
		r.listDepth++
		if node.ListFlags&blackfriday.ListTypeOrdered != 0 {
			r.listCounters = append(r.listCounters, 1)
		}
		out(w, openTag)
	} else {
		if node.ListFlags&blackfriday.ListTypeOrdered != 0 {
			r.listCounters = r.listCounters[:len(r.listCounters)-1]
		}
		out(w, closeTag)
		r.listDepth--
	}
}

func (r *roffRenderer) handleItem(w io.Writer, node *blackfriday.Node, entering bool) {
	if entering {
		if node.ListFlags&blackfriday.ListTypeOrdered != 0 {
			out(w, fmt.Sprintf(".IP \"%3d.\" 5\n", r.listCounters[len(r.listCounters)-1]))
			r.listCounters[len(r.listCounters)-1]++
		} else if node.ListFlags&blackfriday.ListTypeDefinition != 0 {
			// state machine for handling terms and following definitions
			// since blackfriday does not distinguish them properly, nor
			// does it seperate them into separate lists as it should
			if !r.defineTerm {
				out(w, arglistTag)
				r.defineTerm = true
			} else {
				r.defineTerm = false
			}
		} else {
			out(w, ".IP \\(bu 2\n")
		}
	} else {
		out(w, "\n")
	}
}

func (r *roffRenderer) handleTable(w io.Writer, node *blackfriday.Node, entering bool) {
	if entering {
		out(w, tableStart)
		//call walker to count cells (and rows?) so format section can be produced
		columns := countColumns(node)
		out(w, strings.Repeat("l ", columns)+"\n")
		out(w, strings.Repeat("l ", columns)+".\n")
	} else {
		out(w, tableEnd)
	}
}

func (r *roffRenderer) handleTableCell(w io.Writer, node *blackfriday.Node, entering bool) {
	var (
		start, end string
	)
	if node.IsHeader {
		start = codespanTag
		end = codespanCloseTag
	}
	if entering {
		if node.Prev != nil && node.Prev.Type == blackfriday.TableCell {
			out(w, "\t"+start)
		} else {
			out(w, start)
		}
	} else {
		// need to carriage return if we are at the end of the header row
		if node.IsHeader && node.Next == nil {
			end = end + crTag
		}
		out(w, end)
	}
}

// because roff format requires knowing the column count before outputting any table
// data we need to walk a table tree and count the columns
func countColumns(node *blackfriday.Node) int {
	var columns int

	node.Walk(func(node *blackfriday.Node, entering bool) blackfriday.WalkStatus {
		switch node.Type {
		case blackfriday.TableRow:
			if !entering {
				return blackfriday.Terminate
			}
		case blackfriday.TableCell:
			if entering {
				columns++
			}
		default:
		}
		return blackfriday.GoToNext
	})
	return columns
}

func out(w io.Writer, output string) {
	io.WriteString(w, output) // nolint: errcheck
}

func needsBackslash(c byte) bool {
	for _, r := range []byte("-_&\\~") {
		if c == r {
			return true
		}
	}
	return false
}

func escapeSpecialChars(w io.Writer, text []byte) {
	for i := 0; i < len(text); i++ {
		// escape initial apostrophe or period
		if len(text) >= 1 && (text[0] == '\'' || text[0] == '.') {
			out(w, "\\&")
		}

		// directly copy normal characters
		org := i

		for i < len(text) && !needsBackslash(text[i]) {
			i++
		}
		if i > org {
			w.Write(text[org:i]) // nolint: errcheck
		}

		// escape a character
		if i >= len(text) {
			break
		}

		w.Write([]byte{'\\', text[i]}) // nolint: errcheck
	}
}
<commit_msg>RenderNode: fix error printing<commit_after>package md2man

import (
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/russross/blackfriday"
)

// roffRenderer implements the blackfriday.Renderer interface for creating
// roff format (manpages) from markdown text
type roffRenderer struct {
	extensions   blackfriday.Extensions
	listCounters []int
	firstHeader  bool
	defineTerm   bool
	listDepth    int
}

const (
	titleHeader      = ".TH "
	topLevelHeader   = "\n\n.SH "
	secondLevelHdr   = "\n.SH "
	otherHeader      = "\n.SS "
	crTag            = "\n"
	emphTag          = "\\fI"
	emphCloseTag     = "\\fP"
	strongTag        = "\\fB"
	strongCloseTag   = "\\fP"
	breakTag         = "\n.br\n"
	paraTag          = "\n.PP\n"
	hruleTag         = "\n.ti 0\n\\l'\\n(.lu'\n"
	linkTag          = "\n\\[la]"
	linkCloseTag     = "\\[ra]"
	codespanTag      = "\\fB\\fC"
	codespanCloseTag = "\\fR"
	codeTag          = "\n.PP\n.RS\n\n.nf\n"
	codeCloseTag     = "\n.fi\n.RE\n"
	quoteTag         = "\n.PP\n.RS\n"
	quoteCloseTag    = "\n.RE\n"
	listTag          = "\n.RS\n"
	listCloseTag     = "\n.RE\n"
	arglistTag       = "\n.TP\n"
	tableStart       = "\n.TS\nallbox;\n"
	tableEnd         = ".TE\n"
	tableCellStart   = "T{\n"
	tableCellEnd     = "\nT}\n"
)

// NewRoffRenderer creates a new blackfriday Renderer for generating roff documents
// from markdown
func NewRoffRenderer() *roffRenderer { // nolint: golint
	var extensions blackfriday.Extensions

	extensions |= blackfriday.NoIntraEmphasis
	extensions |= blackfriday.Tables
	extensions |= blackfriday.FencedCode
	extensions |= blackfriday.SpaceHeadings
	extensions |= blackfriday.Footnotes
	extensions |= blackfriday.Titleblock
	extensions |= blackfriday.DefinitionLists
	return &roffRenderer{
		extensions: extensions,
	}
}

// GetExtensions returns the list of extensions used by this renderer implementation
func (r *roffRenderer) GetExtensions() blackfriday.Extensions {
	return r.extensions
}

// RenderHeader handles outputting the header at document start
func (r *roffRenderer) RenderHeader(w io.Writer, ast *blackfriday.Node) {
	// disable hyphenation
	out(w, ".nh\n")
}

// RenderFooter handles outputting the footer at the document end; the roff
// renderer has no footer information
func (r *roffRenderer) RenderFooter(w io.Writer, ast *blackfriday.Node) {
}

// RenderNode is called for each node in a markdown document; based on the node
// type the equivalent roff output is sent to the writer
func (r *roffRenderer) RenderNode(w io.Writer, node *blackfriday.Node, entering bool) blackfriday.WalkStatus {

	var walkAction = blackfriday.GoToNext

	switch node.Type {
	case blackfriday.Text:
		r.handleText(w, node, entering)
	case blackfriday.Softbreak:
		out(w, crTag)
	case blackfriday.Hardbreak:
		out(w, breakTag)
	case blackfriday.Emph:
		if entering {
			out(w, emphTag)
		} else {
			out(w, emphCloseTag)
		}
	case blackfriday.Strong:
		if entering {
			out(w, strongTag)
		} else {
			out(w, strongCloseTag)
		}
	case blackfriday.Link:
		if entering {
			out(w, linkTag+string(node.LinkData.Destination))
		} else {
			out(w, linkCloseTag)
		}
	case blackfriday.Image:
		// ignore images
		walkAction = blackfriday.SkipChildren
	case blackfriday.Code:
		out(w, codespanTag)
		escapeSpecialChars(w, node.Literal)
		out(w, codespanCloseTag)
	case blackfriday.Document:
		break
	case blackfriday.Paragraph:
		// roff .PP markers break lists
		if r.listDepth > 0 {
			return blackfriday.GoToNext
		}
		if entering {
			out(w, paraTag)
		} else {
			out(w, crTag)
		}
	case blackfriday.BlockQuote:
		if entering {
			out(w, quoteTag)
		} else {
			out(w, quoteCloseTag)
		}
	case blackfriday.Heading:
		r.handleHeading(w, node, entering)
	case blackfriday.HorizontalRule:
		out(w, hruleTag)
	case blackfriday.List:
		r.handleList(w, node, entering)
	case blackfriday.Item:
		r.handleItem(w, node, entering)
	case blackfriday.CodeBlock:
		out(w, codeTag)
		escapeSpecialChars(w, node.Literal)
		out(w, codeCloseTag)
	case blackfriday.Table:
		r.handleTable(w, node, entering)
	case blackfriday.TableCell:
		r.handleTableCell(w, node, entering)
	case blackfriday.TableHead:
	case blackfriday.TableBody:
	case blackfriday.TableRow:
		// no action as cell entries do all the nroff formatting
		return blackfriday.GoToNext
	default:
		fmt.Fprintln(os.Stderr, "WARNING: go-md2man does not handle node type "+node.Type.String())
	}
	return walkAction
}

func (r *roffRenderer) handleText(w io.Writer, node *blackfriday.Node, entering bool) {
	var (
		start, end string
	)
	// handle special roff table cell text encapsulation
	if node.Parent.Type == blackfriday.TableCell {
		if len(node.Literal) > 30 {
			start = tableCellStart
			end = tableCellEnd
		} else {
			// end rows that aren't terminated by "tableCellEnd" with a cr if end of row
			if node.Parent.Next == nil && !node.Parent.IsHeader {
				end = crTag
			}
		}
	}
	out(w, start)
	escapeSpecialChars(w, node.Literal)
	out(w, end)
}

func (r *roffRenderer) handleHeading(w io.Writer, node *blackfriday.Node, entering bool) {
	if entering {
		switch node.Level {
		case 1:
			if !r.firstHeader {
				out(w, titleHeader)
				r.firstHeader = true
				break
			}
			out(w, topLevelHeader)
		case 2:
			out(w, secondLevelHdr)
		default:
			out(w, otherHeader)
		}
	}
}

func (r *roffRenderer) handleList(w io.Writer, node *blackfriday.Node, entering bool) {
	openTag := listTag
	closeTag := listCloseTag
	if node.ListFlags&blackfriday.ListTypeDefinition != 0 {
		// tags for definition lists handled within Item node
		openTag = ""
		closeTag = ""
	}
	if entering {
		r.listDepth++
		if node.ListFlags&blackfriday.ListTypeOrdered != 0 {
			r.listCounters = append(r.listCounters, 1)
		}
		out(w, openTag)
	} else {
		if node.ListFlags&blackfriday.ListTypeOrdered != 0 {
			r.listCounters = r.listCounters[:len(r.listCounters)-1]
		}
		out(w, closeTag)
		r.listDepth--
	}
}

func (r *roffRenderer) handleItem(w io.Writer, node *blackfriday.Node, entering bool) {
	if entering {
		if node.ListFlags&blackfriday.ListTypeOrdered != 0 {
			out(w, fmt.Sprintf(".IP \"%3d.\" 5\n", r.listCounters[len(r.listCounters)-1]))
			r.listCounters[len(r.listCounters)-1]++
		} else if node.ListFlags&blackfriday.ListTypeDefinition != 0 {
			// state machine for handling terms and following definitions
			// since blackfriday does not distinguish them properly, nor
			// does it seperate them into separate lists as it should
			if !r.defineTerm {
				out(w, arglistTag)
				r.defineTerm = true
			} else {
				r.defineTerm = false
			}
		} else {
			out(w, ".IP \\(bu 2\n")
		}
	} else {
		out(w, "\n")
	}
}

func (r *roffRenderer) handleTable(w io.Writer, node *blackfriday.Node, entering bool) {
	if entering {
		out(w, tableStart)
		//call walker to count cells (and rows?) so format section can be produced
		columns := countColumns(node)
		out(w, strings.Repeat("l ", columns)+"\n")
		out(w, strings.Repeat("l ", columns)+".\n")
	} else {
		out(w, tableEnd)
	}
}

func (r *roffRenderer) handleTableCell(w io.Writer, node *blackfriday.Node, entering bool) {
	var (
		start, end string
	)
	if node.IsHeader {
		start = codespanTag
		end = codespanCloseTag
	}
	if entering {
		if node.Prev != nil && node.Prev.Type == blackfriday.TableCell {
			out(w, "\t"+start)
		} else {
			out(w, start)
		}
	} else {
		// need to carriage return if we are at the end of the header row
		if node.IsHeader && node.Next == nil {
			end = end + crTag
		}
		out(w, end)
	}
}

// because roff format requires knowing the column count before outputting any table
// data we need to walk a table tree and count the columns
func countColumns(node *blackfriday.Node) int {
	var columns int

	node.Walk(func(node *blackfriday.Node, entering bool) blackfriday.WalkStatus {
		switch node.Type {
		case blackfriday.TableRow:
			if !entering {
				return blackfriday.Terminate
			}
		case blackfriday.TableCell:
			if entering {
				columns++
			}
		default:
		}
		return blackfriday.GoToNext
	})
	return columns
}

func out(w io.Writer, output string) {
	io.WriteString(w, output) // nolint: errcheck
}

func needsBackslash(c byte) bool {
	for _, r := range []byte("-_&\\~") {
		if c == r {
			return true
		}
	}
	return false
}

func escapeSpecialChars(w io.Writer, text []byte) {
	for i := 0; i < len(text); i++ {
		// escape initial apostrophe or period
		if len(text) >= 1 && (text[0] == '\'' || text[0] == '.') {
			out(w, "\\&")
		}

		// directly copy normal characters
		org := i

		for i < len(text) && !needsBackslash(text[i]) {
			i++
		}
		if i > org {
			w.Write(text[org:i]) // nolint: errcheck
		}

		// escape a character
		if i >= len(text) {
			break
		}

		w.Write([]byte{'\\', text[i]}) // nolint: errcheck
	}
}
<|endoftext|>


<commit_before>package database

import (
	"os"
	"reflect"
	"testing"

	"errors"

	"fmt"

	"github.com/bouk/monkey"
	"github.com/jinzhu/gorm"
	log "github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
)

func TestEnvironmentDatabase(t *testing.T) {
	currentEnv := os.Getenv("CLAIMR_DATABASE")
	os.Unsetenv("CLAIMR_DATABASE")
	defer func() { os.Setenv("CLAIMR_DATABASE", currentEnv) }()

	expectedMsg := "Claimr mysql database string unset. Set CLAIMR_DATABASE to continue."

	mockLogFatal := func(msg ...interface{}) {
		assert.Equal(t, expectedMsg, msg[0])
		panic("log.Fatal called")
	}
	patchLog := monkey.Patch(log.Fatal, mockLogFatal)
	defer patchLog.Unpatch()
	assert.PanicsWithValue(t, "log.Fatal called", initDB, "log.Fatal was not called")
}

func TestDBError(t *testing.T) {
	expectedMsg := "could not create a database connection - [ERROR]"

	mockLogFatal := func(format string, args ...interface{}) {
		assert.Equal(t, expectedMsg, fmt.Sprintf(format, args))
		panic("log.Fatal called")
	}

	mockGormOpen := func(dialect string, args ...interface{}) (db *gorm.DB, err error) {
		return nil, errors.New("ERROR")
	}

	patchLog := monkey.Patch(log.Fatalf, mockLogFatal)
	patchGorm := monkey.Patch(gorm.Open, mockGormOpen)

	defer patchLog.Unpatch()
	defer patchGorm.Unpatch()

	assert.PanicsWithValue(t, "log.Fatal called", initDB, "log.Fatal was not called")
}

func TestDBClose(t *testing.T) {
	mockLogInfo := func(args ...interface{}) {
		assert.Equal(t, "DB Closed", args)
	}

	mockDBClose := func(db *gorm.DB) error {
		return nil
	}

	patchLog := monkey.Patch(log.Info, mockLogInfo)
	patchClose := monkey.PatchInstanceMethod(reflect.TypeOf(DB), "Close", mockDBClose)

	DB.Close()

	patchLog.Unpatch()
	patchClose.Unpatch()
}
<commit_msg>DB Close error<commit_after>package database

import (
	"os"
	"reflect"
	"testing"

	"errors"

	"fmt"

	"github.com/bouk/monkey"
	"github.com/jinzhu/gorm"
	log "github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
)

func TestEnvironmentDatabase(t *testing.T) {
	currentEnv := os.Getenv("CLAIMR_DATABASE")
	os.Unsetenv("CLAIMR_DATABASE")
	defer func() { os.Setenv("CLAIMR_DATABASE", currentEnv) }()

	expectedMsg := "Claimr mysql database string unset. Set CLAIMR_DATABASE to continue."

	mockLogFatal := func(msg ...interface{}) {
		assert.Equal(t, expectedMsg, msg[0])
		panic("log.Fatal called")
	}
	patchLog := monkey.Patch(log.Fatal, mockLogFatal)
	defer patchLog.Unpatch()
	assert.PanicsWithValue(t, "log.Fatal called", initDB, "log.Fatal was not called")
}

func TestDBError(t *testing.T) {
	expectedMsg := "could not create a database connection - [ERROR]"

	mockLogFatal := func(format string, args ...interface{}) {
		assert.Equal(t, expectedMsg, fmt.Sprintf(format, args))
		panic("log.Fatal called")
	}

	mockGormOpen := func(dialect string, args ...interface{}) (db *gorm.DB, err error) {
		return nil, errors.New("ERROR")
	}

	patchLog := monkey.Patch(log.Fatalf, mockLogFatal)
	patchGorm := monkey.Patch(gorm.Open, mockGormOpen)

	defer patchLog.Unpatch()
	defer patchGorm.Unpatch()

	assert.PanicsWithValue(t, "log.Fatal called", initDB, "log.Fatal was not called")
}

func TestDBClose(t *testing.T) {
	mockLogInfo := func(args ...interface{}) {
		assert.Equal(t, "DB Closed", args)
	}

	mockDBClose := func(db *gorm.DB) error {
		return nil
	}

	patchLog := monkey.Patch(log.Info, mockLogInfo)
	patchClose := monkey.PatchInstanceMethod(reflect.TypeOf(DB), "Close", mockDBClose)

	DB.Close()

	patchLog.Unpatch()
	patchClose.Unpatch()
}

func TestDBCloseError(t *testing.T) {
	mockLogWarnf := func(format string, args ...interface{}) {
		assert.Equal(t, "Error while closing database. Please check for memory leak %s", "Simulated error")
	}

	mockDBClose := func(db *gorm.DB) error {
		return fmt.Errorf("Simulated error")
	}

	patchLog := monkey.Patch(log.Warnf, mockLogWarnf)
	patchClose := monkey.PatchInstanceMethod(reflect.TypeOf(DB), "Close", mockDBClose)

	DB.Close()

	patchLog.Unpatch()
	patchClose.Unpatch()
}
<|endoftext|>