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" ) type TRPClient struct { svrAddr string linkMap map[int]*conn.Link tunnel *conn.Conn bridgeConnType string stop chan bool sync.Mutex vKey string } //new client func NewRPClient(svraddr string, vKey string, bridgeConnType string) *TRPClient { return &TRPClient{ svrAddr: svraddr, linkMap: make(map[int]*conn.Link), Mutex: sync.Mutex{}, vKey: vKey, bridgeConnType: bridgeConnType, stop: make(chan bool), } } //start func (s *TRPClient) Start() { retry: c, err := NewConn(s.bridgeConnType, s.vKey, s.svrAddr, common.WORK_MAIN) if err != nil { lg.Println("The connection server failed and will be reconnected in five seconds") time.Sleep(time.Second * 5) goto retry } 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) { go s.dealChan() for { flags, err := c.ReadFlag() if err != nil { lg.Printf("Accept server data error %s, end this service", err.Error()) break } switch flags { case common.VERIFY_EER: 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("The authentication key is connected by another client or the server closes the client.") case common.RES_MSG: lg.Println("Server-side return error") break default: lg.Println("The error could not be resolved") break } } 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) if err != nil { c.WriteFail(link.Id) lg.Println("connect to ", link.Host, "error:", err) return } c.WriteSuccess(link.Id) 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() }() <-link.Stop } //隧道模式处理 func (s *TRPClient) dealChan() { var err error 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 } go func() { for { if id, err := s.tunnel.GetLen(); err != nil { break } else { s.Lock() if v, ok := s.linkMap[id]; ok { s.Unlock() if content, err := s.tunnel.GetMsgContent(v); err != nil { lg.Println("get msg content error:", err, id) pool.PutBufPoolCopy(content) break } else { if len(content) == len(common.IO_EOF) && string(content) == common.IO_EOF { v.Conn.Close() } else if v.Conn != nil { v.Conn.Write(content) } pool.PutBufPoolCopy(content) } } else { s.Unlock() } } } }() <-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 }