<commit_before>// Teleport是一款适用于分布式系统的高并发API框架，它采用socket全双工通信，实现S/C对等工作，支持长、短两种连接模式，支持断开后自动连接与手动断开连接，内部数据传输格式为JSON。
package teleport

import (
	"encoding/json"
	"github.com/henrylee2cn/teleport/debug"
	"log"
	"time"
)

// mode
const (
	SERVER = iota + 1
	CLIENT
)

// API中定义操作时必须保留的字段
const (
	// 身份登记
	IDENTITY = "+identity+"
	// 心跳操作符
	HEARTBEAT = "+heartbeat+"
	// 默认包头
	DEFAULT_PACK_HEADER = "henrylee2cn"
	// SERVER默认UID
	DEFAULT_SERVER_UID = "server"
	// 默认端口
	DEFAULT_PORT = ":8080"
	// 服务器默认心跳间隔时长
	DEFAULT_TIMEOUT_S = 20e9
	// 客户端默认心跳间隔时长
	DEFAULT_TIMEOUT_C = 15e9
	// 等待连接的轮询时长
	LOOP_TIMEOUT = 1e9
)

type Teleport interface {
	// *以服务器模式运行，端口默认为常量DEFAULT_PORT
	Server(port ...string)
	// *以客户端模式运行，port为空时默认等于常量DEFAULT_PORT
	Client(serverAddr string, port string, isShort ...bool)
	// *主动推送信息，不写nodeuid默认随机发送给一个节点
	Request(body interface{}, operation string, flag string, nodeuid ...string)
	// 指定自定义的应用程序API
	SetAPI(api API) Teleport
	// 断开连接，参数为空则断开所有连接，服务器模式下还将停止监听
	Close(nodeuid ...string)

	// 设置唯一标识符，mine为本节点UID(默认ip:port)
	// server为服务器UID(默认为常量DEFAULT_SERVER_UID，此参数仅客户端模式下有用)
	// 可不调用该方法，此时UID均为默认
	SetUID(mine string, server ...string) Teleport
	// 设置包头字符串，默认为henrylee2cn
	SetPackHeader(string) Teleport
	// 设置指定API处理的数据的接收缓存通道长度
	SetApiRChan(int) Teleport
	// 设置每个连接对象的发送缓存通道长度
	SetConnWChan(int) Teleport
	// 设置每个连接对象的接收缓冲区大小
	SetConnBuffer(int) Teleport
	// 设置连接超时(心跳频率)
	SetTimeout(time.Duration) Teleport

	// 返回运行模式
	GetMode() int
	// 返回当前有效连接节点数
	CountNodes() int
}

type TP struct {
	// 本节点唯一标识符
	uid string
	// 运行模式 1 SERVER  2 CLIENT (用于判断自身模式)
	mode int
	// 服务器端口号，格式如":9988"
	port string
	// 服务器地址（不含端口号），格式如"127.0.0.1"
	serverAddr string
	// 长连接池，刚一连接时key为host:port形式，随后通过身份验证替换为UID
	connPool map[string]*Connect
	// 连接时长，心跳时长的依据，默认为常量DEFAULT_TIMEOUT
	timeout time.Duration
	// 粘包处理
	*Protocol
	// 全局接收缓存通道
	apiReadChan chan *NetData
	// 每个连接对象的发送缓存通道长度
	connWChanCap int
	// 每个连接对象的接收缓冲区大小
	connBufferLen int
	// 应用程序API
	api API
	// 服务器模式专有成员
	*tpServer
	// 客户端模式专有成员
	*tpClient
}

// 每个API方法需判断stutas状态，并做相应处理
type API map[string]Handle

// 请求处理接口
type Handle interface {
	Process(*NetData) *NetData
}

// 创建接口实例
func New() Teleport {
	return &TP{
		connPool:      make(map[string]*Connect),
		api:           API{},
		Protocol:      NewProtocol(DEFAULT_PACK_HEADER),
		apiReadChan:   make(chan *NetData, 4096),
		connWChanCap:  2048,
		connBufferLen: 1024,
		tpServer:      new(tpServer),
		tpClient:      new(tpClient),
	}
}

// ***********************************************实现接口*************************************************** \\

// 设置唯一标识符，mine为本节点UID(默认ip:port)
// server为服务器UID(默认为常量DEFAULT_SERVER_UID，此参数仅客户端模式下有用)
// 可不调用该方法，此时UID均为默认
func (self *TP) SetUID(mine string, server ...string) Teleport {
	if len(server) > 0 {
		self.tpClient.serverUID = server[0]
	}
	self.uid = mine
	return self
}

// 指定应用程序API
func (self *TP) SetAPI(api API) Teleport {
	self.api = api
	return self
}

// *主动推送信息，直到有连接出现开始发送，不写nodeuid默认随机发送给一个节点
func (self *TP) Request(body interface{}, operation string, flag string, nodeuid ...string) {
	var conn *Connect
	var uid string
	if len(nodeuid) == 0 {
		for {
			if self.CountNodes() > 0 {
				break
			}
			time.Sleep(LOOP_TIMEOUT)
		}
		// 一个随机节点的信息
		for uid, conn = range self.connPool {
			if conn.Usable {
				nodeuid = append(nodeuid, uid)
				break
			}
		}
	}
	// 等待并取得连接实例
	conn = self.getConn(nodeuid[0])
	for conn == nil || !conn.Usable {
		conn = self.getConn(nodeuid[0])
		time.Sleep(LOOP_TIMEOUT)
	}
	conn.WriteChan <- NewNetData(self.uid, nodeuid[0], operation, flag, body)
	// log.Println("添加一条请求：", conn.RemoteAddr().String(), operation, body)
}

// 断开连接，参数为空则断开所有连接，服务器模式下停止监听
func (self *TP) Close(nodeuid ...string) {
	if self.mode == CLIENT {
		self.tpClient.mustClose = true

	} else if self.mode == SERVER && self.tpServer.listener != nil {
		self.tpServer.listener.Close()
		log.Printf(" *     —— 服务器已终止监听 %v！ ——", self.port)
	}

	if len(nodeuid) == 0 {
		// 断开全部连接
		for uid, conn := range self.connPool {
			delete(self.connPool, uid)
			conn.Close()
			self.closeMsg(uid, conn.Addr(), conn.Short)
		}
		return
	}

	for _, uid := range nodeuid {
		conn := self.connPool[uid]
		delete(self.connPool, uid)
		conn.Close()
		self.closeMsg(uid, conn.Addr(), conn.Short)
	}
}

// 设置包头字符串，默认为henrylee2cn
func (self *TP) SetPackHeader(header string) Teleport {
	self.Protocol.ReSet(header)
	return self
}

// 设置全局接收缓存通道长度
func (self *TP) SetApiRChan(length int) Teleport {
	self.apiReadChan = make(chan *NetData, length)
	return self
}

// 设置每个连接对象的发送缓存通道长度
func (self *TP) SetConnWChan(length int) Teleport {
	self.connWChanCap = length
	return self
}

// 每个连接对象的接收缓冲区大小
func (self *TP) SetConnBuffer(length int) Teleport {
	self.connBufferLen = length
	return self
}

// 设置连接超长(心跳频率)
func (self *TP) SetTimeout(long time.Duration) Teleport {
	self.timeout = long
	return self
}

// 返回运行模式
func (self *TP) GetMode() int {
	return self.mode
}

// 返回当前有效连接节点数
func (self *TP) CountNodes() int {
	count := 0
	for _, conn := range self.connPool {
		if conn != nil && conn.Usable {
			count++
		}
	}
	return count
}

// ***********************************************公用方法*************************************************** \\

func (self *TP) read(conn *Connect) bool {
	read_len, err := conn.Read(conn.Buffer)
	if err != nil {
		return false
	}
	if read_len == 0 {
		return false // connection already closed by client
	}
	conn.TmpBuffer = append(conn.TmpBuffer, conn.Buffer[:read_len]...)
	self.save(conn)
	return true
}

// 根据uid获取连接对象
func (self *TP) getConn(nodeuid string) *Connect {
	return self.connPool[nodeuid]
}

// 根据uid获取连接对象地址
func (self *TP) getConnAddr(nodeuid string) string {
	conn := self.getConn(nodeuid)
	if conn == nil {
		// log.Printf("已与节点 %v 失去连接，无法完成发送请求！", nodeuid)
		return ""
	}
	return conn.Addr()
}

// 关闭连接，退出协程
func (self *TP) closeConn(nodeuid string, reconnect bool) {
	conn, ok := self.connPool[nodeuid]
	if !ok {
		return
	}

	// 是否允许自动重连
	if reconnect {
		self.connPool[nodeuid] = nil
	} else {
		delete(self.connPool, nodeuid)
	}

	if conn == nil {
		return
	}
	conn.Close()
	self.closeMsg(nodeuid, conn.Addr(), conn.Short)
}

// 关闭连接时log信息
func (self *TP) closeMsg(uid, addr string, short bool) {
	if short {
		return
	}
	switch self.mode {
	case SERVER:
		log.Printf(" *     —— 与客户端 %v (%v) 断开连接！——", uid, addr)
	case CLIENT:
		log.Printf(" *     —— 与服务器 %v 断开连接！——", addr)
	}
}

// 通信数据编码与发送
func (self *TP) send(data *NetData) {
	if data.From == "" {
		data.From = self.uid
	}

	d, err := json.Marshal(*data)
	if err != nil {
		debug.Println("Debug: 发送数据-编码出错", err)
		return
	}
	conn := self.getConn(data.To)
	if conn == nil {
		debug.Printf("Debug: 发送数据-连接已断开: %+v", data)
		return
	}
	// 封包
	end := self.Packet(d)
	// 发送
	conn.Write(end)
	debug.Printf("Debug: 发送数据-成功: %+v", data)
}

// 解码收到的数据并存入缓存
func (self *TP) save(conn *Connect) {
	debug.Printf("Debug: 收到数据-Byte: %v", conn.TmpBuffer)
	// 解包
	dataSlice := make([][]byte, 10)
	dataSlice, conn.TmpBuffer = self.Unpack(conn.TmpBuffer)

	for _, data := range dataSlice {
		debug.Printf("Debug: 收到数据-解码前: %v", string(data))

		d := new(NetData)
		err := json.Unmarshal(data, d)
		if err == nil {
			// 修复缺失请求方地址的请求
			if d.From == "" {
				d.From = conn.Addr()
			}
			// 添加到读取缓存
			self.apiReadChan <- d
			debug.Printf("Debug: 收到数据-NetData: %+v", d)
		} else {
			debug.Printf("Debug: 收到数据-解码错误: %v", err)
		}
	}
}

// 使用API并发处理请求
func (self *TP) apiHandle() {
	for {
		req := <-self.apiReadChan
		go func(req *NetData) {
			var conn *Connect

			operation, from, to, flag := req.Operation, req.To, req.From, req.Flag
			handle, ok := self.api[operation]

			// 非法请求返回错误
			if !ok {
				// log.Printf("%+v", req)
				if self.mode == SERVER {
					self.autoErrorHandle(req, LLLEGAL, "服务器 ("+self.getConn(to).LocalAddr().String()+") 不存在API接口: "+req.Operation+"！", to)
					log.Printf("客户端 %v (%v) 正在请求不存在的API接口: %v！", to, self.getConnAddr(to), req.Operation)
				} else {
					self.autoErrorHandle(req, LLLEGAL, "客户端 "+from+" ("+self.getConn(to).LocalAddr().String()+") 不存在API接口: "+req.Operation+"！", to)
					log.Printf("服务器 (%v) 正在请求不存在的API接口: %v！", self.getConnAddr(to), req.Operation)
				}
				return
			}

			resp := handle.Process(req)
			if resp == nil {
				if conn = self.getConn(to); conn != nil && self.getConn(to).Short {
					self.closeConn(to, false)
				}
				return //continue
			}

			if resp.To == "" {
				resp.To = to
			}

			// 若指定节点连接不存在，则向原请求端返回错误
			if conn = self.getConn(resp.To); conn == nil {
				self.autoErrorHandle(req, FAILURE, "", to)
				return
			}

			// 默认指定与req相同的操作符
			if resp.Operation == "" {
				resp.Operation = operation
			}

			if resp.From == "" {
				resp.From = from
			}

			if resp.Flag == "" {
				resp.Flag = flag
			}

			conn.WriteChan <- resp

		}(req)
	}
}

func (self *TP) autoErrorHandle(data *NetData, status int, msg string, reqFrom string) bool {
	oldConn := self.getConn(reqFrom)
	if oldConn == nil {
		return false
	}
	respErr := ReturnError(data, status, msg)
	respErr.From = self.uid
	respErr.To = reqFrom
	oldConn.WriteChan <- respErr
	return true
}

// 连接权限验证
func (self *TP) checkRights(data *NetData, addr string) bool {
	if data.To != self.uid {
		log.Printf("陌生连接(%v)提供的服务器标识符错误，请求被拒绝！", addr)
		return false
	}
	return true
}

// 强制设定系统保留的API
func (self *TP) reserveAPI() {
	// 添加保留规则——身份识别
	self.api[IDENTITY] = identi
	// 添加保留规则——忽略心跳请求
	self.api[HEARTBEAT] = beat
}

var identi, beat = new(identity), new(heartbeat)

type identity struct{}

func (*identity) Process(receive *NetData) *NetData {
	receive.From, receive.To = receive.To, receive.From
	return nil
}

type heartbeat struct{}

func (*heartbeat) Process(receive *NetData) *NetData {
	return nil
}
<commit_msg>clear<commit_after>// Teleport是一款适用于分布式系统的高并发API框架，它采用socket全双工通信，实现S/C对等工作，支持长、短两种连接模式，支持断开后自动连接与手动断开连接，内部数据传输格式为JSON。
package teleport

import (
	"encoding/json"
	"github.com/henrylee2cn/teleport/debug"
	"log"
	"time"
)

// mode
const (
	SERVER = iota + 1
	CLIENT
)

