客户端配置，端口白名单等

client/client.go

@@ -1,12 +1,16 @@
 package client
 
 import (
+	"errors"
 	"github.com/cnlh/nps/lib/common"
+	"github.com/cnlh/nps/lib/config"
 	"github.com/cnlh/nps/lib/conn"
 	"github.com/cnlh/nps/lib/kcp"
 	"github.com/cnlh/nps/lib/lg"
 	"github.com/cnlh/nps/lib/pool"
+	"io/ioutil"
 	"net"
+	"path/filepath"
 	"sync"
 	"time"
 )
@@ -14,9 +18,9 @@ import (
 type TRPClient struct {
 	svrAddr        string
 	linkMap        map[int]*conn.Link
-	stop           chan bool
 	tunnel         *conn.Conn
 	bridgeConnType string
+	stop           chan bool
 	sync.Mutex
 	vKey string
 }
@@ -26,81 +30,70 @@ func NewRPClient(svraddr string, vKey string, bridgeConnType string) *TRPClient
 	return &TRPClient{
 		svrAddr:        svraddr,
 		linkMap:        make(map[int]*conn.Link),
-		stop:           make(chan bool),
 		Mutex:          sync.Mutex{},
 		vKey:           vKey,
 		bridgeConnType: bridgeConnType,
+		stop:           make(chan bool),
 	}
 }
 
 //start
