补充功能

This commit is contained in:
chenbox
2017-07-23 16:35:21 +08:00
parent 4fe6708a00
commit b28b25dc71
14 changed files with 229 additions and 164 deletions
-4
View File
@@ -81,10 +81,6 @@ func WHandler(conn util.NetConn) {
for {
select {
case msg := <-Wch:
// prefix := string(msg[0:2])
// if prefix != "00" && prefix != "01" && prefix != "11" {
// log.Println("not equal", string(msg))
// } else {
l, err := conn.Write(msg)
if err != nil {
log.Println("写到Nat目录连接出错:", err.Error())
+2
View File
@@ -6,6 +6,8 @@ const (
PROXY_CONTROL_NORMAL byte = 0 // 控制码 正常发送
PROXY_CONTROL_ACK byte = 12 // 控制码 确认
PROXY_CONTROL_QUIT byte = 10 // 控制码 退出
PROXY_CONTROL_HEARTBIT byte = 13 // 控制码 心跳包
PROXY_CONTROL_HEARTBITACK byte = 14 // 心跳包确认
PROXY_CONTROL_ERROR_NO_SERVER byte = 201 // 服务端还没有注册
PROXY_CONTROL_ERROR_CLIENT_EXIST byte = 202 // 客户端已经存在
PROXY_CONTROL_ERROR_SERVER_EXIST byte = 203 // 服务端已经存在
+1 -1
View File
@@ -1,5 +1,5 @@
package main
import "punching/client"
func main() {
client.Main()
}
+1 -1
View File
@@ -1,5 +1,5 @@
package main
import "punching/proxy"
func main() {
proxy.Main()
}
+1 -1
View File
@@ -1,5 +1,5 @@
package main
import "punching/server"
func main() {
server.Main()
}
+55 -67
View File
@@ -28,7 +28,7 @@ type ClientConn struct{
// 全局变量
var (
OnlineServerList map[string]*ServerConn // 服务端连接列表Map
OnlineClientList map[string]*ClientConn // 客户端连接列表Map
OnlineClientList map[string]string // 客户端连接列表Map
RWLockClient *sync.RWMutex //读写锁
RWLockServer *sync.RWMutex
)
@@ -42,7 +42,7 @@ func Main(){
}
OnlineServerList = make(map[string]*ServerConn)
OnlineClientList = make(map[string]*ClientConn)
OnlineClientList = make(map[string]string)
RWLockClient = new(sync.RWMutex)
RWLockServer = new(sync.RWMutex)
@@ -72,90 +72,82 @@ func Main(){
}
}
// parseFirstPackage 解析连接的第一个条
// [1]+[1]+[4-32] (包头标识+类型+32字节的Pairname)
func parseFirstPackage(data []byte)(pairName string, roleType int, err error){
roleType = int(data[1])
if data[0] != PROXY_PACKAGE_HEAD || ( roleType != ROLE_CLIENT && roleType != ROLE_SERVER) {
err = fmt.Errorf("%s","包内容不匹配")
return
}
pairName = string(byte[2:])
return
}
// processRoleClient 处理客户端连接
func processRoleClient(conn net.Connpairname string ){
func processRoleClient(conn net.Conn, pairName string ){
// 判断匹配的服务端是否已经注册
RWLockServer.RLock()
serverConn,ok := OnlineServerList[pairname]
serverConn,ok := OnlineServerList[pairName]
RWLockServer.RUnlock()
if !ok {
// 客户端没有注册
conn.Write([]byte(PROXY_RESP_NO_SERVER))
packErr := util.PackageProxy(PROXY_CONTROL_ERROR_NO_SERVER, []byte(""))
conn.Write(packErr)
return
}
// Check the client with the save pair name exists
RWLockClient.RLock()
ok, _ := OnlineClientList[pairname]
_, ok = OnlineClientList[pairName]
RWLockClient.RUnlock()
if !ok{
conn.Write([]byte(PROXY_RESP_CLIENT_EXIST))
packErr := util.PackageProxy(PROXY_CONTROL_ERROR_CLIENT_EXIST,[]byte(""))
conn.Write(packErr)
return
}
// 添加到客户端列表
RWLockClient.Lock()
OnlineClientList[pairname] = pairname
OnlineClientList[pairName] = pairName
RWLockClient.Unlock()
// byteNat := bytes.NewBuffer(nil)
// byteNat.write([]byte())
// result.Bytes()
// 发送Nat地址和接收确认
toClientAddrs := serverConn.LocalAddr + "," + conn.LocalAddr
conn.Write(byte[]{toClientAddrs})
buf := make([]byte, 128)
// 发送Nat地址和接收确认
toClientAddrs := serverConn.LocalAddr + "," + conn.LocalAddr().String()
pack := util.PackageProxy(PROXY_CONTROL_NORMAL, []byte(toClientAddrs))
conn.Write(pack)
buf := make([]byte, 512)
lenAck, err := conn.Read(buf)
flag := 0
if err != nil {
fmt.Println("读客户端确认数据出错")
return
}
if buf[0] == CLIENT_RESP_ACK {
ackPack := util.UnpackageProxy(buf[0:lenAck])
flag := 0
if ackPack.CotnrolID == PROXY_CONTROL_HEARTBITACK {
flag += 1
}
toServerAddrs := conn.LocalAddr + "," + serverConn.LocalAddr
serverSide.Wch <- toServerAddrs
toServerAddrs := conn.LocalAddr().String() + "," + serverConn.LocalAddr
addrPack := util.PackageProxy(PROXY_CONTROL_NORMAL, []byte(toServerAddrs))
serverConn.Wch <- addrPack
// 等待服务端的确认数据
select {
case bufAck := <- serverSide.Rch
if bufAck[0] == SERVER_RESP_ACK {
flag += 1
break
}
case bufAck := <- serverConn.Rch:
pack := util.UnpackageProxy(bufAck)
if pack.CotnrolID == PROXY_CONTROL_HEARTBITACK {
flag += 1
}
break
}
// 收到服务端的确认数据
if flag == 2 {
RWLockServer.Lock()
serverConn.Dch <- true // 关闭服务端连接
delete( OnlineServerList, pairname)
delete( OnlineServerList, pairName)
RWLockServer.Unlock()
}
RWLockClient.Lock()
delete( OnlineClientList, pairname)
delete( OnlineClientList, pairName)
RWLockClient.Unlock()
return
@@ -169,10 +161,9 @@ func Handler(conn net.Conn) {
if r := recover(); r != nil {
fmt.Printf("连接出现问题:%s",r)
}
}()
}()
defer conn.Close()
buf := make([]byte, 1024)
var pairName string
@@ -187,22 +178,23 @@ func Handler(conn net.Conn) {
for {
i, err := conn.Read(buf)
if err != nil {
fmt.Println("读取数据错误:", err.Error())
return
}
//获取匹配名称和连接类型(服务端或客户端)
pairName, roleType, err := parseFirstPackage(buf[0:i])
if err != nil {
return
firstPack := util.UnpackageProxy(buf[0:i])
clientType := firstPack.Data[0]
var pairName string
if len(firstPack.Data) >1 {
pairName = string(firstPack.Data[1:])
}
// 处理客户端连接
if roleType == ROLE_CLIENT{
processRoleClient(conn)
if clientType == ROLE_CLIENT{
processRoleClient(conn, pairName)
return // 退出客户端连接
}
break
@@ -279,7 +271,8 @@ func WHandler(conn net.Conn, C *ServerConn) {
}
// 读客户端数据 + 心跳检测
func RHandler(conn net.Conn, C *ServerSide) {
func RHandler(conn net.Conn, C *ServerConn) {
// 心跳ack
// 业务数据 写入Wch
@@ -290,27 +283,22 @@ func RHandler(conn net.Conn, C *ServerSide) {
if err != nil {
fmt.Println(err)
}
if _, derr := conn.Read(data); derr == nil {
if i, derr := conn.Read(data); derr == nil {
// 可能是来自客户端的消息确认
// 数据消息
fmt.Println(data)
if data[0] == Res {
fmt.Println("recv client data ack")
} else if data[0] == Req {
fmt.Println("recv client data")
fmt.Println(data)
conn.Write([]byte{Res, '#'})
// C.Rch <- data
pack := util.UnpackageProxy(data[0:i])
if pack.CotnrolID == PROXY_CONTROL_HEARTBITACK {
fmt.Println("Received hartbeat ack")//// C.Rch <- data
}
continue
}
//如果等待10秒没有读到客户端数据或读写出错,写心跳包
// 写心跳包
conn.Write([]byte{Req_HEARTBEAT, '#'})
fmt.Println("send ht packet")
heartPack := util.PackageProxy(PROXY_CONTROL_HEARTBIT, []byte(""))
conn.Write(heartPack)
conn.SetReadDeadline(time.Now().Add(5 * time.Second))
if _, herr := conn.Read(data); herr == nil {
@@ -326,7 +314,7 @@ func RHandler(conn net.Conn, C *ServerSide) {
fmt.Println("resv ht packet ack")
} else {
RWLockServer.Lock()
delete(OnlineServerSide, C.Pairname)
delete(OnlineServerList, C.Pairname)
RWLockServer.Unlock()
fmt.Println("delete user!")
return
+33 -48
View File
@@ -4,20 +4,15 @@ import (
"log"
"net"
"os"
"reflect"
"sync"
"unsafe"
"punching/util"
. "punching/constant"
)
var ExitChanMap map[string]chan bool
var RWLock *sync.RWMutex
var DialTargetMap map[string]net.Conn
func b2s(buf []byte) string {
return *(*string)(unsafe.Pointer(&buf))
}
func s2b(s *string) []byte {
return *(*[]byte)(unsafe.Pointer((*reflect.SliceHeader)(unsafe.Pointer(s))))
}
func handleServerConn() {
@@ -25,77 +20,67 @@ func handleServerConn() {
ExitChanMap = make(map[string]chan bool)
RWLock = new(sync.RWMutex)
//var target net.Conn
//var err error
// defer func() {
// if target != nil {
// target.Close()
// }
// //source.Close()
// }()
var targetIP string
targetIP = "192.168.3.2:5901"
var targetAddr = Config.Dial
for {
select {
case r := <-Rch:
case pack := <-Rch:
//确定target是否存在,如果不存在,重新生成target
//分析数据包
log.Println("接收到数据:", len(r.Data))
controlid := r.ControlID
uniqueid := r.Key
pack := r.Data
controlID := pack.ControlID
sessionID := string(pack.SessionID)
data := pack.Data
//log.Println("读取Nat接收包:handleReadConn", string(r[0:34]), "长度为", len(r))
if controlid == "01" {
if controlID == PAIR_CONTROL_QUIT {
RWLock.RLock()
if c, ok := ExitChanMap[uniqueid]; ok {
if c, ok := ExitChanMap[sessionID]; ok {
log.Println("发送退出信号")
c <- true
} else {
log.Println("在ExitChanMap里找不到Key为:", uniqueid)
log.Println("在ExitChanMap里找不到Key为:", sessionID)
}
RWLock.RUnlock()
break
}
//第一次
if controlid == "11" {
log.Println("准备连接:", targetIP)
target, err := net.Dial("tcp", targetIP)
if controlID == PAIR_CONTROL_FIRST {
log.Println("准备连接:", targetAddr)
target, err := net.Dial("tcp", targetAddr)
if err != nil {
log.Println("连接目标出错", targetIP)
log.Println("连接目标出错", targetAddr)
break
}
ExitChanMap[uniqueid] = make(chan bool)
DialTargetMap[uniqueid] = target
ExitChanMap[sessionID] = make(chan bool)
DialTargetMap[sessionID] = target
log.Println("连接目标成功:", targetIP)
log.Println("连接目标成功:", targetAddr)
_, err2 := target.Write(pack)
if err2 != nil {
log.Println("连接成功后写目标出错", err2.Error())
break
}
go ReadFromTarget(target, uniqueid)
go ReadFromTarget(target, sessionID)
} else {
if dialtarget, ok := DialTargetMap[uniqueid]; ok {
if dialtarget, ok := DialTargetMap[sessionID]; ok {
len2, err2 := dialtarget.Write(pack)
len2, err2 := dialtarget.Write(data)
log.Println("已写入:", len2)
if err2 != nil {
log.Println("写目标出错", targetIP, err2.Error())
log.Println("写目标出错", targetAddr, err2.Error())
//发送控制
pack01 := Packet(uniqueid, "01", []byte(""))
Wch <- pack01
quitPack := util.PackageNat(PAIR_CONTROL_QUIT, [4]byte(sessionID),[]byte(""))
Wch <- quitPack
break
}
@@ -114,13 +99,13 @@ func handleServerConn() {
}
// 读取目标流到源
func ReadFromTarget(target net.Conn, uniqueid string) {
func ReadFromTarget(target net.Conn, sessionID string) {
defer func() {
target.Close()
RWLock.Lock()
delete(DialTargetMap, uniqueid)
delete(ExitChanMap, uniqueid)
delete(DialTargetMap, sessionID)
delete(ExitChanMap, sessionID)
RWLock.Unlock()
}()
@@ -134,13 +119,13 @@ func ReadFromTarget(target net.Conn, uniqueid string) {
if err != nil || j == 0 {
log.Println("读取目标连接数据出错,原因为:", err.Error())
pack := Packet(uniqueid, "01", []byte(""))
pack := util.PackageNat(PAIR_CONTROL_QUIT, [4]byte(sessionID),[]byte(""))
Wch <- pack
return
}
pack := Packet(uniqueid, "00", buf[0:j])
pack := util.PackageNat(PAIR_CONTROL_NORMAL,[4]byte(sessionID), buf[0:j])
Wch <- pack
@@ -149,7 +134,7 @@ func ReadFromTarget(target net.Conn, uniqueid string) {
//接受到退出标识
select {
case <-ExitChanMap[uniqueid]:
case <-ExitChanMap[sessionID]:
log.Println("需要退出Accept")
return
}
+67 -2
View File
@@ -7,6 +7,14 @@ import (
"time"
)
var (
Dch chan bool
Rch chan util.PairPackage
Wch chan []byte
)
func Main() {
// 加载配置信息
@@ -22,6 +30,8 @@ func Main() {
}
pairName := Config.Key
var connPeer util.NetConn
// 如果跟Peer连接出错,要重新连接Proxy
for {
@@ -34,7 +44,7 @@ func Main() {
continue
}
localAddr, remoteAddr, errWait := WaitForPeer(conn)
localAddr, remoteAddr, _, errWait := WaitForPeer(conn)
if errWait != nil {
log.Println(errWait)
time.Sleep(5 * time.Second)
@@ -42,7 +52,8 @@ func Main() {
}
//连接对方
connPeer, errPeer := util.DialPeer(localAddr, destAddr)
var errPeer error
connPeer, errPeer = util.DialPeer(localAddr, remoteAddr)
if errPeer != nil { //无法连接上
log.Println(errPeer)
continue
@@ -50,6 +61,60 @@ func Main() {
//已经连接上
// 连接要开启的服务
Dch = make(chan bool)
Rch = make(chan util.PairPackage)
Wch = make(chan []byte)
go RHandler(connPeer) //Nat端写通道
go WHandler(connPeer) //Nat端读通道
// 如果P2P端通讯出错,退出
select {
case <-Dch:
continue
}
}
}
func RHandler(conn util.NetConn) {
//声明一个临时缓冲区,用来存储被截断的数据
tmpBuffer := make([]byte, 0)
buff := make([]byte, 1024)
for {
j, err := conn.Read(buff)
if err != nil {
log.Println("读取连接数据出错,原因为:", err.Error())
Dch <- true
break
}
log.Println("准备解包数据:", j)
// 解包
tmpBuffer = util.UnpackageNat(append(tmpBuffer, buff[:j]...), Rch)
}
}
func WHandler(conn util.NetConn) {
for {
select {
case msg := <-Wch:
l, err := conn.Write(msg)
if err != nil {
log.Println("写到Nat目录连接出错:", err.Error())
Dch <- true
} else {
log.Println(time.Now().UnixNano(), "已写入到Nat", l)
}
// }
}
}
}
+34 -23
View File
@@ -9,13 +9,17 @@ import (
"strings"
)
var (
ProxyDch chan byte
ProxyDch chan util.ProxyPackage
ProxyRch chan []byte
ProxyWch chan []byte
)
func WaitForPeer(conn util.NetConn) (localAddr string, remoteAddr string, err error) {
// 等待,直到P2P客户端连入,在此期间会一直接受Proxy解析端发来的心跳包
func WaitForPeer(conn util.NetConn) (localAddr string, remoteAddr string, pairName string, err error) {
defer conn.Close()
// 接收心跳和客户端连接确认
ProxyRch = make(chan []byte)
@@ -23,31 +27,29 @@ func WaitForPeer(conn util.NetConn) (localAddr string, remoteAddr string, err e
go RProxyHandler(conn)
select {
case ret := <-ProxyDch:
switch ret {
case pack := <-ProxyDch:
switch pack.CotnrolID {
case PROXY_CONTROL_QUIT: //无法跟proxy连接
//关闭连接
err = error("出现错误")
break
case PROXY_CONTROL_NORMAL: // 获取客户端发来信息
localAddr = ""
remoteAddr = ""
pairName = ""
data := pack.Data
str := string(data)
parts := strings.Split(str,",")
localAddr = parts[0]
remoteAddr = parts[1]
pairName = parts[2]
break
}
return
}
}
// ServerDialProxy P2P服务端连接Proxy解析
func ServerDialProxy(proxyAddr string, pairName string) (retConn util.NetConn, err error) {
// 发第一个包
// 地址接收确认
// 心跑包接收和确认
var conn util.NetConn
// 不指定端口,让系统自动分配
@@ -116,21 +118,30 @@ func RProxyHandler(conn util.NetConn) {
for {
// 心跳包,回复ack
data := make([]byte, 128)
data := make([]byte, 512)
i, _ := conn.Read(data)
if i == 0 {
Dch <- true
ProxyDch <- util.ProxyPackage{CotnrolID:PROXY_CONTROL_QUIT}
return
}
if data[0] == Req_HEARTBEAT {
fmt.Println("recv ht pack")
conn.Write([]byte{Res_REGISTER, '#', 'h'})
fmt.Println("send ht pack ack")
} else if data[0] == Req { // 接收到确认信息
fmt.Println("recv data pack")
fmt.Printf("%v\n", string(data[2:]))
conn.Write([]byte{Res, '#'})
// Invalid package
pack := util.UnpackageProxy(data[0:i])
if pack.Head != PROXY_PACKAGE_HEAD {
ProxyDch <- util.ProxyPackage{CotnrolID:PROXY_CONTROL_QUIT}
return
}
if pack.CotnrolID == PROXY_CONTROL_HEARTBIT {
// Received heartbeat package
// 确认
ackPack := util.PackageProxy(PROXY_CONTROL_HEARTBITACK, []byte(""))
conn.Write(ackPack)
}
ProxyDch <- pack
}
}
-1
View File
@@ -1 +0,0 @@
package server
+15 -11
View File
@@ -1,18 +1,16 @@
package util
import (
"flag"
"fmt"
"os"
"github.com/BurntSushi/toml"
"os"
"fmt"
)
// LoadTomlFile 加载配置文件
func LoadTomlFile(filename string)(sections map[string]toml.Primitive, meta MetaData err error){
func LoadTomlFile(fileName string)(sections map[string]toml.Primitive, m toml.MetaData, err error){
// 判断配置文件是否存在
if _, err := os.Stat(fileName); err != nil {
if _, err = os.Stat(fileName); err != nil {
if os.IsNotExist(err) {
err = fmt.Errorf("configuration file %s does not exist.\r\n", fileName)
} else {
@@ -21,21 +19,27 @@ func LoadTomlFile(filename string)(sections map[string]toml.Primitive, meta Meta
return
}
// 加载配置文件
var sections map[string]toml.Primitive
var file toml.Primitive
var meta toml.MetaData
if meta, err = toml.DecodeFile(fileName, &file); err != nil {
err = fmt.Errorf("load configuration file %s failed:%s", fileName, err.Error())
}else{
err = meta.PrimitiveDecode(file, &sections)
}
return
}
m = meta
return
}
// DecodeSection 解码一个节点的配置信息
func DecodeSection(filename, name string, v interface{}) (err error) {
sections, _err := loadTomlFile(filename)
sections, meta, err := LoadTomlFile(filename)
if err != nil{
return
}
+2 -2
View File
@@ -114,13 +114,13 @@ func DialPeer(localAddr string, remoteAddr string) (netconn NetConn, err error)
"remote address", localAddr, remoteAddr)
return
}
err02 := conn.Connect(util.InetAddr(remoteTCPAddr.IP.String()), remoteTCPAddr.Port)
err02 := conn.Connect(InetAddr(remoteTCPAddr.IP.String()), remoteTCPAddr.Port)
if err02 != nil {
log.Printf("第%d次不能连接远程服务器:%s", tryCount, err02.Error())
time.Sleep(1 * time.Second)
continue
} else {
log.Println("已经连接到peer: ", remoteAddr.String())
log.Println("已经连接到peer: ", remoteTCPAddr.String())
break
}
}
+9
View File
@@ -8,10 +8,19 @@ import (
"time"
)
func MAKEWORD(low, high uint8) uint32 {
var ret uint16 = uint16(high)<<8 + uint16(low)
return uint32(ret)
}
type NetConn struct {
sock syscall.Handle
}
func (hole *NetConn) Close() {
syscall.WSACleanup()
+9 -3
View File
@@ -14,6 +14,12 @@ type PairPackage struct {
Data []byte // 数据
}
type ProxyPackage struct {
Head byte // 头
CotnrolID byte // 控制ID
Data []byte // 数据
}
// 跟代理解析端通讯封包
func PackageProxy(control byte, data []byte) []byte {
pack := bytes.NewBuffer(nil)
@@ -24,14 +30,14 @@ func PackageProxy(control byte, data []byte) []byte {
}
// 跟代理解析端拆包
func UnpackageProxy(pact []byte) (control byte, data []byte) {
return pact[1], pact[1:]
func UnpackageProxy(pack []byte) (ProxyPackage) {
return ProxyPackage{pack[0], pack[1], pack[1:]}
}
// Customize P2P data package
// The format of package defined below:
// head(6)+control(1)+session id(4) + data length (4) + data
func PackageNat(control byte, sessionID [4]byte, data []byte) []byte {
func PackageNat(control byte, sessionID string, data []byte) []byte {
pack := bytes.NewBuffer(nil)
pack.Write([]byte(PAIR_PACKAGE_HEAD)) // Head [6]byte
pack.Write([]byte{control})