// API中定义操作时必须保留的字段
const (
	// 身份登记
	IDENTITY = "+identity+"
	// 心跳操作符
	HEARTBEAT = "+heartbeat+"
	// 默认包头
	DEFAULT_PACK_HEADER = "henrylee2cn"
	// SERVER默认UID
	DEFAULT_SERVER_UID = "server"
	// 默认端口
	DEFAULT_PORT = ":8080"
	// 服务器默认心跳间隔时长
	DEFAULT_TIMEOUT_S = 20e9
	// 客户端默认心跳间隔时长
	DEFAULT_TIMEOUT_C = 15e9
	// 等待连接的轮询时长
	LOOP_TIMEOUT = 1e9
)

type Teleport interface {
	// *以服务器模式运行，端口默认为常量DEFAULT_PORT
	Server(port ...string)
	// *以客户端模式运行，port为空时默认等于常量DEFAULT_PORT
	Client(serverAddr string, port string, isShort ...bool)
	// *主动推送信息，不写nodeuid默认随机发送给一个节点
	Request(body interface{}, operation string, flag string, nodeuid ...string)
	// 指定自定义的应用程序API
	SetAPI(api API) Teleport
	// 断开连接，参数为空则断开所有连接，服务器模式下还将停止监听
	Close(nodeuid ...string)

	// 设置唯一标识符，mine为本节点UID(默认ip:port)
	// server为服务器UID(默认为常量DEFAULT_SERVER_UID，此参数仅客户端模式下有用)
	// 可不调用该方法，此时UID均为默认
	SetUID(mine string, server ...string) Teleport
	// 设置包头字符串，默认为henrylee2cn
	SetPackHeader(string) Teleport
	// 设置指定API处理的数据的接收缓存通道长度
	SetApiRChan(int) Teleport
	// 设置每个连接对象的发送缓存通道长度
	SetConnWChan(int) Teleport
	// 设置每个连接对象的接收缓冲区大小
	SetConnBuffer(int) Teleport
	// 设置连接超时(心跳频率)
	SetTimeout(time.Duration) Teleport

	// 返回运行模式
	GetMode() int
	// 返回当前有效连接节点数
	CountNodes() int
}

type TP struct {
	// 本节点唯一标识符
	uid string
	// 运行模式 1 SERVER  2 CLIENT (用于判断自身模式)
	mode int
	// 服务器端口号，格式如":9988"
	port string
	// 服务器地址（不含端口号），格式如"127.0.0.1"
	serverAddr string
	// 长连接池，刚一连接时key为host:port形式，随后通过身份验证替换为UID
	connPool map[string]*Connect
	// 连接时长，心跳时长的依据，默认为常量DEFAULT_TIMEOUT
	timeout time.Duration
	// 粘包处理
	*Protocol
	// 全局接收缓存通道
	apiReadChan chan *NetData
	// 每个连接对象的发送缓存通道长度
	connWChanCap int
	// 每个连接对象的接收缓冲区大小
	connBufferLen int
	// 应用程序API
	api API
	// 服务器模式专有成员
	*tpServer
	// 客户端模式专有成员
	*tpClient
}

// 每个API方法需判断stutas状态，并做相应处理
type API map[string]Handle

// 请求处理接口
type Handle interface {
	Process(*NetData) *NetData
}

// 创建接口实例
func New() Teleport {
	return &TP{
		connPool:      make(map[string]*Connect),
		api:           API{},
		Protocol:      NewProtocol(DEFAULT_PACK_HEADER),
		apiReadChan:   make(chan *NetData, 4096),
		connWChanCap:  2048,
		connBufferLen: 1024,
		tpServer:      new(tpServer),
		tpClient:      new(tpClient),
	}
}

// ***********************************************实现接口*************************************************** \\

// 设置唯一标识符，mine为本节点UID(默认ip:port)
// server为服务器UID(默认为常量DEFAULT_SERVER_UID，此参数仅客户端模式下有用)
// 可不调用该方法，此时UID均为默认
func (self *TP) SetUID(mine string, server ...string) Teleport {
	if len(server) > 0 {
		self.tpClient.serverUID = server[0]
	}
	self.uid = mine
	return self
}

// 指定应用程序API
func (self *TP) SetAPI(api API) Teleport {
	self.api = api
	return self
}

// *主动推送信息，直到有连接出现开始发送，不写nodeuid默认随机发送给一个节点
func (self *TP) Request(body interface{}, operation string, flag string, nodeuid ...string) {
	var conn *Connect
	var uid string
	if len(nodeuid) == 0 {
		for {
			if self.CountNodes() > 0 {
				break
			}
			time.Sleep(LOOP_TIMEOUT)
		}
		// 一个随机节点的信息
		for uid, conn = range self.connPool {
			if conn.Usable {
				nodeuid = append(nodeuid, uid)
				break
			}
		}
	}
	// 等待并取得连接实例
	conn = self.getConn(nodeuid[0])
	for conn == nil || !conn.Usable {
		conn = self.getConn(nodeuid[0])
		time.Sleep(LOOP_TIMEOUT)
	}
	conn.WriteChan <- NewNetData(self.uid, nodeuid[0], operation, flag, body)
	// log.Println("添加一条请求：", conn.RemoteAddr().String(), operation, body)
}

// 断开连接，参数为空则断开所有连接，服务器模式下停止监听
func (self *TP) Close(nodeuid ...string) {
	if self.mode == CLIENT {
		self.tpClient.mustClose = true

	} else if self.mode == SERVER && self.tpServer.listener != nil {
		self.tpServer.listener.Close()
		log.Printf(" *     —— 服务器已终止监听 %v！ ——", self.port)
	}

	if len(nodeuid) == 0 {
		// 断开全部连接
		for uid, conn := range self.connPool {
			delete(self.connPool, uid)
			conn.Close()
			self.closeMsg(uid, conn.Addr(), conn.Short)
		}
		return
	}

	for _, uid := range nodeuid {
		conn := self.connPool[uid]
		delete(self.connPool, uid)
		conn.Close()
		self.closeMsg(uid, conn.Addr(), conn.Short)
	}
}

// 设置包头字符串，默认为henrylee2cn
func (self *TP) SetPackHeader(header string) Teleport {
	self.Protocol.ReSet(header)
	return self
}

// 设置全局接收缓存通道长度
func (self *TP) SetApiRChan(length int) Teleport {
	self.apiReadChan = make(chan *NetData, length)
	return self
}

// 设置每个连接对象的发送缓存通道长度
func (self *TP) SetConnWChan(length int) Teleport {
	self.connWChanCap = length
	return self
}

// 每个连接对象的接收缓冲区大小
func (self *TP) SetConnBuffer(length int) Teleport {
	self.connBufferLen = length
	return self
}

// 设置连接超长(心跳频率)
func (self *TP) SetTimeout(long time.Duration) Teleport {
	self.timeout = long
	return self
}

// 返回运行模式
func (self *TP) GetMode() int {
	return self.mode
}

// 返回当前有效连接节点数
func (self *TP) CountNodes() int {
	count := 0
	for _, conn := range self.connPool {
		if conn != nil && conn.Usable {
			count++
		}
	}
	return count
}

// ***********************************************公用方法*************************************************** \\

func (self *TP) read(conn *Connect) bool {
	read_len, err := conn.Read(conn.Buffer)
	if err != nil {
		return false
	}
	if read_len == 0 {
		return false // connection already closed by client
	}
	conn.TmpBuffer = append(conn.TmpBuffer, conn.Buffer[:read_len]...)
	self.save(conn)
	return true
}

// 根据uid获取连接对象
func (self *TP) getConn(nodeuid string) *Connect {
	return self.connPool[nodeuid]
}

// 根据uid获取连接对象地址
func (self *TP) getConnAddr(nodeuid string) string {
	conn := self.getConn(nodeuid)
	if conn == nil {
		// log.Printf("已与节点 %v 失去连接，无法完成发送请求！", nodeuid)
		return ""
	}
	return conn.Addr()
}

// 关闭连接，退出协程
func (self *TP) closeConn(nodeuid string, reconnect bool) {
	conn, ok := self.connPool[nodeuid]
	if !ok {
		return
	}

	// 是否允许自动重连
	if reconnect {
		self.connPool[nodeuid] = nil
	} else {
		delete(self.connPool, nodeuid)
	}

	if conn == nil {
		return
	}
	conn.Close()
	self.closeMsg(nodeuid, conn.Addr(), conn.Short)
}

// 关闭连接时log信息
func (self *TP) closeMsg(uid, addr string, short bool) {
	if short {
		return
	}
	switch self.mode {
	case SERVER:
		log.Printf(" *     —— 与客户端 %v (%v) 断开连接！——", uid, addr)
	case CLIENT:
		log.Printf(" *     —— 与服务器 %v 断开连接！——", addr)
	}
}

// 通信数据编码与发送
func (self *TP) send(data *NetData) {
	if data.From == "" {
		data.From = self.uid
	}

	d, err := json.Marshal(*data)
	if err != nil {
		debug.Println("Debug: 发送数据-编码出错", err)
		return
	}
	conn := self.getConn(data.To)
	if conn == nil {
		debug.Printf("Debug: 发送数据-连接已断开: %+v", data)
		return
	}
	// 封包
	end := self.Packet(d)
	// 发送
	conn.Write(end)
	debug.Printf("Debug: 发送数据-成功: %+v", data)
}

// 解码收到的数据并存入缓存
func (self *TP) save(conn *Connect) {
	debug.Printf("Debug: 收到数据-Byte: %v", conn.TmpBuffer)
	// 解包
	dataSlice := make([][]byte, 10)
	dataSlice, conn.TmpBuffer = self.Unpack(conn.TmpBuffer)

	for _, data := range dataSlice {
		debug.Printf("Debug: 收到数据-解码前: %v", string(data))

		d := new(NetData)
		err := json.Unmarshal(data, d)
		if err == nil {
			// 修复缺失请求方地址的请求
			if d.From == "" {
				d.From = conn.Addr()
			}
			// 添加到读取缓存
			self.apiReadChan <- d
			debug.Printf("Debug: 收到数据-NetData: %+v", d)
		} else {
			debug.Printf("Debug: 收到数据-解码错误: %v", err)
		}
	}
}

// 使用API并发处理请求
func (self *TP) apiHandle() {
	for {
		req := <-self.apiReadChan
		go func(req *NetData) {
			var conn *Connect

			operation, from, to, flag := req.Operation, req.To, req.From, req.Flag
			handle, ok := self.api[operation]

			// 非法请求返回错误
			if !ok {
				// log.Printf("%+v", req)
				if self.mode == SERVER {
					self.autoErrorHandle(req, LLLEGAL, "服务器 ("+self.getConn(to).LocalAddr().String()+") 不存在API接口: "+req.Operation+"！", to)
					log.Printf("客户端 %v (%v) 正在请求不存在的API接口: %v！", to, self.getConnAddr(to), req.Operation)
				} else {
					self.autoErrorHandle(req, LLLEGAL, "客户端 "+from+" ("+self.getConn(to).LocalAddr().String()+") 不存在API接口: "+req.Operation+"！", to)
					log.Printf("服务器 (%v) 正在请求不存在的API接口: %v！", self.getConnAddr(to), req.Operation)
				}
				return
			}

			resp := handle.Process(req)
			if resp == nil {
				if conn = self.getConn(to); conn != nil && self.getConn(to).Short {
					self.closeConn(to, false)
				}
				return //continue
			}

			if resp.To == "" {
				resp.To = to
			}

			// 若指定节点连接不存在，则向原请求端返回错误
			if conn = self.getConn(resp.To); conn == nil {
				self.autoErrorHandle(req, FAILURE, "", to)
				return
			}

			// 默认指定与req相同的操作符
			if resp.Operation == "" {
				resp.Operation = operation
			}

			if resp.From == "" {
				resp.From = from
			}

			if resp.Flag == "" {
				resp.Flag = flag
			}

			conn.WriteChan <- resp

		}(req)
	}
}

func (self *TP) autoErrorHandle(data *NetData, status int, msg string, reqFrom string) bool {
	oldConn := self.getConn(reqFrom)
	if oldConn == nil {
		return false
	}
	respErr := ReturnError(data, status, msg)
	respErr.From = self.uid
	respErr.To = reqFrom
	oldConn.WriteChan <- respErr
	return true
}

// 连接权限验证
func (self *TP) checkRights(data *NetData, addr string) bool {
	if data.To != self.uid {
		log.Printf("陌生连接(%v)提供的服务器标识符错误，请求被拒绝！", addr)
		return false
	}
	return true
}

// 强制设定系统保留的API
func (self *TP) reserveAPI() {
	// 添加保留规则——身份识别
	self.api[IDENTITY] = identi
	// 添加保留规则——忽略心跳请求
	self.api[HEARTBEAT] = beat
}

var identi, beat = new(identity), new(heartbeat)

type identity struct{}

func (*identity) Process(receive *NetData) *NetData {
	return nil
}

type heartbeat struct{}

func (*heartbeat) Process(receive *NetData) *NetData {
	return nil
}
<|endoftext|>


<commit_before>package main

