add functions

server/proxy/base.go

@@ -17,8 +17,8 @@ type Service interface {
 	Close() error
 }
 
-//server base struct
-type server struct {
+//Server BaseServer struct
+type BaseServer struct {
 	id           int
 	bridge       *bridge.Bridge
 	task         *file.Tunnel
@@ -26,21 +26,30 @@ type server struct {
 	sync.Mutex
 }
 
-func (s *server) FlowAdd(in, out int64) {
+func NewBaseServer(bridge *bridge.Bridge, task *file.Tunnel) *BaseServer {
+	return &BaseServer{
+		bridge:       bridge,
+		task:         task,
+		errorContent: nil,
+		Mutex:        sync.Mutex{},
+	}
+}
+
+func (s *BaseServer) FlowAdd(in, out int64) {
 	s.Lock()
 	defer s.Unlock()
 	s.task.Flow.ExportFlow += out
 	s.task.Flow.InletFlow += in
 }
 
-func (s *server) FlowAddHost(host *file.Host, in, out int64) {
+func (s *BaseServer) FlowAddHost(host *file.Host, in, out int64) {
 	s.Lock()
 	defer s.Unlock()
 	host.Flow.ExportFlow += out
 	host.Flow.InletFlow += in
 }
 
-func (s *server) linkCopy(link *conn.Link, c *conn.Conn, rb []byte, tunnel *conn.Conn, flow *file.Flow) {
+func (s *BaseServer) linkCopy(link *conn.Link, c *conn.Conn, rb []byte, tunnel *conn.Conn, flow *file.Flow) {
 	if rb != nil {
 		if _, err := tunnel.SendMsg(rb, link); err != nil {
 			c.Close()
@@ -68,16 +77,17 @@ func (s *server) linkCopy(link *conn.Link, c *conn.Conn, rb []byte, tunnel *conn
 		}
 		<-link.StatusCh
 	}
+	s.task.Client.AddConn()
 	pool.PutBufPoolCopy(buf)
 }
 
-func (s *server) writeConnFail(c net.Conn) {
+func (s *BaseServer) writeConnFail(c net.Conn) {
 	c.Write([]byte(common.ConnectionFailBytes))
 	c.Write(s.errorContent)
 }
 
 //权限认证
-func (s *server) auth(r *http.Request, c *conn.Conn, u, p string) error {
+func (s *BaseServer) auth(r *http.Request, c *conn.Conn, u, p string) error {
 	if u != "" && p != "" && !common.CheckAuth(r, u, p) {
 		c.Write([]byte(common.UnauthorizedBytes))
 		c.Close()
@@ -86,9 +96,23 @@ func (s *server) auth(r *http.Request, c *conn.Conn, u, p string) error {
 	return nil
 }
 
-func (s *server) checkFlow() error {
+func (s *BaseServer) checkFlow() error {
 	if s.task.Client.Flow.FlowLimit > 0 && (s.task.Client.Flow.FlowLimit<<20) < (s.task.Client.Flow.ExportFlow+s.task.Client.Flow.InletFlow) {
 		return errors.New("Traffic exceeded")
 	}
 	return nil
 }
+
+//与客户端建立通道
+func (s *BaseServer) DealClient(c *conn.Conn, addr string, rb []byte) error {
+	link := conn.NewLink(s.task.Client.GetId(), common.CONN_TCP, addr, s.task.Client.Cnf.CompressEncode, s.task.Client.Cnf.CompressDecode, s.task.Client.Cnf.Crypt, c, s.task.Flow, nil, s.task.Client.Rate, nil)
+
+	if tunnel, err := s.bridge.SendLinkInfo(s.task.Client.Id, link, c.Conn.RemoteAddr().String()); err != nil {
+		c.Close()
+		return err
+	} else {
+		link.Run(true)
+		s.linkCopy(link, c, rb, tunnel, s.task.Flow)
+	}
+	return nil
+}

server/proxy/http.go

@@ -7,17 +7,18 @@ import (
 	"github.com/cnlh/nps/lib/common"
 	"github.com/cnlh/nps/lib/conn"
 	"github.com/cnlh/nps/lib/file"
-	"github.com/cnlh/nps/lib/lg"
 	"github.com/cnlh/nps/vender/github.com/astaxie/beego"
+	"github.com/cnlh/nps/vender/github.com/astaxie/beego/logs"
 	"net/http"
 	"net/http/httputil"
+	"os"
 	"path/filepath"
 	"strconv"
 	"sync"
 )
 
 type httpServer struct {
-	server
+	BaseServer
 	httpPort  int //http端口
 	httpsPort int //https监听端口
 	pemPath   string
@@ -31,7 +32,7 @@ func NewHttp(bridge *bridge.Bridge, c *file.Tunnel) *httpServer {
 	pemPath := beego.AppConfig.String("pemPath")
 	keyPath := beego.AppConfig.String("keyPath")
 	return &httpServer{
-		server: server{
+		BaseServer: BaseServer{
 			task:   c,
 			bridge: bridge,
 			Mutex:  sync.Mutex{},
@@ -54,26 +55,30 @@ func (s *httpServer) Start() error {
 	if s.httpPort > 0 {
 		http = s.NewServer(s.httpPort)
 		go func() {
-			lg.Println("Start http listener, port is", s.httpPort)
+			logs.Info("Start http listener, port is", s.httpPort)
 			err := http.ListenAndServe()
 			if err != nil {
-				lg.Fatalln(err)
+				logs.Error(err)
+				os.Exit(0)
 			}
 		}()
 	}
 	if s.httpsPort > 0 {
 		if !common.FileExists(s.pemPath) {
-			lg.Fatalf("ssl certFile %s is not exist", s.pemPath)
+			logs.Error("ssl certFile %s is not exist", s.pemPath)
+			os.Exit(0)
 		}
 		if !common.FileExists(s.keyPath) {
-			lg.Fatalf("ssl keyFile %s exist", s.keyPath)
+			logs.Error("ssl keyFile %s exist", s.keyPath)
+			os.Exit(0)
 		}
 		https = s.NewServer(s.httpsPort)
 		go func() {
-			lg.Println("Start https listener, port is", s.httpsPort)
+			logs.Info("Start https listener, port is", s.httpsPort)
 			err := https.ListenAndServeTLS(s.pemPath, s.keyPath)
 			if err != nil {
-				lg.Fatalln(err)
+				logs.Error(err)
+				os.Exit(0)
 			}
 		}()
 	}
@@ -118,9 +123,14 @@ func (s *httpServer) process(c *conn.Conn, r *http.Request) {
 		err      error
 	)
 	if host, err = file.GetCsvDb().GetInfoByHost(r.Host, r); err != nil {
-		lg.Printf("the url %s %s can't be parsed!", r.Host, r.RequestURI)
+		logs.Notice("the url %s %s can't be parsed!", r.Host, r.RequestURI)
 		goto end
+	} else if !host.Client.GetConn() {
+		logs.Notice("Connections exceed the current client %d limit", host.Client.Id)
+		c.Close()
+		return
 	} else {
+		logs.Trace("New http(s) connection,clientId %d,host %s,url %s,remote address %s", host.Client.Id, r.Host, r.URL, r.RemoteAddr)
 		lastHost = host
 	}
 	for {
@@ -137,22 +147,24 @@ func (s *httpServer) process(c *conn.Conn, r *http.Request) {
 			}
 			lk = conn.NewLink(host.Client.GetId(), common.CONN_TCP, host.GetRandomTarget(), host.Client.Cnf.CompressEncode, host.Client.Cnf.CompressDecode, host.Client.Cnf.Crypt, c, host.Flow, nil, host.Client.Rate, nil)
 			if tunnel, err = s.bridge.SendLinkInfo(host.Client.Id, lk, c.Conn.RemoteAddr().String()); err != nil {
-				lg.Println(err)
+				logs.Notice(err)
 				break
 			}
 			lk.Run(true)
 			isConn = false
 		} else {
 			r, err = http.ReadRequest(bufio.NewReader(c))
+			logs.Trace("New http(s) connection,clientId %d,host %s,url %s,remote address %s", host.Client.Id, r.Host, r.URL, r.RemoteAddr)
 			if err != nil {
 				break
 			}
 			if host, err = file.GetCsvDb().GetInfoByHost(r.Host, r); err != nil {
-				lg.Printf("the url %s %s is not found !", r.Host, r.RequestURI)
+				logs.Notice("the url %s %s can't be parsed!", r.Host, r.RequestURI)
 				break
 			} else if host != lastHost {
 				lastHost = host
 				isConn = true
+				host.Client.AddConn()
 				goto start
 			}
 		}
@@ -176,6 +188,9 @@ end:
 		tunnel.SendMsg([]byte(common.IO_EOF), lk)
 	}
 	c.Close()
+	if host != nil {
+		host.Client.AddConn()
+	}
 }
 
 func (s *httpServer) NewServer(port int) *http.Server {

server/proxy/socks5.go

@@ -7,7 +7,7 @@ import (
 	"github.com/cnlh/nps/lib/common"
 	"github.com/cnlh/nps/lib/conn"
 	"github.com/cnlh/nps/lib/file"
-	"github.com/cnlh/nps/lib/lg"
+	"github.com/cnlh/nps/vender/github.com/astaxie/beego/logs"
 	"io"
 	"net"
 	"strconv"
@@ -48,7 +48,7 @@ const (
 )
 
 type Sock5ModeServer struct {
-	server
+	BaseServer
 	listener net.Listener
 }
 
@@ -67,7 +67,7 @@ func (s *Sock5ModeServer) handleRequest(c net.Conn) {
 	_, err := io.ReadFull(c, header)
 
 	if err != nil {
-		lg.Println("illegal request", err)
+		logs.Warn("illegal request", err)
 		c.Close()
 		return
 	}
@@ -165,7 +165,6 @@ func (s *Sock5ModeServer) handleBind(c net.Conn) {
 
 //udp
 func (s *Sock5ModeServer) handleUDP(c net.Conn) {
-	lg.Println("UDP Associate")
 	/*
 	   +----+------+------+----------+----------+----------+
 	   |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
@@ -178,7 +177,7 @@ func (s *Sock5ModeServer) handleUDP(c net.Conn) {
 	// relay udp datagram silently, without any notification to the requesting client
 	if buf[2] != 0 {
 		// does not support fragmentation, drop it
-		lg.Println("does not support fragmentation, drop")
+		logs.Warn("does not support fragmentation, drop")
 		dummy := make([]byte, maxUDPPacketSize)
 		c.Read(dummy)
 	}
@@ -190,13 +189,13 @@ func (s *Sock5ModeServer) handleUDP(c net.Conn) {
 func (s *Sock5ModeServer) handleConn(c net.Conn) {
 	buf := make([]byte, 2)
 	if _, err := io.ReadFull(c, buf); err != nil {
-		lg.Println("negotiation err", err)
+		logs.Warn("negotiation err", err)
 		c.Close()
 		return
 	}
 
 	if version := buf[0]; version != 5 {
-		lg.Println("only support socks5, request from: ", c.RemoteAddr())
+		logs.Warn("only support socks5, request from: ", c.RemoteAddr())
 		c.Close()
 		return
 	}
@@ -204,7 +203,7 @@ func (s *Sock5ModeServer) handleConn(c net.Conn) {
 
 	methods := make([]byte, nMethods)
 	if len, err := c.Read(methods); len != int(nMethods) || err != nil {
-		lg.Println("wrong method")
+		logs.Warn("wrong method")
 		c.Close()
 		return
 	}
@@ -213,7 +212,7 @@ func (s *Sock5ModeServer) handleConn(c net.Conn) {
 		c.Write(buf)
 		if err := s.Auth(c); err != nil {
 			c.Close()
-			lg.Println("Validation failed:", err)
+			logs.Warn("Validation failed:", err)
 			return
 		}
 	} else {
@@ -271,9 +270,15 @@ func (s *Sock5ModeServer) Start() error {
 			if strings.Contains(err.Error(), "use of closed network connection") {
 				break
 			}
-			lg.Fatalln("accept error: ", err)
+			logs.Warn("accept error: ", err)
+		}
+		if s.task.Client.GetConn() {
+			logs.Trace("New socks5 connection,client %d,remote address %s", s.task.Client.Id, conn.RemoteAddr())
+			go s.handleConn(conn)
+		} else {
+			logs.Warn("Connections exceed the current client %d limit", s.task.Client.Id)
+			conn.Close()
 		}
-		go s.handleConn(conn)
 	}
 	return nil
 }

server/proxy/tcp.go

@@ -6,15 +6,16 @@ import (
 	"github.com/cnlh/nps/lib/common"
 	"github.com/cnlh/nps/lib/conn"
 	"github.com/cnlh/nps/lib/file"
-	"github.com/cnlh/nps/lib/lg"
 	"github.com/cnlh/nps/vender/github.com/astaxie/beego"
+	"github.com/cnlh/nps/vender/github.com/astaxie/beego/logs"
 	"net"
+	"os"
 	"path/filepath"
 	"strings"
 )
 
 type TunnelModeServer struct {
-	server
+	BaseServer
 	process  process
 	listener *net.TCPListener
 }
@@ -41,24 +42,16 @@ func (s *TunnelModeServer) Start() error {
 			if strings.Contains(err.Error(), "use of closed network connection") {
 				break
 			}
-			lg.Println(err)
+			logs.Info(err)
 			continue
 		}
-		go s.process(conn.NewConn(c), s)
-	}
-	return nil
-}
-
-//与客户端建立通道
-func (s *TunnelModeServer) dealClient(c *conn.Conn, cnf *file.Config, addr string, method string, rb []byte) error {
-	link := conn.NewLink(s.task.Client.GetId(), common.CONN_TCP, addr, cnf.CompressEncode, cnf.CompressDecode, cnf.Crypt, c, s.task.Flow, nil, s.task.Client.Rate, nil)
-
-	if tunnel, err := s.bridge.SendLinkInfo(s.task.Client.Id, link, c.Conn.RemoteAddr().String()); err != nil {
-		c.Close()
-		return err
-	} else {
-		link.Run(true)
-		s.linkCopy(link, c, rb, tunnel, s.task.Flow)
+		if s.task.Client.GetConn() {
+			logs.Trace("New tcp connection,client %d,remote address %s", s.task.Client.Id, c.RemoteAddr())
+			go s.process(conn.NewConn(c), s)
+		} else {
+			logs.Info("Connections exceed the current client %d limit", s.task.Client.Id)
+			c.Close()
+		}
 	}
 	return nil
 }
@@ -70,17 +63,18 @@ func (s *TunnelModeServer) Close() error {
 
 //web管理方式
 type WebServer struct {
-	server
+	BaseServer
 }
 
 //开始
 func (s *WebServer) Start() error {
 	p, _ := beego.AppConfig.Int("httpport")
 	if !common.TestTcpPort(p) {
-		lg.Fatalf("Web management port %d is occupied", p)
+		logs.Error("Web management port %d is occupied", p)
+		os.Exit(0)
 	}
 	beego.BConfig.WebConfig.Session.SessionOn = true
-	lg.Println("Web management start, access port is", p)
+	logs.Info("Web management start, access port is", p)
 	beego.SetStaticPath("/static", filepath.Join(common.GetRunPath(), "web", "static"))
 	beego.SetViewsPath(filepath.Join(common.GetRunPath(), "web", "views"))
 	beego.Run()
@@ -102,23 +96,23 @@ type process func(c *conn.Conn, s *TunnelModeServer) error
 
 //tcp隧道模式
 func ProcessTunnel(c *conn.Conn, s *TunnelModeServer) error {
-	return s.dealClient(c, s.task.Client.Cnf, s.task.Target, "", nil)
+	return s.DealClient(c, s.task.Target, nil)
 }
 
 //http代理模式
 func ProcessHttp(c *conn.Conn, s *TunnelModeServer) error {
-	method, addr, rb, err, r := c.GetHost()
+	_, addr, rb, err, r := c.GetHost()
 	if err != nil {
 		c.Close()
-		lg.Println(err)
+		logs.Info(err)
 		return err
 	}
 	if r.Method == "CONNECT" {
 		c.Write([]byte("HTTP/1.1 200 Connection Established\r\n"))
-		rb = nil //reset
+		rb = nil
 	}
 	if err := s.auth(r, c, s.task.Client.Cnf.U, s.task.Client.Cnf.P); err != nil {
 		return err
 	}
-	return s.dealClient(c, s.task.Client.Cnf, addr, method, rb)
+	return s.DealClient(c, addr, rb)
 }

server/proxy/udp.go

@@ -6,12 +6,13 @@ import (
 	"github.com/cnlh/nps/lib/conn"
 	"github.com/cnlh/nps/lib/file"
 	"github.com/cnlh/nps/lib/pool"
+	"github.com/cnlh/nps/vender/github.com/astaxie/beego/logs"
 	"net"
 	"strings"
 )
 
 type UdpModeServer struct {
-	server
+	BaseServer
 	listener *net.UDPConn
 	udpMap   map[string]*conn.Conn
 }
@@ -40,6 +41,7 @@ func (s *UdpModeServer) Start() error {
 			}
 			continue
 		}
+		logs.Trace("New ydo connection,client %d,remote address %s", s.task.Client.Id, addr)
 		go s.process(addr, buf[:n])
 	}
 	return nil