-func (s *TRPClient) Start() error {
-	s.NewConn()
-	return nil
-}
-
-//新建
-func (s *TRPClient) NewConn() {
-	var err error
-	var c net.Conn
+func (s *TRPClient) Start() {
 retry:
-	if s.bridgeConnType == "tcp" {
-		c, err = net.Dial("tcp", s.svrAddr)
-	} else {
-		var sess *kcp.UDPSession
-		sess, err = kcp.DialWithOptions(s.svrAddr, nil, 150, 3)
-		conn.SetUdpSession(sess)
-		c = sess
-	}
+	c, err := NewConn(s.bridgeConnType, s.vKey, s.svrAddr, common.WORK_MAIN)
 	if err != nil {
-		lg.Println("连接服务端失败,五秒后将重连")
+		lg.Println("The connection server failed and will be reconnected in five seconds")
 		time.Sleep(time.Second * 5)
 		goto retry
-		return
 	}
-	s.processor(conn.NewConn(c))
+	lg.Printf("Successful connection with server %s", s.svrAddr)
+	s.processor(c)
+}
+
+func (s *TRPClient) Close() {
+	s.tunnel.Close()
+	s.stop <- true
+	for _, v := range s.linkMap {
+		v.Stop <- true
+	}
 }
 
 //处理
 func (s *TRPClient) processor(c *conn.Conn) {
-	c.SetAlive(s.bridgeConnType)
-	if _, err := c.Write([]byte(common.Getverifyval(s.vKey))); err != nil {
-		return
-	}
-	c.WriteMain()
 	go s.dealChan()
 	for {
 		flags, err := c.ReadFlag()
 		if err != nil {
-			lg.Println("服务端断开,正在重新连接")
+			lg.Printf("Accept server data error %s, end this service", err.Error())
 			break
 		}
 		switch flags {
 		case common.VERIFY_EER:
-			lg.Fatalf("vKey:%s不正确,服务端拒绝连接,请检查", s.vKey)
+			lg.Fatalf("VKey:%s is incorrect, the server refuses to connect, please check", s.vKey)
 		case common.NEW_CONN:
 			if link, err := c.GetLinkInfo(); err != nil {
 				break
 			} else {
+				link.Stop = make(chan bool)
 				s.Lock()
 				s.linkMap[link.Id] = link
 				s.Unlock()
 				go s.linkProcess(link, c)
 			}
 		case common.RES_CLOSE:
-			lg.Fatalln("该vkey被另一客户连接")
+			lg.Fatalln("The authentication key is connected by another client or the server closes the client.")
 		case common.RES_MSG:
-			lg.Println("服务端返回错误，重新连接")
+			lg.Println("Server-side return error")
 			break
 		default:
-			lg.Println("无法解析该错误，重新连接")
+			lg.Println("The error could not be resolved")
 			break
 		}
 	}
-	s.stop <- true
-	s.linkMap = make(map[int]*conn.Link)
-	go s.NewConn()
+	c.Close()
+	s.Close()
 }
+
 func (s *TRPClient) linkProcess(link *conn.Link, c *conn.Conn) {
 	//与目标建立连接
 	server, err := net.DialTimeout(link.ConnType, link.Host, time.Second*3)
@@ -113,56 +106,40 @@ func (s *TRPClient) linkProcess(link *conn.Link, c *conn.Conn) {
 
 	c.WriteSuccess(link.Id)
 
-	link.Conn = conn.NewConn(server)
-	buf := pool.BufPoolCopy.Get().([]byte)
-	for {
-		if n, err := server.Read(buf); err != nil {
-			s.tunnel.SendMsg([]byte(common.IO_EOF), link)
-			break
-		} else {
-			if _, err := s.tunnel.SendMsg(buf[:n], link); err != nil {
-				c.Close()
+	go func() {
+		link.Conn = conn.NewConn(server)
+		buf := pool.BufPoolCopy.Get().([]byte)
+		for {
+			if n, err := server.Read(buf); err != nil {
+				s.tunnel.SendMsg([]byte(common.IO_EOF), link)
 				break
+			} else {
+				if _, err := s.tunnel.SendMsg(buf[:n], link); err != nil {
+					c.Close()
+					break
+				}
 			}
 		}
-	}
-	pool.PutBufPoolCopy(buf)
-	s.Lock()
-	delete(s.linkMap, link.Id)
-	s.Unlock()
+		pool.PutBufPoolCopy(buf)
+		s.Lock()
+		delete(s.linkMap, link.Id)
+		s.Unlock()
+	}()
+	<-link.Stop
 }
 
 //隧道模式处理
 func (s *TRPClient) dealChan() {
 	var err error
-	var c net.Conn
-	var sess *kcp.UDPSession
-	if s.bridgeConnType == "tcp" {
-		c, err = net.Dial("tcp", s.svrAddr)
-	} else {
-		sess, err = kcp.DialWithOptions(s.svrAddr, nil, 10, 3)
-		conn.SetUdpSession(sess)
-		c = sess
-	}
+	s.tunnel, err = NewConn(s.bridgeConnType, s.vKey, s.svrAddr, common.WORK_CHAN)
 	if err != nil {
 		lg.Println("connect to ", s.svrAddr, "error:", err)
 		return
 	}
-	//验证
-	if _, err := c.Write([]byte(common.Getverifyval(s.vKey))); err != nil {
-		lg.Println("connect to ", s.svrAddr, "error:", err)
-		return
-	}
-	//默认长连接保持
-	s.tunnel = conn.NewConn(c)
-	s.tunnel.SetAlive(s.bridgeConnType)
-	//写标志
-	s.tunnel.WriteChan()
 
 	go func() {
 		for {
 			if id, err := s.tunnel.GetLen(); err != nil {
-				lg.Println("get msg id error")
 				break
 			} else {
 				s.Lock()
@@ -186,8 +163,103 @@ func (s *TRPClient) dealChan() {
 			}
 		}
 	}()
-	select {
-	case <-s.stop:
-		break
-	}
+	<-s.stop
+}
+
+var errAdd = errors.New("The server returned an error, which port or host may have been occupied or not allowed to open.")
+
+func StartFromFile(path string) {
+	first := true
+	cnf, err := config.NewConfig(path)
+	if err != nil {
+		lg.Fatalln(err)
+	}
+	lg.Printf("Loading configuration file %s successfully", path)
+re:
+	if first || cnf.CommonConfig.AutoReconnection {
+		if !first {
+			lg.Println("Reconnecting...")
+			time.Sleep(time.Second * 5)
+		}
+	} else {
+		return
+	}
+	first = false
+	c, err := NewConn(cnf.CommonConfig.Tp, cnf.CommonConfig.VKey, cnf.CommonConfig.Server, common.WORK_CONFIG)
+	if err != nil {
+		lg.Println(err)
+		goto re
+	}
+	if _, err := c.SendConfigInfo(cnf.CommonConfig.Cnf); err != nil {
+		lg.Println(err)
+		goto re
+	}
+	var b []byte
+	if b, err = c.ReadLen(16); err != nil {
+		lg.Println(err)
+		goto re
+	} else {
+		ioutil.WriteFile(filepath.Join(common.GetTmpPath(), "npc_vkey.txt"), []byte(string(b)), 0600)
+	}
+	if !c.GetAddStatus() {
+		lg.Println(errAdd)
+		goto re
+	}
+	for _, v := range cnf.Hosts {
+		if _, err := c.SendHostInfo(v); err != nil {
+			lg.Println(err)
+			goto re
+		}
+		if !c.GetAddStatus() {
+			lg.Println(errAdd, v.Host)
+			goto re
+		}
+	}
+	for _, v := range cnf.Tasks {
+		if _, err := c.SendTaskInfo(v); err != nil {
+			lg.Println(err)
+			goto re
+		}
+		if !c.GetAddStatus() {
+			lg.Println(errAdd, v.Port)
+			goto re
+		}
+	}
+
+	c.Close()
+
+	NewRPClient(cnf.CommonConfig.Server, string(b), cnf.CommonConfig.Tp).Start()
+	goto re
+}
+
+//Create a new connection with the server and verify it
+func NewConn(tp string, vkey string, server string, connType string) (*conn.Conn, error) {
+	var err error
+	var connection net.Conn
+	var sess *kcp.UDPSession
+	if tp == "tcp" {
+		connection, err = net.Dial("tcp", server)
+	} else {
+		sess, err = kcp.DialWithOptions(server, nil, 10, 3)
+		conn.SetUdpSession(sess)
+		connection = sess
+	}
+	if err != nil {
+		return nil, err
+	}
+	c := conn.NewConn(connection)
+	if _, err := c.Write([]byte(common.Getverifyval(vkey))); err != nil {
+		lg.Println(err)
+	}
+	if s, err := c.ReadFlag(); err != nil {
+		lg.Println(err)
+	} else if s == common.VERIFY_EER {
+		lg.Fatalf("Validation key %s incorrect", vkey)
+	}
+	if _, err := c.Write([]byte(connType)); err != nil {
+		lg.Println(err)
+	}
+	c.SetAlive(tp)
+
+	return c, nil
 }

client/client_test.go

@@ -0,0 +1,74 @@
+package client
+
+import (
+	"github.com/cnlh/nps/lib/common"
+	conn2 "github.com/cnlh/nps/lib/conn"
+	"github.com/cnlh/nps/lib/file"
+	"net"
+	"sync"
+	"testing"
+)
+
+func TestConfig(t *testing.T) {
+	conn, err := net.Dial("tcp", "127.0.0.1:8284")
+	if err != nil {
+		t.Fail()
+	}
+	c := conn2.NewConn(conn)
+	c.SetAlive("tcp")
+	if _, err := c.Write([]byte(common.Getverifyval("123"))); err != nil {
+		t.Fail()
+	}
+	c.WriteConfig()
+	config := &file.Config{
+		U:              "1",
+		P:              "2",
+		Compress:       "snappy",
+		Crypt:          true,
+		CompressEncode: 0,
+		CompressDecode: 0,
+	}
+	host := &file.Host{
+		Host:         "a.o.com",
+		Target:       "127.0.0.1:8080",
+		HeaderChange: "",
+		HostChange:   "",
+		Flow:         nil,
+		Client:       nil,
+		Remark:       "111",
+		NowIndex:     0,
+		TargetArr:    nil,
+		NoStore:      false,
+		RWMutex:      sync.RWMutex{},
+	}
+	tunnel := &file.Tunnel{
+		Port: 9001,
+		Mode:    "tunnelServer",
+		Target:  "127.0.0.1:8082",
+		Remark:  "333",
+	}
+	var b []byte
+	if b, err = c.ReadLen(16); err != nil {
+		t.Fail()
+	}
+	if _, err := c.SendConfigInfo(config); err != nil {
+		t.Fail()
+	}
+	if !c.GetAddStatus() {
+		t.Fail()
+	}
+	if _, err := c.SendHostInfo(host); err != nil {
+		t.Fail()
+	}
+	if !c.GetAddStatus() {
+		t.Fail()
+	}
+	if _, err := c.SendTaskInfo(tunnel); err != nil {
+		t.Fail()
+	}
+	if !c.GetAddStatus() {
+		t.Fail()
+	}
+	c.Close()
+	NewRPClient("127.0.0.1:8284", string(b), "tcp").Start()
+}