var pkgTemplate = `{{with .PDoc}}
{{if $.IsMain}}
> {{ base .ImportPath }}
{{comment_md .Doc}}
{{else}}
# {{ .Name }}
` + "`" + `import "{{.ImportPath | import_as}}"` + "`" + `

* [Overview](#pkg-overview)
* [Imported Packages](#pkg-imports)
* [Index](#pkg-index){{if $.Examples}}
* [Examples](#pkg-examples){{- end}}

## Overview
{{pkgdoc_md .Doc}}
{{example_text $ "" "#### "}}

## Imported Packages

{{list_imports $ .Imports (.ImportPath|import_as)}}

## Index{{if .Consts}}
* [Constants](#pkg-constants){{end}}{{if .Vars}}
* [Variables](#pkg-variables){{end}}{{- range .Funcs -}}{{$name_html := html .Name}}
* [{{node_html $ .Decl false | sanitize | md}}](#{{$name_html}}){{- end}}{{- range .Types}}{{$tname_html := html .Name}}
* [type {{$tname_html}}](#{{$tname_html}}){{- range .Funcs}}{{$name_html := html .Name}}
  * [{{node_html $ .Decl false | sanitize | md}}](#{{$name_html}}){{- end}}{{- range .Methods}}{{$name_html := html .Name}}
  * [{{node_html $ .Decl false | sanitize | md}}](#{{$tname_html}}.{{$name_html}}){{- end}}{{- end}}{{- if $.Notes}}{{- range $marker, $item := $.Notes}}
* [{{noteTitle $marker | html}}s](#pkg-note-{{$marker}}){{end}}{{end}}
{{if $.Examples}}
#### Examples{{- range $.Examples}}
* [{{example_name .Name}}](#example_{{.Name}}){{- end}}{{- end}}

{{with .Consts}}## Constants
{{range .}}{{node $ .Decl | pre}}
{{comment_md .Doc}}{{end}}{{end}}
{{with .Vars}}## Variables
{{range .}}{{node $ .Decl | pre}}
{{comment_md .Doc}}{{end}}{{end}}

{{range .Funcs}}{{$name_html := html .Name}}## func [{{$name_html}}]({{gh_url $ .Decl}})
{{node $ .Decl | pre}}
{{comment_md .Doc}}
{{example_text $ "" "#### "}}
{{callgraph_html $ "" .Name}}{{end}}
{{range .Types}}{{$tname := .Name}}{{$tname_html := html .Name}}## type [{{$tname_html}}]({{gh_url $ .Decl}})
{{node $ .Decl | pre}}
{{comment_md .Doc}}{{range .Consts}}
{{node $ .Decl | pre }}
{{comment_md .Doc}}{{end}}{{range .Vars}}
{{node $ .Decl | pre }}
{{comment_md .Doc}}{{end}}

{{example_text $ "" "#### "}}
{{implements_html $ $tname}}
{{methodset_html $ $tname}}

{{range .Funcs}}{{$name_html := html .Name}}### func [{{$name_html}}]({{gh_url $ .Decl}})
{{node $ .Decl | pre}}
{{comment_md .Doc}}
{{example_text $ "" "#### "}}
{{callgraph_html $ "" .Name}}

{{range .Methods}}{{$name_html := html .Name}}### func ({{md .Recv}}) [{{$name_html}}]({{gh_url $ .Decl}})
{{node $ .Decl | pre}}
{{comment_md .Doc}}
{{$name := printf "%s_%s" $tname .Name}}
{{example_text $ "" "#### "}}
{{callgraph_html $ .Recv .Name}}
{{end}}{{end}}{{end}}

{{with $.Notes}}
{{range $marker, $content := .}}
## {{noteTitle $marker | html}}s

{{range .}}

☞ {{html .Body}}

{{end}}

{{end}}
{{end}}
{{end}}
`
<commit_msg>Ugly hack to fix missing example code when using the -ex option. Fix for the ugly-hack  to correct some fat-finger cut&paste!<commit_after>package main

var pkgTemplate = `{{with .PDoc}}
{{if $.IsMain}}
> {{ base .ImportPath }}
{{comment_md .Doc}}
{{else}}
# {{ .Name }}
` + "`" + `import "{{.ImportPath | import_as}}"` + "`" + `

* [Overview](#pkg-overview)
* [Imported Packages](#pkg-imports)
* [Index](#pkg-index){{if $.Examples}}
* [Examples](#pkg-examples){{- end}}

## Overview
{{pkgdoc_md .Doc}}
{{example_text $ "" "#### "}}

## Imported Packages

{{list_imports $ .Imports (.ImportPath|import_as)}}

## Index{{if .Consts}}
* [Constants](#pkg-constants){{end}}{{if .Vars}}
* [Variables](#pkg-variables){{end}}{{- range .Funcs -}}{{$name_html := html .Name}}
* [{{node_html $ .Decl false | sanitize | md}}](#{{$name_html}}){{- end}}{{- range .Types}}{{$tname_html := html .Name}}
* [type {{$tname_html}}](#{{$tname_html}}){{- range .Funcs}}{{$name_html := html .Name}}
  * [{{node_html $ .Decl false | sanitize | md}}](#{{$name_html}}){{- end}}{{- range .Methods}}{{$name_html := html .Name}}
  * [{{node_html $ .Decl false | sanitize | md}}](#{{$tname_html}}.{{$name_html}}){{- end}}{{- end}}{{- if $.Notes}}{{- range $marker, $item := $.Notes}}
* [{{noteTitle $marker | html}}s](#pkg-note-{{$marker}}){{end}}{{end}}
{{if $.Examples}}
#### Examples{{- range $.Examples}}
* [{{example_name .Name}}](#example_{{.Name}}){{- end}}{{- end}}

{{with .Consts}}## Constants
{{range .}}{{node $ .Decl | pre}}
{{comment_md .Doc}}{{end}}{{end}}
{{with .Vars}}## Variables
{{range .}}{{node $ .Decl | pre}}
{{comment_md .Doc}}{{end}}{{end}}

{{range .Funcs}}{{$name_html := html .Name}}## func [{{$name_html}}]({{gh_url $ .Decl}})
{{node $ .Decl | pre}}
{{comment_md .Doc}}
{{example_text $ .Name "#### "}}
{{callgraph_html $ "" .Name}}{{end}}
{{range .Types}}{{$tname := .Name}}{{$tname_html := html .Name}}## type [{{$tname_html}}]({{gh_url $ .Decl}})
{{node $ .Decl | pre}}
{{comment_md .Doc}}{{range .Consts}}
{{node $ .Decl | pre }}
{{comment_md .Doc}}{{end}}{{range .Vars}}
{{node $ .Decl | pre }}
{{comment_md .Doc}}{{end}}

{{example_text $ $tname "#### "}}
{{implements_html $ $tname}}
{{methodset_html $ $tname}}

{{range .Funcs}}{{$name_html := html .Name}}### func [{{$name_html}}]({{gh_url $ .Decl}})
{{node $ .Decl | pre}}
{{comment_md .Doc}}
{{example_text $ .Name "#### "}}{{end}}
{{callgraph_html $ "" .Name}}

{{range .Methods}}{{$name_html := html .Name}}### func ({{md .Recv}}) [{{$name_html}}]({{gh_url $ .Decl}})
{{node $ .Decl | pre}}
{{comment_md .Doc}}
{{$name := printf "%s_%s" $tname .Name}}
{{example_text $ $name "#### "}}
{{callgraph_html $ .Recv .Name}}
{{end}}{{end}}{{end}}

{{with $.Notes}}
{{range $marker, $content := .}}
## {{noteTitle $marker | html}}s

{{range .}}

☞ {{html .Body}}

{{end}}

{{end}}
{{end}}
{{end}}
`
<|endoftext|>


<commit_before>package revel

import (
	"fmt"
	"html"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"time"
)

var ERROR_CLASS = "hasError"

// This object handles loading and parsing of templates.
// Everything below the application's views directory is treated as a template.
type TemplateLoader struct {
	// This is the set of all templates under views
	templateSet *template.Template
	// If an error was encountered parsing the templates, it is stored here.
	compileError *Error
	// Paths to search for templates, in priority order.
	paths []string
	// Map from template name to the path from whence it was loaded.
	templatePaths map[string]string
}

type Template interface {
	Name() string
	Content() []string
	Render(wr io.Writer, arg interface{}) error
}

var invalidSlugPattern = regexp.MustCompile(`[^a-z0-9 _-]`)
var whiteSpacePattern = regexp.MustCompile(`\s+`)

var (
	// The functions available for use in the templates.
	TemplateFuncs = map[string]interface{}{
		"url": ReverseUrl,
		"eq":  Equal,
		"set": func(renderArgs map[string]interface{}, key string, value interface{}) template.HTML {
			renderArgs[key] = value
			return template.HTML("")
		},
		"append": func(renderArgs map[string]interface{}, key string, value interface{}) template.HTML {
			if renderArgs[key] == nil {
				renderArgs[key] = []interface{}{value}
			} else {
				renderArgs[key] = append(renderArgs[key].([]interface{}), value)
			}
			return template.HTML("")
		},
		"field": NewField,
		"option": func(f *Field, val, label string) template.HTML {
			selected := ""
			if f.Flash() == val {
				selected = " selected"
			}
			return template.HTML(fmt.Sprintf(`<option value="%s"%s>%s</option>`,
				html.EscapeString(val), selected, html.EscapeString(label)))
		},
		"radio": func(f *Field, val string) template.HTML {
			checked := ""
			if f.Flash() == val {
				checked = " checked"
			}
			return template.HTML(fmt.Sprintf(`<input type="radio" name="%s" value="%s"%s>`,
				html.EscapeString(f.Name), html.EscapeString(val), checked))
		},
		"checkbox": func(f *Field, val string) template.HTML {
			checked := ""
			if f.Flash() == val {
				checked = " checked"
			}
			return template.HTML(fmt.Sprintf(`<input type="checkbox" name="%s" value="%s"%s>`,
				html.EscapeString(f.Name), html.EscapeString(val), checked))
		},
		// Pads the given string with  's up to the given width.
		"pad": func(str string, width int) template.HTML {
			if len(str) >= width {
				return template.HTML(html.EscapeString(str))
			}
			return template.HTML(html.EscapeString(str) + strings.Repeat(" ", width-len(str)))
		},

		"errorClass": func(name string, renderArgs map[string]interface{}) template.HTML {
			errorMap, ok := renderArgs["errors"].(map[string]*ValidationError)
			if !ok || errorMap == nil {
				WARN.Println("Called 'errorClass' without 'errors' in the render args.")
				return template.HTML("")
			}
			valError, ok := errorMap[name]
			if !ok || valError == nil {
				return template.HTML("")
			}
			return template.HTML(ERROR_CLASS)
		},

		"msg": func(renderArgs map[string]interface{}, message string, args ...interface{}) template.HTML {
			return template.HTML(Message(renderArgs[CurrentLocaleRenderArg].(string), message, args...))
		},

		// Replaces newlines with

		"nl2br": func(text string) template.HTML {
			return template.HTML(strings.Replace(template.HTMLEscapeString(text), "\n", "
", -1))
		},

		// Skips sanitation on the parameter.  Do not use with dynamic data.
		"raw": func(text string) template.HTML {
			return template.HTML(text)
		},

		// Pluralize, a helper for pluralizing words to correspond to data of dynamic length.
		// items - a slice of items, or an integer indicating how many items there are.
		// pluralOverrides - optional arguments specifying the output in the
		//     singular and plural cases.  by default "" and "s"
		"pluralize": func(items interface{}, pluralOverrides ...string) string {
			singular, plural := "", "s"
			if len(pluralOverrides) >= 1 {
				singular = pluralOverrides[0]
				if len(pluralOverrides) == 2 {
					plural = pluralOverrides[1]
				}
			}

			switch v := reflect.ValueOf(items); v.Kind() {
			case reflect.Int:
				if items.(int) != 1 {
					return plural
				}
			case reflect.Slice:
				if v.Len() != 1 {
					return plural
				}
			default:
				ERROR.Println("pluralize: unexpected type: ", v)
			}
			return singular
		},

		// Format a date according to the application's default date(time) format.
		"date": func(date time.Time) string {
			return date.Format(DateFormat)
		},
		"datetime": func(date time.Time) string {
			return date.Format(DateTimeFormat)
		},
		"slug": Slug,
	}
)

func NewTemplateLoader(paths []string) *TemplateLoader {
	loader := &TemplateLoader{
		paths: paths,
	}
	return loader
}

// This scans the views directory and parses all templates as Go Templates.
// If a template fails to parse, the error is set on the loader.
// (It's awkward to refresh a single Go Template)
func (loader *TemplateLoader) Refresh() *Error {
	TRACE.Printf("Refreshing templates from %s", loader.paths)

	loader.compileError = nil
	loader.templatePaths = map[string]string{}

	// Set the template delimiters for the project if present, then split into left
	// and right delimiters around a space character
	var splitDelims []string
	if TemplateDelims != "" {
		splitDelims = strings.Split(TemplateDelims, " ")
		if len(splitDelims) != 2 {
			log.Fatalln("app.conf: Incorrect format for template.delimiters")
		}
	}

	// Walk through the template loader's paths and build up a template set.
	var templateSet *template.Template = nil
	for _, basePath := range loader.paths {
		// Walk only returns an error if the template loader is completely unusable
		// (namely, if one of the TemplateFuncs does not have an acceptable signature).
		funcErr := filepath.Walk(basePath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				ERROR.Println("error walking templates:", err)
				return nil
			}

			// Walk into watchable directories
			if info.IsDir() {
				if !loader.WatchDir(info) {
					return filepath.SkipDir
				}
				return nil
			}

			// Only add watchable
			if !loader.WatchFile(info.Name()) {
				return nil
			}

			var fileStr string
			var fileBytes []byte

			// addTemplate allows the same template to be added multiple
			// times with different template names.
			addTemplate := func(templateName string) (err error) {
				// Convert template names to use forward slashes, even on Windows.
				if os.PathSeparator == '\\' {
					templateName = strings.Replace(templateName, `\`, `/`, -1) // `
				}

				// If we already loaded a template of this name, skip it.
				if _, ok := loader.templatePaths[templateName]; ok {
					return nil
				}
				loader.templatePaths[templateName] = path

				// Load the file if we haven't already
				if fileBytes == nil {
					fileBytes, err = ioutil.ReadFile(path)
					if err != nil {
						ERROR.Println("Failed reading file:", path)
						return nil
					}

					fileStr = string(fileBytes)
				}

				if templateSet == nil {
					// Create the template set.  This panics if any of the funcs do not
					// conform to expectations, so we wrap it in a func and handle those
					// panics by serving an error page.
					var funcError *Error
					func() {
						defer func() {
							if err := recover(); err != nil {
								funcError = &Error{
									Title:       "Panic (Template Loader)",
									Description: fmt.Sprintln(err),
								}
							}
						}()
						templateSet = template.New(templateName).Funcs(TemplateFuncs)
						// If alternate delimiters set for the project, change them for this set
						if splitDelims != nil && basePath == ViewsPath {
							templateSet.Delims(splitDelims[0], splitDelims[1])
						} else {
							// Reset to default otherwise
							templateSet.Delims("", "")
						}
						_, err = templateSet.Parse(fileStr)
					}()

					if funcError != nil {
						return funcError
					}

				} else {
					if splitDelims != nil && basePath == ViewsPath {
						templateSet.Delims(splitDelims[0], splitDelims[1])
					} else {
						templateSet.Delims("", "")
					}
					_, err = templateSet.New(templateName).Parse(fileStr)
				}
				return err
			}

			templateName := path[len(basePath)+1:]

			// Lower case the file name for case-insensitive matching
			lowerCaseTemplateName := strings.ToLower(templateName)

			err = addTemplate(templateName)
			err = addTemplate(lowerCaseTemplateName)

			// Store / report the first error encountered.
			if err != nil && loader.compileError == nil {
				_, line, description := parseTemplateError(err)
				loader.compileError = &Error{
					Title:       "Template Compilation Error",
					Path:        templateName,
					Description: description,
					Line:        line,
					SourceLines: strings.Split(fileStr, "\n"),
				}
				ERROR.Printf("Template compilation error (In %s around line %d):\n%s",
					templateName, line, description)
			}
			return nil
		})

		// If there was an error with the Funcs, set it and return immediately.
		if funcErr != nil {
			loader.compileError = funcErr.(*Error)
			return loader.compileError
		}
	}

	// Note: compileError may or may not be set.
	loader.templateSet = templateSet
	return loader.compileError
}

func (loader *TemplateLoader) WatchDir(info os.FileInfo) bool {
	// Watch all directories, except the ones starting with a dot.
	return !strings.HasPrefix(info.Name(), ".")
}

func (loader *TemplateLoader) WatchFile(basename string) bool {
	// Watch all files, except the ones starting with a dot.
	return !strings.HasPrefix(basename, ".")
}

// Parse the line, and description from an error message like:
// html/template:Application/Register.html:36: no such template "footer.html"
func parseTemplateError(err error) (templateName string, line int, description string) {
	description = err.Error()
	i := regexp.MustCompile(`:\d+:`).FindStringIndex(description)
	if i != nil {
		line, err = strconv.Atoi(description[i[0]+1 : i[1]-1])
		if err != nil {
			ERROR.Println("Failed to parse line number from error message:", err)
		}
		templateName = description[:i[0]]
		if colon := strings.Index(templateName, ":"); colon != -1 {
			templateName = templateName[colon+1:]
		}
		templateName = strings.TrimSpace(templateName)
		description = description[i[1]+1:]
	}
	return templateName, line, description
}

// Return the Template with the given name.  The name is the template's path
// relative to a template loader root.
//
// An Error is returned if there was any problem with any of the templates.  (In
// this case, if a template is returned, it may still be usable.)
func (loader *TemplateLoader) Template(name string) (Template, error) {
	// Lower case the file name to support case-insensitive matching
	name = strings.ToLower(name)
	// Look up and return the template.
	tmpl := loader.templateSet.Lookup(name)

	// This is necessary.
	// If a nil loader.compileError is returned directly, a caller testing against
	// nil will get the wrong result.  Something to do with casting *Error to error.
	var err error
	if loader.compileError != nil {
		err = loader.compileError
	}

	if tmpl == nil && err == nil {
		return nil, fmt.Errorf("Template %s not found.", name)
	}

	return GoTemplate{tmpl, loader}, err
}

// Adapter for Go Templates.
type GoTemplate struct {
	*template.Template
	loader *TemplateLoader
}

// return a 'revel.Template' from Go's template.
func (gotmpl GoTemplate) Render(wr io.Writer, arg interface{}) error {
	return gotmpl.Execute(wr, arg)
}

func (gotmpl GoTemplate) Content() []string {
	content, _ := ReadLines(gotmpl.loader.templatePaths[gotmpl.Name()])
	return content
}

/////////////////////
// Template functions
/////////////////////

// Return a url capable of invoking a given controller method:
// "Application.ShowApp 123" => "/app/123"
func ReverseUrl(args ...interface{}) (string, error) {
	if len(args) == 0 {
		return "", fmt.Errorf("no arguments provided to reverse route")
	}

	action := args[0].(string)
	actionSplit := strings.Split(action, ".")
	if len(actionSplit) != 2 {
		return "", fmt.Errorf("reversing '%s', expected 'Controller.Action'", action)
	}

	// Look up the types.
	var c Controller
	if err := c.SetAction(actionSplit[0], actionSplit[1]); err != nil {
		return "", fmt.Errorf("reversing %s: %s", action, err)
	}

	// Unbind the arguments.
	argsByName := make(map[string]string)
	for i, argValue := range args[1:] {
		Unbind(argsByName, c.MethodType.Args[i].Name, argValue)
	}

	return MainRouter.Reverse(args[0].(string), argsByName).Url, nil
}

func Slug(text string) string {
	separator := "-"
	text = strings.ToLower(text)
	text = invalidSlugPattern.ReplaceAllString(text, "")
	text = whiteSpacePattern.ReplaceAllString(text, separator)
	text = strings.Trim(text, separator)
	return text
}
<commit_msg>there's no reason to store both fileStr and fileBytes<commit_after>package revel

import (
	"fmt"
	"html"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"time"
)

var ERROR_CLASS = "hasError"

// This object handles loading and parsing of templates.
// Everything below the application's views directory is treated as a template.
type TemplateLoader struct {
	// This is the set of all templates under views
	templateSet *template.Template
	// If an error was encountered parsing the templates, it is stored here.
	compileError *Error
	// Paths to search for templates, in priority order.
	paths []string
	// Map from template name to the path from whence it was loaded.
	templatePaths map[string]string
}

type Template interface {
	Name() string
	Content() []string
	Render(wr io.Writer, arg interface{}) error
}

var invalidSlugPattern = regexp.MustCompile(`[^a-z0-9 _-]`)
var whiteSpacePattern = regexp.MustCompile(`\s+`)

var (
	// The functions available for use in the templates.
	TemplateFuncs = map[string]interface{}{
		"url": ReverseUrl,
		"eq":  Equal,
		"set": func(renderArgs map[string]interface{}, key string, value interface{}) template.HTML {
			renderArgs[key] = value
			return template.HTML("")
		},
		"append": func(renderArgs map[string]interface{}, key string, value interface{}) template.HTML {
			if renderArgs[key] == nil {
				renderArgs[key] = []interface{}{value}
			} else {
				renderArgs[key] = append(renderArgs[key].([]interface{}), value)
			}
			return template.HTML("")
		},
		"field": NewField,
		"option": func(f *Field, val, label string) template.HTML {
			selected := ""
			if f.Flash() == val {
				selected = " selected"
			}
			return template.HTML(fmt.Sprintf(`<option value="%s"%s>%s</option>`,
				html.EscapeString(val), selected, html.EscapeString(label)))
		},
		"radio": func(f *Field, val string) template.HTML {
			checked := ""
			if f.Flash() == val {
				checked = " checked"
			}
			return template.HTML(fmt.Sprintf(`<input type="radio" name="%s" value="%s"%s>`,
				html.EscapeString(f.Name), html.EscapeString(val), checked))
		},
		"checkbox": func(f *Field, val string) template.HTML {
			checked := ""
			if f.Flash() == val {
				checked = " checked"
			}
			return template.HTML(fmt.Sprintf(`<input type="checkbox" name="%s" value="%s"%s>`,
				html.EscapeString(f.Name), html.EscapeString(val), checked))
		},
		// Pads the given string with  's up to the given width.
		"pad": func(str string, width int) template.HTML {
			if len(str) >= width {
				return template.HTML(html.EscapeString(str))
			}
			return template.HTML(html.EscapeString(str) + strings.Repeat(" ", width-len(str)))
		},

		"errorClass": func(name string, renderArgs map[string]interface{}) template.HTML {
			errorMap, ok := renderArgs["errors"].(map[string]*ValidationError)
			if !ok || errorMap == nil {
				WARN.Println("Called 'errorClass' without 'errors' in the render args.")
				return template.HTML("")
			}
			valError, ok := errorMap[name]
			if !ok || valError == nil {
				return template.HTML("")
			}
			return template.HTML(ERROR_CLASS)
		},

		"msg": func(renderArgs map[string]interface{}, message string, args ...interface{}) template.HTML {
			return template.HTML(Message(renderArgs[CurrentLocaleRenderArg].(string), message, args...))
		},

		// Replaces newlines with

		"nl2br": func(text string) template.HTML {
			return template.HTML(strings.Replace(template.HTMLEscapeString(text), "\n", "
", -1))
		},

		// Skips sanitation on the parameter.  Do not use with dynamic data.
		"raw": func(text string) template.HTML {
			return template.HTML(text)
		},

		// Pluralize, a helper for pluralizing words to correspond to data of dynamic length.
		// items - a slice of items, or an integer indicating how many items there are.
		// pluralOverrides - optional arguments specifying the output in the
		//     singular and plural cases.  by default "" and "s"
		"pluralize": func(items interface{}, pluralOverrides ...string) string {
			singular, plural := "", "s"
			if len(pluralOverrides) >= 1 {
				singular = pluralOverrides[0]
				if len(pluralOverrides) == 2 {
					plural = pluralOverrides[1]
				}
			}

			switch v := reflect.ValueOf(items); v.Kind() {
			case reflect.Int:
				if items.(int) != 1 {
					return plural
				}
			case reflect.Slice:
				if v.Len() != 1 {
					return plural
				}
			default:
				ERROR.Println("pluralize: unexpected type: ", v)
			}
			return singular
		},

		// Format a date according to the application's default date(time) format.
		"date": func(date time.Time) string {
			return date.Format(DateFormat)
		},
		"datetime": func(date time.Time) string {
			return date.Format(DateTimeFormat)
		},
		"slug": Slug,
	}
)

func NewTemplateLoader(paths []string) *TemplateLoader {
	loader := &TemplateLoader{
		paths: paths,
	}
	return loader
}

// This scans the views directory and parses all templates as Go Templates.
// If a template fails to parse, the error is set on the loader.
// (It's awkward to refresh a single Go Template)
func (loader *TemplateLoader) Refresh() *Error {
	TRACE.Printf("Refreshing templates from %s", loader.paths)

	loader.compileError = nil
	loader.templatePaths = map[string]string{}

	// Set the template delimiters for the project if present, then split into left
	// and right delimiters around a space character
	var splitDelims []string
	if TemplateDelims != "" {
		splitDelims = strings.Split(TemplateDelims, " ")
		if len(splitDelims) != 2 {
			log.Fatalln("app.conf: Incorrect format for template.delimiters")
		}
	}

	// Walk through the template loader's paths and build up a template set.
	var templateSet *template.Template = nil
	for _, basePath := range loader.paths {
		// Walk only returns an error if the template loader is completely unusable
		// (namely, if one of the TemplateFuncs does not have an acceptable signature).
		funcErr := filepath.Walk(basePath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				ERROR.Println("error walking templates:", err)
				return nil
			}

			// Walk into watchable directories
			if info.IsDir() {
				if !loader.WatchDir(info) {
					return filepath.SkipDir
				}
				return nil
			}

			// Only add watchable
			if !loader.WatchFile(info.Name()) {
				return nil
			}

			var fileStr string

			// addTemplate allows the same template to be added multiple
			// times with different template names.
			addTemplate := func(templateName string) (err error) {
				// Convert template names to use forward slashes, even on Windows.
				if os.PathSeparator == '\\' {
					templateName = strings.Replace(templateName, `\`, `/`, -1) // `
				}

				// If we already loaded a template of this name, skip it.
				if _, ok := loader.templatePaths[templateName]; ok {
					return nil
				}
				loader.templatePaths[templateName] = path

				// Load the file if we haven't already
				if fileStr == "" {
					fileBytes, err := ioutil.ReadFile(path)
					if err != nil {
						ERROR.Println("Failed reading file:", path)
						return nil
					}

					fileStr = string(fileBytes)
				}

				if templateSet == nil {
					// Create the template set.  This panics if any of the funcs do not
					// conform to expectations, so we wrap it in a func and handle those
					// panics by serving an error page.
					var funcError *Error
					func() {
						defer func() {
							if err := recover(); err != nil {
								funcError = &Error{
									Title:       "Panic (Template Loader)",
									Description: fmt.Sprintln(err),
								}
							}
						}()
						templateSet = template.New(templateName).Funcs(TemplateFuncs)
						// If alternate delimiters set for the project, change them for this set
						if splitDelims != nil && basePath == ViewsPath {
							templateSet.Delims(splitDelims[0], splitDelims[1])
						} else {
							// Reset to default otherwise
							templateSet.Delims("", "")
						}
						_, err = templateSet.Parse(fileStr)
					}()

					if funcError != nil {
						return funcError
					}

				} else {
					if splitDelims != nil && basePath == ViewsPath {
						templateSet.Delims(splitDelims[0], splitDelims[1])
					} else {
						templateSet.Delims("", "")
					}
					_, err = templateSet.New(templateName).Parse(fileStr)
				}
				return err
			}

			templateName := path[len(basePath)+1:]

			// Lower case the file name for case-insensitive matching
			lowerCaseTemplateName := strings.ToLower(templateName)

			err = addTemplate(templateName)
			err = addTemplate(lowerCaseTemplateName)

			// Store / report the first error encountered.
			if err != nil && loader.compileError == nil {
				_, line, description := parseTemplateError(err)
				loader.compileError = &Error{
					Title:       "Template Compilation Error",
					Path:        templateName,
					Description: description,
					Line:        line,
					SourceLines: strings.Split(fileStr, "\n"),
				}
				ERROR.Printf("Template compilation error (In %s around line %d):\n%s",
					templateName, line, description)
			}
			return nil
		})

		// If there was an error with the Funcs, set it and return immediately.
		if funcErr != nil {
			loader.compileError = funcErr.(*Error)
			return loader.compileError
		}
	}

	// Note: compileError may or may not be set.
	loader.templateSet = templateSet
	return loader.compileError
}

func (loader *TemplateLoader) WatchDir(info os.FileInfo) bool {
	// Watch all directories, except the ones starting with a dot.
	return !strings.HasPrefix(info.Name(), ".")
}

func (loader *TemplateLoader) WatchFile(basename string) bool {
	// Watch all files, except the ones starting with a dot.
	return !strings.HasPrefix(basename, ".")
}

// Parse the line, and description from an error message like:
// html/template:Application/Register.html:36: no such template "footer.html"
func parseTemplateError(err error) (templateName string, line int, description string) {
	description = err.Error()
	i := regexp.MustCompile(`:\d+:`).FindStringIndex(description)
	if i != nil {
		line, err = strconv.Atoi(description[i[0]+1 : i[1]-1])
		if err != nil {
			ERROR.Println("Failed to parse line number from error message:", err)
		}
		templateName = description[:i[0]]
		if colon := strings.Index(templateName, ":"); colon != -1 {
			templateName = templateName[colon+1:]
		}
		templateName = strings.TrimSpace(templateName)
		description = description[i[1]+1:]
	}
	return templateName, line, description
}

// Return the Template with the given name.  The name is the template's path
// relative to a template loader root.
//
// An Error is returned if there was any problem with any of the templates.  (In
// this case, if a template is returned, it may still be usable.)
func (loader *TemplateLoader) Template(name string) (Template, error) {
	// Lower case the file name to support case-insensitive matching
	name = strings.ToLower(name)
	// Look up and return the template.
	tmpl := loader.templateSet.Lookup(name)

	// This is necessary.
	// If a nil loader.compileError is returned directly, a caller testing against
	// nil will get the wrong result.  Something to do with casting *Error to error.
	var err error
	if loader.compileError != nil {
		err = loader.compileError
	}

	if tmpl == nil && err == nil {
		return nil, fmt.Errorf("Template %s not found.", name)
	}

	return GoTemplate{tmpl, loader}, err
}

// Adapter for Go Templates.
type GoTemplate struct {
	*template.Template
	loader *TemplateLoader
}

// return a 'revel.Template' from Go's template.
func (gotmpl GoTemplate) Render(wr io.Writer, arg interface{}) error {
	return gotmpl.Execute(wr, arg)
}

func (gotmpl GoTemplate) Content() []string {
	content, _ := ReadLines(gotmpl.loader.templatePaths[gotmpl.Name()])
	return content
}

/////////////////////
// Template functions
/////////////////////

// Return a url capable of invoking a given controller method:
// "Application.ShowApp 123" => "/app/123"
func ReverseUrl(args ...interface{}) (string, error) {
	if len(args) == 0 {
		return "", fmt.Errorf("no arguments provided to reverse route")
	}

	action := args[0].(string)
	actionSplit := strings.Split(action, ".")
	if len(actionSplit) != 2 {
		return "", fmt.Errorf("reversing '%s', expected 'Controller.Action'", action)
	}

	// Look up the types.
	var c Controller
	if err := c.SetAction(actionSplit[0], actionSplit[1]); err != nil {
		return "", fmt.Errorf("reversing %s: %s", action, err)
	}

	// Unbind the arguments.
	argsByName := make(map[string]string)
	for i, argValue := range args[1:] {
		Unbind(argsByName, c.MethodType.Args[i].Name, argValue)
	}

	return MainRouter.Reverse(args[0].(string), argsByName).Url, nil
}

func Slug(text string) string {
	separator := "-"
	text = strings.ToLower(text)
	text = invalidSlugPattern.ReplaceAllString(text, "")
	text = whiteSpacePattern.ReplaceAllString(text, separator)
	text = strings.Trim(text, separator)
	return text
}
<|endoftext|>


<commit_before>package main

import (
	"golang.org/x/crypto/ssh/terminal"

	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"

	"github.com/cafebazaar/kupak/pkg/pak"
	"github.com/cafebazaar/kupak/pkg/util"
	"github.com/codegangsta/cli"
	"gopkg.in/yaml.v2"
)

func install(c *cli.Context) error {
	pakURL := c.Args().First()
	valuesFile := c.Args().Get(1)
	if pakURL == "" {
		return cli.NewExitError("please specify the pak", -1)
	}

	if strings.Index(pakURL, "/") == -1 &&
		!strings.HasSuffix(pakURL, ".json") &&
		!strings.HasSuffix(pakURL, ".yaml") {

		nameOfPakToInstall := pakURL
		repoAddr := c.GlobalString("repo")

		if len(repoAddr) > 0 {
			// TODO: change JoinURL
			repoPaks, err := pak.RepoFromURL(repoAddr)
			if err != nil {
				return cli.NewExitError(fmt.Sprintf("can't fetch the repo: %v", err.Error()), -1)
			}
			for _, pak := range repoPaks.Paks {
				if pak.Name == nameOfPakToInstall {
					pakURL = pak.URL
					if util.Relative(pakURL) {
						pakURL = util.JoinURL(repoAddr, pakURL)
					}
				}
			}
		}
	}

	p, err := pak.FromURL(pakURL)
	if err != nil {
		return cli.NewExitError(fmt.Sprintf("can't fetch the pak: %v", err.Error()), -1)
	}

	values := make(map[string]interface{})
	if c.Bool("interactive") || terminal.IsTerminal(int(os.Stdin.Fd())) {
		values, err = readValuesInteractively(p)
	} else {
		values, err = readValuesFromFile(p, valuesFile)
	}
	if err != nil {
		return cli.NewExitError(fmt.Sprintf("error in reading value: %v", err.Error()), -1)
	}

	_, err = pakManager.Install(p, c.GlobalString("namespace"), values)
	if err != nil {
		return cli.NewExitError(fmt.Sprintf("installation error: %v", err.Error()), -1)
	}
	return nil
}

func readValuesFromFile(p *pak.Pak, path string) (map[string]interface{}, error) {
	values := make(map[string]interface{})
	var valuesData []byte
	var err error
	if path == "" {
		valuesData, err = ioutil.ReadAll(os.Stdin)
	} else {
		valuesData, err = ioutil.ReadFile(path)
	}
	if err != nil {
		return nil, err
	}
	err = yaml.Unmarshal(valuesData, &values)
	return values, err
}

func readValuesInteractively(p *pak.Pak) (map[string]interface{}, error) {
	values := make(map[string]interface{})

	// ask for group
	groupValue, err := scanValue("Group Name (return for random): ", "string", false)
	if err != nil {
		return nil, err
	}
	if groupValue != nil {
		values["group"] = groupValue.(string)
	}

	// ask for all properties
	for i := range p.Properties {
		var prompt string
		if p.Properties[i].Default == nil {
			prompt = fmt.Sprintf("Field \"%s\" [type: %s]? ", p.Properties[i].Name, p.Properties[i].Type)
		} else {
			prompt = fmt.Sprintf("Field \"%s\" [type: %s, default: %v] (return for default)? ", p.Properties[i].Name, p.Properties[i].Type, p.Properties[i].Default)
		}
		value, err := scanValue(prompt, p.Properties[i].Type, p.Properties[i].Default == nil)
		if err != nil {
			return nil, err
		}
		if value != nil {
			values[p.Properties[i].Name] = value
		}
	}
	return values, nil
}

func scanValue(prompt string, valueType string, required bool) (interface{}, error) {
	bio := bufio.NewReader(os.Stdin)
	var value []byte
	var err error
	for {
		fmt.Printf(prompt)
		value, _, err = bio.ReadLine()
		if err != nil {
			return nil, err
		}
		if len(value) == 0 && !required {
			return nil, nil
		} else if len(value) == 0 && required {
			continue
		}

		switch valueType {
		case "int":
			i, err := strconv.Atoi(string(value))
			if err != nil {
				fmt.Println("given value is not an int, try again")
				continue
			}
			return i, nil
		case "bool":
			b, err := util.StringToBool(string(value))
			if err != nil {
				fmt.Println(err)
				continue
			}
			return b, nil
		case "string":
			fallthrough
		default:
			return string(value), nil
		}
	}
}
<commit_msg>Added repos file functionality to install sub-command<commit_after>package main

import (
	"golang.org/x/crypto/ssh/terminal"

	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"

	"github.com/cafebazaar/kupak/pkg/pak"
	"github.com/cafebazaar/kupak/pkg/util"
	"github.com/codegangsta/cli"
	"gopkg.in/yaml.v2"
)

type installationSource struct {
	repo    string
	address string
}

func install(c *cli.Context) error {
	pakURL := c.Args().First()
	valuesFile := c.Args().Get(1)
	if pakURL == "" {
		return cli.NewExitError("please specify the pak", -1)
	}

	if strings.Index(pakURL, "/") == -1 &&
		!strings.HasSuffix(pakURL, ".json") &&
		!strings.HasSuffix(pakURL, ".yaml") {

		nameOfPakToInstall := pakURL

		repoAddresses := make([]string, 0)
		if c.GlobalIsSet("repo") {
			repoAddresses = append(repoAddresses, c.GlobalString("repo"))
		} else {
			reposFileCreateIfNotExist()
			for _, repoEntry := range reposFileEntries() {
				trimmed := strings.TrimSpace(repoEntry)
				if len(trimmed) > 0 {
					repoAddresses = append(repoAddresses, trimmed)
				}
			}
			if len(repoAddresses) == 0 {
				repoAddresses = append(repoAddresses, "github.com/cafebazaar/paks")
			}
		}
		sources := make([]installationSource, 0)
		for _, repoAddr := range repoAddresses {
			if len(repoAddr) > 0 {
				// TODO: change JoinURL
				repoPaks, err := pak.RepoFromURL(repoAddr)
				if err != nil {
					return cli.NewExitError(fmt.Sprintf("can't fetch the repo: %v", err.Error()), -1)
				}
				for _, pak := range repoPaks.Paks {
					if pak.Name == nameOfPakToInstall {
						pakURL = pak.URL
						if util.Relative(pakURL) {
							pakURL = util.JoinURL(repoAddr, pakURL)
						}
						sources = append(sources, installationSource{repoAddr, pakURL})
					}
				}
			}
		}
		if len(sources) > 1 {
			fmt.Println("Pak", nameOfPakToInstall, "exists in multiple repositories")
			fmt.Println("Enter the number of the pak you want to install")
			for i, source := range sources {
				fmt.Printf("%d) %q from repository %q\n", i, source.address, source.repo)
			}
			val, err := scanValue("default=0 ?", "int", false)
			if err != nil {
				return err
			}
			if choice, good := val.(int32); good {
				pakURL = sources[choice].address
			} else {
				pakURL = sources[0].address
			}
		}
	}

	p, err := pak.FromURL(pakURL)
	if err != nil {
		return cli.NewExitError(fmt.Sprintf("can't fetch the pak: %v", err.Error()), -1)
	}

	values := make(map[string]interface{})
	if c.Bool("interactive") || terminal.IsTerminal(int(os.Stdin.Fd())) {
		values, err = readValuesInteractively(p)
	} else {
		values, err = readValuesFromFile(p, valuesFile)
	}
	if err != nil {
		return cli.NewExitError(fmt.Sprintf("error in reading value: %v", err.Error()), -1)
	}

	_, err = pakManager.Install(p, c.GlobalString("namespace"), values)
	if err != nil {
		return cli.NewExitError(fmt.Sprintf("installation error: %v", err.Error()), -1)
	}
	return nil
}

func readValuesFromFile(p *pak.Pak, path string) (map[string]interface{}, error) {
	values := make(map[string]interface{})
	var valuesData []byte
	var err error
	if path == "" {
		valuesData, err = ioutil.ReadAll(os.Stdin)
	} else {
		valuesData, err = ioutil.ReadFile(path)
	}
	if err != nil {
		return nil, err
	}
	err = yaml.Unmarshal(valuesData, &values)
	return values, err
}

func readValuesInteractively(p *pak.Pak) (map[string]interface{}, error) {
	values := make(map[string]interface{})

	// ask for group
	groupValue, err := scanValue("Group Name (return for random): ", "string", false)
	if err != nil {
		return nil, err
	}
	if groupValue != nil {
		values["group"] = groupValue.(string)
	}

	// ask for all properties
	for i := range p.Properties {
		var prompt string
		if p.Properties[i].Default == nil {
			prompt = fmt.Sprintf("Field \"%s\" [type: %s]? ", p.Properties[i].Name, p.Properties[i].Type)
		} else {
			prompt = fmt.Sprintf("Field \"%s\" [type: %s, default: %v] (return for default)? ", p.Properties[i].Name, p.Properties[i].Type, p.Properties[i].Default)
		}
		value, err := scanValue(prompt, p.Properties[i].Type, p.Properties[i].Default == nil)
		if err != nil {
			return nil, err
		}
		if value != nil {
			values[p.Properties[i].Name] = value
		}
	}
	return values, nil
}

func scanValue(prompt string, valueType string, required bool) (interface{}, error) {
	bio := bufio.NewReader(os.Stdin)
	var value []byte
	var err error
	for {
		fmt.Printf(prompt)
		value, _, err = bio.ReadLine()
		if err != nil {
			return nil, err
		}
		if len(value) == 0 && !required {
			return nil, nil
		} else if len(value) == 0 && required {
			continue
		}

		switch valueType {
		case "int":
			i, err := strconv.Atoi(string(value))
			if err != nil {
				fmt.Println("given value is not an int, try again")
				continue
			}
			return i, nil
		case "bool":
			b, err := util.StringToBool(string(value))
			if err != nil {
				fmt.Println(err)
				continue
			}
			return b, nil
		case "string":
			fallthrough
		default:
			return string(value), nil
		}
	}
}
<|endoftext|>


<commit_before>package publish

import (
	"fmt"

	"github.com/jinzhu/gorm"
)

type Publish struct {
	*gorm.DB
	SupportedModels []interface{}
}

func Open(driver, source string) (*Publish, error) {
	db, err := gorm.Open(driver, source)
	return &Publish{DB: &db}, err
}

func DraftTableName(scope *gorm.Scope) string {
	return scope.TableName() + "_draft"
}

func (publish *Publish) Support(models ...interface{}) {
	fmt.Println(models[0])
	publish.SupportedModels = append(publish.SupportedModels, models...)
}

func (publish *Publish) AutoMigrateDrafts() {
	for _, value := range publish.SupportedModels {
		publish.Table(DraftTableName(&gorm.Scope{Value: value})).AutoMigrate(value)
	}
}

func (publish *Publish) ProductionMode() *gorm.DB {
	return publish.Set("publish_draft_mode", false)
}

func (publish *Publish) DraftMode() *gorm.DB {
	return publish.Set("publish_draft_mode", true)
}

func SetTable(scope *gorm.Scope) {
	var inDraft bool
	if draftMode, ok := scope.Get("publish_draft_mode"); ok {
		if value, ok := draftMode.(bool); ok && value {
			inDraft = true
		}
	}

	if inDraft {
		scope.Search.TableName = DraftTableName(scope)
	}
}
<commit_msg>[publish] Change setting key for publish<commit_after>package publish

import "github.com/jinzhu/gorm"

type Publish struct {
	*gorm.DB
	SupportedModels []interface{}
}

func Open(driver, source string) (*Publish, error) {
	db, err := gorm.Open(driver, source)
	return &Publish{DB: &db}, err
}

func DraftTableName(scope *gorm.Scope) string {
	return scope.TableName() + "_draft"
}

func (publish *Publish) Support(models ...interface{}) {
	publish.SupportedModels = append(publish.SupportedModels, models...)
}

func (publish *Publish) AutoMigrateDrafts() {
	for _, value := range publish.SupportedModels {
		publish.Table(DraftTableName(&gorm.Scope{Value: value})).AutoMigrate(value)
	}
}

func (publish *Publish) ProductionMode() *gorm.DB {
	return publish.Set("qor_publish:draft_mode", false)
}

func (publish *Publish) DraftMode() *gorm.DB {
	return publish.Set("qor_publish:draft_mode", true)
}

func SetTable(scope *gorm.Scope) {
	if draftMode, ok := scope.Get("qor_publish:draft_mode"); ok {
		if value, ok := draftMode.(bool); ok && value {
			scope.Search.TableName = DraftTableName(scope)
		}
	}
}
<|endoftext|>


<commit_before>package publish

import (
	"encoding/json"
	"fmt"
	"os"
	"sync/atomic"
	"time"

	MQTT "git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"
	"github.com/antlinker/go-cmap"
	"github.com/antlinker/mqttgroup/config"
	"gopkg.in/alog.v1"
	"gopkg.in/mgo.v2"
)

// Pub 执行Publish操作
func Pub(cfg *Config) {
	pub := &Publish{
		cfg:          cfg,
		lg:           alog.NewALog(),
		groupData:    make(map[string]int),
		clients:      make(map[string]*MQTT.Client),
		disClients:   cmap.NewConcurrencyMap(),
		execComplete: make(chan bool, 1),
		end:          make(chan bool, 1),
		startTime:    time.Now(),
	}
	pub.lg.SetLogTag("PUBLISH")
	session, err := mgo.Dial(cfg.MongoUrl)
	if err != nil {
		pub.lg.Errorf("数据库连接发生异常:%s", err.Error())
		return
	}
	pub.session = session
	pub.database = session.DB(config.Database)
	err = pub.Init()
	if err != nil {
		pub.lg.Error(err)
		return
	}
	pub.ExecPublish()
	<-pub.end
	pub.lg.InfoC("执行完成.")
	os.Exit(0)
}

type Publish struct {
	cfg               *Config
	lg                *alog.ALog
	session           *mgo.Session
	database          *mgo.Database
	userData          []config.User
	groupData         map[string]int
	clients           map[string]*MQTT.Client
	disClients        cmap.ConcurrencyMap
	publishID         int64
	execNum           int64
	execComplete      chan bool
	publishTotalNum   int64
	publishNum        int64
	prePublishNum     int64
	maxPublishNum     int64
	arrPublishNum     []int64
	receiveTotalNum   int64
	receiveNum        int64
	preReceiveNum     int64
	maxReceiveNum     int64
	arrReceiveNum     []int64
	end               chan bool
	startTime         time.Time
	publishPacketTime time.Time
}

func (p *Publish) Init() error {
	err := p.initData()
	if err != nil {
		return err
	}
	err = p.initConnection()
	if err != nil {
		return err
	}
	p.lg.InfoC("初始化操作完成.")
	p.publishPacketTime = time.Now()
	calcTicker := time.NewTicker(time.Second * 1)
	go p.calcMaxAndMinPacketNum(calcTicker)
	go p.checkComplete()
	prTicker := time.NewTicker(time.Duration(p.cfg.Interval) * time.Second)
	go p.pubAndRecOutput(prTicker)
	return nil
}

func (p *Publish) initData() error {
	p.lg.InfoC("开始组成员数据初始化...")
	err := p.database.C(config.CUser).Find(nil).All(&p.userData)
	if err != nil {
		return fmt.Errorf("Init group user data error:%s", err.Error())
	}
	for i, l := 0, len(p.userData); i < l; i++ {
		user := p.userData[i]
		for j := 0; j < len(user.Groups); j++ {
			gid := user.Groups[j]
			group, ok := p.groupData[gid]
			if ok {
				p.groupData[gid] = group + 1
				continue
			}
			p.groupData[gid] = 1
		}
	}
	p.lg.InfoC("组成员初始化完成.")
	return nil
}

func (p *Publish) initConnection() error {
	p.lg.InfoC("开始组成员建立MQTT数据连接初始化...")
	for i, l := 0, len(p.userData); i < l; i++ {
		user := p.userData[i]
		clientID := user.UserID
		opts := MQTT.NewClientOptions()
		opts.AddBroker(fmt.Sprintf("%s://%s", p.cfg.Network, p.cfg.Address))
		opts.SetClientID(clientID)
		opts.SetStore(MQTT.NewMemoryStore())
		opts.SetProtocolVersion(4)
		autoReconnect := true
		if v := p.cfg.DisconnectScale; v > 0 {
			autoReconnect = false
		}
		opts.SetAutoReconnect(autoReconnect)
		if name, pwd := p.cfg.UserName, p.cfg.Password; name != "" && pwd != "" {
			opts.SetUsername(name)
			opts.SetPassword(pwd)
		}
		if v := p.cfg.KeepAlive; v > 0 {
			opts.SetKeepAlive(time.Duration(v))
		}
		if v := p.cfg.CleanSession; v {
			opts.SetCleanSession(v)
		}
		clientHandle := NewHandleConnect(clientID, p)
		opts.SetConnectionLostHandler(func(cli *MQTT.Client, err error) {
			clientHandle.ErrorHandle(err)
		})
		cli := MQTT.NewClient(opts)
		connectToken := cli.Connect()
		if connectToken.Wait() && connectToken.Error() != nil {
			return fmt.Errorf("客户端[%s]建立连接发生异常:%s", clientID, connectToken.Error().Error())
		}
		subTopics := make(map[string]byte)
		for j := 0; j < len(user.Groups); j++ {
			topic := "G/" + user.Groups[j]
			subTopics[topic] = p.cfg.Qos
		}
		subToken := cli.SubscribeMultiple(subTopics, func(cli *MQTT.Client, msg MQTT.Message) {
			clientHandle.Subscribe([]byte(msg.Topic()), msg.Payload())
		})
		if subToken.Wait() && subToken.Error() != nil {
			return fmt.Errorf("客户端[%s]订阅主题发生异常:%s", clientID, subToken.Error().Error())
		}
		p.clients[clientID] = cli
	}
	p.lg.InfoC("组成员建立MQTT数据连接完成.")
	return nil
}

func (p *Publish) checkComplete() {
	<-p.execComplete
	go func() {
		ticker := time.NewTicker(time.Millisecond * 500)
		for range ticker.C {
			if p.receiveNum == p.receiveTotalNum {
				ticker.Stop()
				p.end <- true
			}
		}
	}()
}

func (p *Publish) calcMaxAndMinPacketNum(ticker *time.Ticker) {
	for range ticker.C {
		pubNum := p.publishNum
		recNum := p.receiveNum
		pNum := pubNum - p.prePublishNum
		rNum := recNum - p.preReceiveNum
		if pNum > p.maxPublishNum {
			p.maxPublishNum = pNum
		}
		p.arrPublishNum = append(p.arrPublishNum, pNum)
		if rNum > p.maxReceiveNum {
			p.maxReceiveNum = rNum
		}
		p.arrReceiveNum = append(p.arrReceiveNum, rNum)
		p.prePublishNum = pubNum
		p.preReceiveNum = recNum
	}
}

func (p *Publish) pubAndRecOutput(ticker *time.Ticker) {
	for ct := range ticker.C {
		totalSecond := float64(ct.Sub(p.startTime)) / float64(time.Second)
		packetSecond := (float64(ct.Sub(p.publishPacketTime)) - float64(time.Duration(p.cfg.Interval)*time.Second)) / float64(time.Second)
		var psNum int64
		arrPNum := p.arrPublishNum
		for i := 0; i < len(arrPNum); i++ {
			psNum += arrPNum[i]
		}
		avgPNum := psNum / int64(len(arrPNum))
		var rsNum int64
		arrRNum := p.arrReceiveNum
		for i := 0; i < len(arrRNum); i++ {
			rsNum += arrRNum[i]
		}
		avgRNum := rsNum / int64(len(arrRNum))
		clientNum := len(p.clients) - p.disClients.Len()
		output := `
总耗时                      %.2fs
发包耗时                    %.2fs
执行次数                    %d
客户端数量                  %d
应发包量                    %d
实际发包量                  %d
每秒平均的发包量            %d
每秒最大的发包量            %d
应收包量                    %d
实际收包量                  %d
每秒平均的接包量            %d
每秒最大的接包量            %d`

		fmt.Printf(output,
			totalSecond, packetSecond, p.execNum, clientNum, p.publishTotalNum, p.publishNum, avgPNum, p.maxPublishNum, p.receiveTotalNum, p.receiveNum, avgRNum, p.maxReceiveNum)
		fmt.Printf("\n")
	}
}

func (p *Publish) ExecPublish() {
	time.AfterFunc(time.Duration(p.cfg.Interval)*time.Second, func() {
		if p.execNum == int64(p.cfg.ExecNum) {
			p.lg.InfoC("发布已执行完成，等待接收订阅...")
			p.disconnect()
			p.execComplete <- true
			return
		}
		atomic.AddInt64(&p.execNum, 1)
		p.publish()
		p.ExecPublish()
	})
}

func (p *Publish) disconnect() {
	if v := p.cfg.DisconnectScale; v > 100 || v <= 0 {
		return
	}
	clientCount := len(p.clients)
	disCount := int(float32(clientCount) * (float32(p.cfg.DisconnectScale) / 100))
	for k, v := range p.clients {
		exist, _ := p.disClients.Contains(k)
		if exist {
			continue
		}
		v.Disconnect(100)
		disCount--
		if disCount == 0 {
			break
		}
	}
}

func (p *Publish) publish() {
	for i, l := 0, len(p.userData); i < l; i++ {
		user := p.userData[i]
		groups := user.Groups
		for j := 0; j < len(groups); j++ {
			p.publishTotalNum++
			p.receiveTotalNum += int64(p.groupData[groups[j]])
		}
		go p.userPublish(user)
		if v := p.cfg.UserInterval; v > 0 {
			time.Sleep(time.Duration(v) * time.Millisecond)
		}
	}
}

func (p *Publish) userPublish(user config.User) {
	v, _ := p.disClients.Contains(user.UserID)
	if v {
		return
	}
	cli := p.clients[user.UserID]
	for i := 0; i < len(user.Groups); i++ {
		group := user.Groups[i]
		cTime := time.Now()
		sendPacket := config.SendPacket{
			SendID:   fmt.Sprintf("%d_%d", cTime.Unix(), atomic.AddInt64(&p.publishID, 1)),
			SendUser: user.UserID,
			ToGroup:  group,
			SendTime: cTime,
			Receives: make([]config.ReceivePacket, 0),
		}
		if p.cfg.IsStore {
			err := p.database.C(config.CPacket).Insert(sendPacket)
			if err != nil {
				p.lg.Errorf("Publish store error:%s", err.Error())
			}
		}
		bufData, err := json.Marshal(sendPacket)
		if err != nil {
			p.lg.Errorf("Publish json marshal error:%s", err.Error())
			continue
		}
		topic := "G/" + group
		publishToken := cli.Publish(topic, p.cfg.Qos, false, bufData)
		if publishToken.Wait() && publishToken.Error() != nil {
			p.lg.Errorf("用户[%s]发布消息出现异常:%s", user.UserID, publishToken.Error().Error())
			continue
		}
		atomic.AddInt64(&p.publishNum, 1)
		if v := p.cfg.UserGroupInterval; v > 0 {
			time.Sleep(time.Duration(v) * time.Millisecond)
		}
	}
}
<commit_msg>修复保持心跳bug<commit_after>package publish

import (
	"encoding/json"
	"fmt"
	"os"
	"sync/atomic"
	"time"

	MQTT "git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"
	"github.com/antlinker/go-cmap"
	"github.com/antlinker/mqttgroup/config"
	"gopkg.in/alog.v1"
	"gopkg.in/mgo.v2"
)

// Pub 执行Publish操作
func Pub(cfg *Config) {
	pub := &Publish{
		cfg:          cfg,
		lg:           alog.NewALog(),
		groupData:    make(map[string]int),
		clients:      make(map[string]*MQTT.Client),
		disClients:   cmap.NewConcurrencyMap(),
		execComplete: make(chan bool, 1),
		end:          make(chan bool, 1),
		startTime:    time.Now(),
	}
	pub.lg.SetLogTag("PUBLISH")
	session, err := mgo.Dial(cfg.MongoUrl)
	if err != nil {
		pub.lg.Errorf("数据库连接发生异常:%s", err.Error())
		return
	}
	pub.session = session
	pub.database = session.DB(config.Database)
	err = pub.Init()
	if err != nil {
		pub.lg.Error(err)
		return
	}
	pub.ExecPublish()
	<-pub.end
	pub.lg.InfoC("执行完成.")
	os.Exit(0)
}

type Publish struct {
	cfg               *Config
	lg                *alog.ALog
	session           *mgo.Session
	database          *mgo.Database
	userData          []config.User
	groupData         map[string]int
	clients           map[string]*MQTT.Client
	disClients        cmap.ConcurrencyMap
	publishID         int64
	execNum           int64
	execComplete      chan bool
	publishTotalNum   int64
	publishNum        int64
	prePublishNum     int64
	maxPublishNum     int64
	arrPublishNum     []int64
	receiveTotalNum   int64
	receiveNum        int64
	preReceiveNum     int64
	maxReceiveNum     int64
	arrReceiveNum     []int64
	end               chan bool
	startTime         time.Time
	publishPacketTime time.Time
}

func (p *Publish) Init() error {
	err := p.initData()
	if err != nil {
		return err
	}
	err = p.initConnection()
	if err != nil {
		return err
	}
	p.lg.InfoC("初始化操作完成.")
	p.publishPacketTime = time.Now()
	calcTicker := time.NewTicker(time.Second * 1)
	go p.calcMaxAndMinPacketNum(calcTicker)
	go p.checkComplete()
	prTicker := time.NewTicker(time.Duration(p.cfg.Interval) * time.Second)
	go p.pubAndRecOutput(prTicker)
	return nil
}

func (p *Publish) initData() error {
	p.lg.InfoC("开始组成员数据初始化...")
	err := p.database.C(config.CUser).Find(nil).All(&p.userData)
	if err != nil {
		return fmt.Errorf("Init group user data error:%s", err.Error())
	}
	for i, l := 0, len(p.userData); i < l; i++ {
		user := p.userData[i]
		for j := 0; j < len(user.Groups); j++ {
			gid := user.Groups[j]
			group, ok := p.groupData[gid]
			if ok {
				p.groupData[gid] = group + 1
				continue
			}
			p.groupData[gid] = 1
		}
	}
	p.lg.InfoC("组成员初始化完成.")
	return nil
}

func (p *Publish) initConnection() error {
	p.lg.InfoC("开始组成员建立MQTT数据连接初始化...")
	for i, l := 0, len(p.userData); i < l; i++ {
		user := p.userData[i]
		clientID := user.UserID
		opts := MQTT.NewClientOptions()
		opts.AddBroker(fmt.Sprintf("%s://%s", p.cfg.Network, p.cfg.Address))
		opts.SetClientID(clientID)
		opts.SetStore(MQTT.NewMemoryStore())
		opts.SetProtocolVersion(4)
		autoReconnect := true
		if v := p.cfg.DisconnectScale; v > 0 {
			autoReconnect = false
		}
		opts.SetAutoReconnect(autoReconnect)
		if name, pwd := p.cfg.UserName, p.cfg.Password; name != "" && pwd != "" {
			opts.SetUsername(name)
			opts.SetPassword(pwd)
		}
		if v := p.cfg.KeepAlive; v > 0 {
			opts.SetKeepAlive(time.Duration(v) * time.Second)
		}
		if v := p.cfg.CleanSession; v {
			opts.SetCleanSession(v)
		}
		clientHandle := NewHandleConnect(clientID, p)
		opts.SetConnectionLostHandler(func(cli *MQTT.Client, err error) {
			clientHandle.ErrorHandle(err)
		})
		cli := MQTT.NewClient(opts)
		connectToken := cli.Connect()
		if connectToken.Wait() && connectToken.Error() != nil {
			return fmt.Errorf("客户端[%s]建立连接发生异常:%s", clientID, connectToken.Error().Error())
		}
		subTopics := make(map[string]byte)
		for j := 0; j < len(user.Groups); j++ {
			topic := "G/" + user.Groups[j]
			subTopics[topic] = p.cfg.Qos
		}
		subToken := cli.SubscribeMultiple(subTopics, func(cli *MQTT.Client, msg MQTT.Message) {
			clientHandle.Subscribe([]byte(msg.Topic()), msg.Payload())
		})
		if subToken.Wait() && subToken.Error() != nil {
			return fmt.Errorf("客户端[%s]订阅主题发生异常:%s", clientID, subToken.Error().Error())
		}
		p.clients[clientID] = cli
	}
	p.lg.InfoC("组成员建立MQTT数据连接完成.")
	return nil
}

func (p *Publish) checkComplete() {
	<-p.execComplete
	go func() {
		ticker := time.NewTicker(time.Millisecond * 500)
		for range ticker.C {
			if p.receiveNum == p.receiveTotalNum {
				ticker.Stop()
				p.end <- true
			}
		}
	}()
}

func (p *Publish) calcMaxAndMinPacketNum(ticker *time.Ticker) {
	for range ticker.C {
		pubNum := p.publishNum
		recNum := p.receiveNum
		pNum := pubNum - p.prePublishNum
		rNum := recNum - p.preReceiveNum
		if pNum > p.maxPublishNum {
			p.maxPublishNum = pNum
		}
		p.arrPublishNum = append(p.arrPublishNum, pNum)
		if rNum > p.maxReceiveNum {
			p.maxReceiveNum = rNum
		}
		p.arrReceiveNum = append(p.arrReceiveNum, rNum)
		p.prePublishNum = pubNum
		p.preReceiveNum = recNum
	}
}

func (p *Publish) pubAndRecOutput(ticker *time.Ticker) {
	for ct := range ticker.C {
		totalSecond := float64(ct.Sub(p.startTime)) / float64(time.Second)
		packetSecond := (float64(ct.Sub(p.publishPacketTime)) - float64(time.Duration(p.cfg.Interval)*time.Second)) / float64(time.Second)
		var psNum int64
		arrPNum := p.arrPublishNum
		for i := 0; i < len(arrPNum); i++ {
			psNum += arrPNum[i]
		}
		avgPNum := psNum / int64(len(arrPNum))
		var rsNum int64
		arrRNum := p.arrReceiveNum
		for i := 0; i < len(arrRNum); i++ {
			rsNum += arrRNum[i]
		}
		avgRNum := rsNum / int64(len(arrRNum))
		clientNum := len(p.clients) - p.disClients.Len()
		output := `
总耗时                      %.2fs
发包耗时                    %.2fs
执行次数                    %d
客户端数量                  %d
应发包量                    %d
实际发包量                  %d
每秒平均的发包量            %d
每秒最大的发包量            %d
应收包量                    %d
实际收包量                  %d
每秒平均的接包量            %d
每秒最大的接包量            %d`

		fmt.Printf(output,
			totalSecond, packetSecond, p.execNum, clientNum, p.publishTotalNum, p.publishNum, avgPNum, p.maxPublishNum, p.receiveTotalNum, p.receiveNum, avgRNum, p.maxReceiveNum)
		fmt.Printf("\n")
	}
}

func (p *Publish) ExecPublish() {
	time.AfterFunc(time.Duration(p.cfg.Interval)*time.Second, func() {
		if p.execNum == int64(p.cfg.ExecNum) {
			p.lg.InfoC("发布已执行完成，等待接收订阅...")
			p.disconnect()
			p.execComplete <- true
			return
		}
		atomic.AddInt64(&p.execNum, 1)
		p.publish()
		p.ExecPublish()
	})
}

func (p *Publish) disconnect() {
	if v := p.cfg.DisconnectScale; v > 100 || v <= 0 {
		return
	}
	clientCount := len(p.clients)
	disCount := int(float32(clientCount) * (float32(p.cfg.DisconnectScale) / 100))
	for k, v := range p.clients {
		exist, _ := p.disClients.Contains(k)
		if exist {
			continue
		}
		v.Disconnect(100)
		disCount--
		if disCount == 0 {
			break
		}
	}
}

func (p *Publish) publish() {
	for i, l := 0, len(p.userData); i < l; i++ {
		user := p.userData[i]
		groups := user.Groups
		for j := 0; j < len(groups); j++ {
			p.publishTotalNum++
			p.receiveTotalNum += int64(p.groupData[groups[j]])
		}
		go p.userPublish(user)
		if v := p.cfg.UserInterval; v > 0 {
			time.Sleep(time.Duration(v) * time.Millisecond)
		}
	}
}

func (p *Publish) userPublish(user config.User) {
	v, _ := p.disClients.Contains(user.UserID)
	if v {
		return
	}
	cli := p.clients[user.UserID]
	for i := 0; i < len(user.Groups); i++ {
		group := user.Groups[i]
		cTime := time.Now()
		sendPacket := config.SendPacket{
			SendID:   fmt.Sprintf("%d_%d", cTime.Unix(), atomic.AddInt64(&p.publishID, 1)),
			SendUser: user.UserID,
			ToGroup:  group,
			SendTime: cTime,
			Receives: make([]config.ReceivePacket, 0),
		}
		if p.cfg.IsStore {
			err := p.database.C(config.CPacket).Insert(sendPacket)
			if err != nil {
				p.lg.Errorf("Publish store error:%s", err.Error())
			}
		}
		bufData, err := json.Marshal(sendPacket)
		if err != nil {
			p.lg.Errorf("Publish json marshal error:%s", err.Error())
			continue
		}
		topic := "G/" + group
		publishToken := cli.Publish(topic, p.cfg.Qos, false, bufData)
		if publishToken.Wait() && publishToken.Error() != nil {
			p.lg.Errorf("用户[%s]发布消息出现异常:%s", user.UserID, publishToken.Error().Error())
			continue
		}
		atomic.AddInt64(&p.publishNum, 1)
		if v := p.cfg.UserGroupInterval; v > 0 {
			time.Sleep(time.Duration(v) * time.Millisecond)
		}
	}
}
<|endoftext|>


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

// Package blockfile provides methods for reading packets from blockfiles
// generated by stenotype.
package blockfile

import (
	"errors"
	"fmt"
	"io"
	"sync"
	"time"
	"unsafe"

	"github.com/google/gopacket"
	"github.com/google/stenographer/base"
	"github.com/google/stenographer/filecache"
	"github.com/google/stenographer/indexfile"
	"github.com/google/stenographer/query"
	"github.com/google/stenographer/stats"
	"golang.org/x/net/context"
)

// #include <linux/if_packet.h>
import "C"

var (
	v                = base.V // Verbose logging
	packetReadNanos  = stats.S.Get("packet_read_nanos")
	packetScanNanos  = stats.S.Get("packet_scan_nanos")
	packetsRead      = stats.S.Get("packets_read")
	packetsScanned   = stats.S.Get("packets_scanned")
	packetBlocksRead = stats.S.Get("packets_blocks_read")
)

// BlockFile provides an interface to a single stenotype file on disk and its
// associated index.
type BlockFile struct {
	name string
	f    *filecache.CachedFile
	i    *indexfile.IndexFile
	mu   sync.RWMutex // Stops Close() from invalidating a file before a current query is done with it.
	done chan struct{}
}

// NewBlockFile opens up a named block file (and its index), returning a handle
// which can be used to look up packets.
func NewBlockFile(filename string, fc *filecache.Cache) (*BlockFile, error) {
	v(1, "Blockfile opening: %q", filename)
	i, err := indexfile.NewIndexFile(indexfile.IndexPathFromBlockfilePath(filename), fc)
	if err != nil {
		return nil, fmt.Errorf("could not open index for %q: %v", filename, err)
	}
	return &BlockFile{
		f:    fc.Open(filename),
		i:    i,
		name: filename,
		done: make(chan struct{}),
	}, nil
}

// Name returns the name of the file underlying this blockfile.
func (b *BlockFile) Name() string {
	return b.name
}

// readPacket reads a single packet from the file at the given position.
// It updates the passed in CaptureInfo with information on the packet.
func (b *BlockFile) readPacket(pos int64, ci *gopacket.CaptureInfo) ([]byte, error) {
	// 28 bytes actually isn't the entire packet header, but it's all the fields
	// that we care about.
	packetsRead.Increment()
	defer packetReadNanos.NanoTimer()()
	var dataBuf [28]byte
	_, err := b.f.ReadAt(dataBuf[:], pos)
	if err != nil {
		return nil, err
	}
	pkt := (*C.struct_tpacket3_hdr)(unsafe.Pointer(&dataBuf[0]))
	*ci = gopacket.CaptureInfo{
		Timestamp:     time.Unix(int64(pkt.tp_sec), int64(pkt.tp_nsec)),
		Length:        int(pkt.tp_len),
		CaptureLength: int(pkt.tp_snaplen),
	}
	out := make([]byte, ci.CaptureLength)
	pos += int64(pkt.tp_mac)
	_, err = b.f.ReadAt(out, pos)
	return out, err
}

// Close cleans up this blockfile.
func (b *BlockFile) Close() (err error) {
	v(2, "Blockfile closing: %q", b.name)
	close(b.done)
	b.mu.Lock()
	defer b.mu.Unlock()
	v(3, "Blockfile closing file descriptors: %q", b.name)
	if e := b.i.Close(); e != nil {
		err = e
	}
	if e := b.f.Close(); e != nil {
		err = e
	}
	b.i, b.f = nil, nil
	return
}

// allPacketsIter implements Iter.
type allPacketsIter struct {
	*BlockFile
	blockData        []byte
	block            *C.struct_tpacket_hdr_v1
	pkt              *C.struct_tpacket3_hdr
	blockPacketsRead int
	blockOffset      int64
	packetOffset     int // offset of packet in block
	err              error
	done             bool
}

func (a *allPacketsIter) Next() bool {
	defer packetScanNanos.NanoTimer()()
	if a.err != nil || a.done {
		return false
	}
	for a.block == nil || a.blockPacketsRead == int(a.block.num_pkts) {
		packetBlocksRead.Increment()
		a.blockData = make([]byte, 1 << 20)
		_, err := a.f.ReadAt(a.blockData[:], a.blockOffset)
		if err == io.EOF {
			a.done = true
			return false
		} else if err != nil {
			a.err = fmt.Errorf("could not read block at %v: %v", a.blockOffset, err)
			return false
		}
		baseHdr := (*C.struct_tpacket_block_desc)(unsafe.Pointer(&a.blockData[0]))
		a.block = (*C.struct_tpacket_hdr_v1)(unsafe.Pointer(&baseHdr.hdr[0]))
		a.blockOffset += 1 << 20
		a.blockPacketsRead = 0
		a.pkt = nil
	}
	a.blockPacketsRead++
	if a.pkt == nil {
		a.packetOffset = int(a.block.offset_to_first_pkt)
	} else if a.pkt.tp_next_offset != 0 {
		a.packetOffset += int(a.pkt.tp_next_offset)
	} else {
		a.err = errors.New("block format currently not supported")
		return false
	}
	a.pkt = (*C.struct_tpacket3_hdr)(unsafe.Pointer(&a.blockData[a.packetOffset]))
	packetsScanned.Increment()
	return true
}

func (a *allPacketsIter) Packet() *base.Packet {
	start := a.packetOffset + int(a.pkt.tp_mac)
	buf := a.blockData[start : start+int(a.pkt.tp_snaplen)]
	p := &base.Packet{Data: buf}
	p.CaptureInfo.Timestamp = time.Unix(int64(a.pkt.tp_sec), int64(a.pkt.tp_nsec))
	p.CaptureInfo.Length = int(a.pkt.tp_len)
	p.CaptureInfo.CaptureLength = int(a.pkt.tp_snaplen)
	return p
}

func (a *allPacketsIter) Err() error {
	return a.err
}

// AllPackets returns a packet channel to which all packets in the blockfile are
// sent.
func (b *BlockFile) AllPackets() *base.PacketChan {
	b.mu.RLock()
	c := base.NewPacketChan(100)
	go func() {
		defer b.mu.RUnlock()
		pkts := &allPacketsIter{BlockFile: b}
		for pkts.Next() {
			c.Send(pkts.Packet())
		}
		c.Close(pkts.Err())
	}()
	return c
}

// Positions returns the positions in the blockfile of all packets matched by
// the passed-in query.
func (b *BlockFile) Positions(ctx context.Context, q query.Query) (base.Positions, error) {
	b.mu.RLock()
	defer b.mu.RUnlock()
	return b.positionsLocked(ctx, q)
}

// positionsLocked returns the positions in the blockfile of all packets matched by
// the passed-in query.  b.mu must be locked.
func (b *BlockFile) positionsLocked(ctx context.Context, q query.Query) (base.Positions, error) {
	if b.i == nil || b.f == nil {
		// If we're closed, just return nothing.
		return nil, nil
	}
	return q.LookupIn(ctx, b.i)
}

// Lookup returns all packets in the blockfile matched by the passed-in query.
func (b *BlockFile) Lookup(ctx context.Context, q query.Query, out *base.PacketChan) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	var ci gopacket.CaptureInfo
	v(2, "Blockfile %q looking up query %q", q.String(), b.name)
	start := time.Now()
	positions, err := b.positionsLocked(ctx, q)
	if err != nil {
		out.Close(fmt.Errorf("index lookup failure: %v", err))
		return
	}
	if positions.IsAllPositions() {
		v(2, "Blockfile %q reading all packets", b.name)
		iter := &allPacketsIter{BlockFile: b}
	all_packets_loop:
		for iter.Next() {
			select {
			case <-ctx.Done():
				v(2, "Blockfile %q canceling packet read", b.name)
				break all_packets_loop
			case <-b.done:
				v(2, "Blockfile %q closing, breaking out of query", b.name)
				break all_packets_loop
			case out.C <- iter.Packet():
			}
		}
		if iter.Err() != nil {
			out.Close(fmt.Errorf("error reading all packets from %q: %v", b.name, iter.Err()))
			return
		}
	} else {
		v(2, "Blockfile %q reading %v packets", b.name, len(positions))
	query_packets_loop:
		for _, pos := range positions {
			buffer, err := b.readPacket(pos, &ci)
			if err != nil {
				v(2, "Blockfile %q error reading packet: %v", b.name, err)
				out.Close(fmt.Errorf("error reading packets from %q @ %v: %v", b.name, pos, err))
				return
			}
			select {
			case <-ctx.Done():
				v(2, "Blockfile %q canceling packet read", b.name)
				break query_packets_loop
			case <-b.done:
				v(2, "Blockfile %q closing, breaking out of query", b.name)
				break query_packets_loop
			case out.C <- &base.Packet{Data: buffer, CaptureInfo: ci}:
			}
		}
	}
	v(2, "Blockfile %q finished reading all packets in %v", b.name, time.Since(start))
	out.Close(ctx.Err())
}

// DumpIndex dumps out a "human-readable" debug version of the blockfile's index
// to the given writer.
func (b *BlockFile) DumpIndex(out io.Writer, start, finish []byte) {
	b.mu.RLock()
	defer b.mu.RUnlock()
	b.i.Dump(out, start, finish)
}
<commit_msg>blockfile: Correct order of strings in debug log<commit_after>// Copyright 2014 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package blockfile provides methods for reading packets from blockfiles
// generated by stenotype.
package blockfile

import (
	"errors"
	"fmt"
	"io"
	"sync"
	"time"
	"unsafe"

	"github.com/google/gopacket"
	"github.com/google/stenographer/base"
	"github.com/google/stenographer/filecache"
	"github.com/google/stenographer/indexfile"
	"github.com/google/stenographer/query"
	"github.com/google/stenographer/stats"
	"golang.org/x/net/context"
)

// #include <linux/if_packet.h>
import "C"

var (
	v                = base.V // Verbose logging
	packetReadNanos  = stats.S.Get("packet_read_nanos")
	packetScanNanos  = stats.S.Get("packet_scan_nanos")
	packetsRead      = stats.S.Get("packets_read")
	packetsScanned   = stats.S.Get("packets_scanned")
	packetBlocksRead = stats.S.Get("packets_blocks_read")
)

// BlockFile provides an interface to a single stenotype file on disk and its
// associated index.
type BlockFile struct {
	name string
	f    *filecache.CachedFile
	i    *indexfile.IndexFile
	mu   sync.RWMutex // Stops Close() from invalidating a file before a current query is done with it.
	done chan struct{}
}

// NewBlockFile opens up a named block file (and its index), returning a handle
// which can be used to look up packets.
func NewBlockFile(filename string, fc *filecache.Cache) (*BlockFile, error) {
	v(1, "Blockfile opening: %q", filename)
	i, err := indexfile.NewIndexFile(indexfile.IndexPathFromBlockfilePath(filename), fc)
	if err != nil {
		return nil, fmt.Errorf("could not open index for %q: %v", filename, err)
	}
	return &BlockFile{
		f:    fc.Open(filename),
		i:    i,
		name: filename,
		done: make(chan struct{}),
	}, nil
}

// Name returns the name of the file underlying this blockfile.
func (b *BlockFile) Name() string {
	return b.name
}

// readPacket reads a single packet from the file at the given position.
// It updates the passed in CaptureInfo with information on the packet.
func (b *BlockFile) readPacket(pos int64, ci *gopacket.CaptureInfo) ([]byte, error) {
	// 28 bytes actually isn't the entire packet header, but it's all the fields
	// that we care about.
	packetsRead.Increment()
	defer packetReadNanos.NanoTimer()()
	var dataBuf [28]byte
	_, err := b.f.ReadAt(dataBuf[:], pos)
	if err != nil {
		return nil, err
	}
	pkt := (*C.struct_tpacket3_hdr)(unsafe.Pointer(&dataBuf[0]))
	*ci = gopacket.CaptureInfo{
		Timestamp:     time.Unix(int64(pkt.tp_sec), int64(pkt.tp_nsec)),
		Length:        int(pkt.tp_len),
		CaptureLength: int(pkt.tp_snaplen),
	}
	out := make([]byte, ci.CaptureLength)
	pos += int64(pkt.tp_mac)
	_, err = b.f.ReadAt(out, pos)
	return out, err
}

// Close cleans up this blockfile.
func (b *BlockFile) Close() (err error) {
	v(2, "Blockfile closing: %q", b.name)
	close(b.done)
	b.mu.Lock()
	defer b.mu.Unlock()
	v(3, "Blockfile closing file descriptors: %q", b.name)
	if e := b.i.Close(); e != nil {
		err = e
	}
	if e := b.f.Close(); e != nil {
		err = e
	}
	b.i, b.f = nil, nil
	return
}

// allPacketsIter implements Iter.
type allPacketsIter struct {
	*BlockFile
	blockData        []byte
	block            *C.struct_tpacket_hdr_v1
	pkt              *C.struct_tpacket3_hdr
	blockPacketsRead int
	blockOffset      int64
	packetOffset     int // offset of packet in block
	err              error
	done             bool
}

func (a *allPacketsIter) Next() bool {
	defer packetScanNanos.NanoTimer()()
	if a.err != nil || a.done {
		return false
	}
	for a.block == nil || a.blockPacketsRead == int(a.block.num_pkts) {
		packetBlocksRead.Increment()
		a.blockData = make([]byte, 1<<20)
		_, err := a.f.ReadAt(a.blockData[:], a.blockOffset)
		if err == io.EOF {
			a.done = true
			return false
		} else if err != nil {
			a.err = fmt.Errorf("could not read block at %v: %v", a.blockOffset, err)
			return false
		}
		baseHdr := (*C.struct_tpacket_block_desc)(unsafe.Pointer(&a.blockData[0]))
		a.block = (*C.struct_tpacket_hdr_v1)(unsafe.Pointer(&baseHdr.hdr[0]))
		a.blockOffset += 1 << 20
		a.blockPacketsRead = 0
		a.pkt = nil
	}
	a.blockPacketsRead++
	if a.pkt == nil {
		a.packetOffset = int(a.block.offset_to_first_pkt)
	} else if a.pkt.tp_next_offset != 0 {
		a.packetOffset += int(a.pkt.tp_next_offset)
	} else {
		a.err = errors.New("block format currently not supported")
		return false
	}
	a.pkt = (*C.struct_tpacket3_hdr)(unsafe.Pointer(&a.blockData[a.packetOffset]))
	packetsScanned.Increment()
	return true
}

func (a *allPacketsIter) Packet() *base.Packet {
	start := a.packetOffset + int(a.pkt.tp_mac)
	buf := a.blockData[start : start+int(a.pkt.tp_snaplen)]
	p := &base.Packet{Data: buf}
	p.CaptureInfo.Timestamp = time.Unix(int64(a.pkt.tp_sec), int64(a.pkt.tp_nsec))
	p.CaptureInfo.Length = int(a.pkt.tp_len)
	p.CaptureInfo.CaptureLength = int(a.pkt.tp_snaplen)
	return p
}

func (a *allPacketsIter) Err() error {
	return a.err
}

// AllPackets returns a packet channel to which all packets in the blockfile are
// sent.
func (b *BlockFile) AllPackets() *base.PacketChan {
	b.mu.RLock()
	c := base.NewPacketChan(100)
	go func() {
		defer b.mu.RUnlock()
		pkts := &allPacketsIter{BlockFile: b}
		for pkts.Next() {
			c.Send(pkts.Packet())
		}
		c.Close(pkts.Err())
	}()
	return c
}

// Positions returns the positions in the blockfile of all packets matched by
// the passed-in query.
func (b *BlockFile) Positions(ctx context.Context, q query.Query) (base.Positions, error) {
	b.mu.RLock()
	defer b.mu.RUnlock()
	return b.positionsLocked(ctx, q)
}

// positionsLocked returns the positions in the blockfile of all packets matched by
// the passed-in query.  b.mu must be locked.
func (b *BlockFile) positionsLocked(ctx context.Context, q query.Query) (base.Positions, error) {
	if b.i == nil || b.f == nil {
		// If we're closed, just return nothing.
		return nil, nil
	}
	return q.LookupIn(ctx, b.i)
}

// Lookup returns all packets in the blockfile matched by the passed-in query.
func (b *BlockFile) Lookup(ctx context.Context, q query.Query, out *base.PacketChan) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	var ci gopacket.CaptureInfo
	v(2, "Blockfile %q looking up query %q", b.name, q.String())
	start := time.Now()
	positions, err := b.positionsLocked(ctx, q)
	if err != nil {
		out.Close(fmt.Errorf("index lookup failure: %v", err))
		return
	}
	if positions.IsAllPositions() {
		v(2, "Blockfile %q reading all packets", b.name)
		iter := &allPacketsIter{BlockFile: b}
	all_packets_loop:
		for iter.Next() {
			select {
			case <-ctx.Done():
				v(2, "Blockfile %q canceling packet read", b.name)
				break all_packets_loop
			case <-b.done:
				v(2, "Blockfile %q closing, breaking out of query", b.name)
				break all_packets_loop
			case out.C <- iter.Packet():
			}
		}
		if iter.Err() != nil {
			out.Close(fmt.Errorf("error reading all packets from %q: %v", b.name, iter.Err()))
			return
		}
	} else {
		v(2, "Blockfile %q reading %v packets", b.name, len(positions))
	query_packets_loop:
		for _, pos := range positions {
			buffer, err := b.readPacket(pos, &ci)
			if err != nil {
				v(2, "Blockfile %q error reading packet: %v", b.name, err)
				out.Close(fmt.Errorf("error reading packets from %q @ %v: %v", b.name, pos, err))
				return
			}
			select {
			case <-ctx.Done():
				v(2, "Blockfile %q canceling packet read", b.name)
				break query_packets_loop
			case <-b.done:
				v(2, "Blockfile %q closing, breaking out of query", b.name)
				break query_packets_loop
			case out.C <- &base.Packet{Data: buffer, CaptureInfo: ci}:
			}
		}
	}
	v(2, "Blockfile %q finished reading all packets in %v", b.name, time.Since(start))
	out.Close(ctx.Err())
}

// DumpIndex dumps out a "human-readable" debug version of the blockfile's index
// to the given writer.
func (b *BlockFile) DumpIndex(out io.Writer, start, finish []byte) {
	b.mu.RLock()
	defer b.mu.RUnlock()
	b.i.Dump(out, start, finish)
}
<|endoftext|>


<commit_before><commit_msg>cmd/ebitenmobile: Use Go modules in gomobile<commit_after><|endoftext|>