package conn import ( "bufio" "bytes" "encoding/binary" "errors" "github.com/cnlh/nps/lib/common" "github.com/cnlh/nps/lib/config" "github.com/cnlh/nps/lib/crypt" "github.com/cnlh/nps/lib/file" "github.com/cnlh/nps/lib/mux" "github.com/cnlh/nps/lib/pool" "github.com/cnlh/nps/lib/rate" "github.com/cnlh/nps/vender/github.com/xtaci/kcp" "io" "net" "net/http" "net/url" "strconv" "strings" "sync" "time" ) type Conn struct { Conn net.Conn } //new conn func NewConn(conn net.Conn) *Conn { return &Conn{Conn: conn} } //get host 、connection type、method...from connection func (s *Conn) GetHost() (method, address string, rb []byte, err error, r *http.Request) { var b [32 * 1024]byte var n int if n, err = s.Read(b[:]); err != nil { return } rb = b[:n] r, err = http.ReadRequest(bufio.NewReader(bytes.NewReader(rb))) if err != nil { return } hostPortURL, err := url.Parse(r.Host) if err != nil { address = r.Host err = nil return } if hostPortURL.Opaque == "443" { if strings.Index(r.Host, ":") == -1 { address = r.Host + ":443" } else { address = r.Host } } else { if strings.Index(r.Host, ":") == -1 { address = r.Host + ":80" } else { address = r.Host } } return } func (s *Conn) GetShortLenContent() (b []byte, err error) { var l int if l, err = s.GetLen(); err != nil { return } if l < 0 || l > 32<<10 { err = errors.New("read length error") return } return s.GetShortContent(l) } func (s *Conn) GetShortContent(l int) (b []byte, err error) { buf := make([]byte, l) return buf, binary.Read(s, binary.LittleEndian, &buf) } //读取指定长度内容 func (s *Conn) ReadLen(cLen int, buf []byte) (int, error) { if cLen > len(buf) { return 0, errors.New("长度错误" + strconv.Itoa(cLen)) } if n, err := io.ReadFull(s, buf[:cLen]); err != nil || n != cLen { return n, errors.New("Error reading specified length " + err.Error()) } return cLen, nil } func (s *Conn) GetLen() (int, error) { var l int32 err := binary.Read(s, binary.LittleEndian, &l) return int(l), err } func (s *Conn) WriteLenContent(buf []byte) (err error) { var b []byte if b, err = GetLenBytes(buf); err != nil { return } return binary.Write(s.Conn, binary.LittleEndian, b) } //read flag func (s *Conn) ReadFlag() (string, error) { buf := make([]byte, 4) return string(buf), binary.Read(s, binary.LittleEndian, &buf) } //set alive func (s *Conn) SetAlive(tp string) { switch s.Conn.(type) { case *kcp.UDPSession: s.Conn.(*kcp.UDPSession).SetReadDeadline(time.Time{}) case *net.TCPConn: conn := s.Conn.(*net.TCPConn) conn.SetReadDeadline(time.Time{}) conn.SetKeepAlive(true) conn.SetKeepAlivePeriod(time.Duration(2 * time.Second)) case *mux.PortConn: s.Conn.(*mux.PortConn).SetReadDeadline(time.Time{}) } } //set read deadline func (s *Conn) SetReadDeadline(t time.Duration, tp string) { switch s.Conn.(type) { case *kcp.UDPSession: s.Conn.(*kcp.UDPSession).SetReadDeadline(time.Now().Add(time.Duration(t) * time.Second)) case *net.TCPConn: s.Conn.(*net.TCPConn).SetReadDeadline(time.Now().Add(time.Duration(t) * time.Second)) case *mux.PortConn: s.Conn.(*mux.PortConn).SetReadDeadline(time.Now().Add(time.Duration(t) * time.Second)) } } //send info for link func (s *Conn) SendLinkInfo(link *Link) (int, error) { raw := bytes.NewBuffer([]byte{}) common.BinaryWrite(raw, link.ConnType, link.Host, common.GetStrByBool(link.Compress), common.GetStrByBool(link.Crypt), link.RemoteAddr) return s.Write(raw.Bytes()) } //get link info from conn func (s *Conn) GetLinkInfo() (lk *Link, err error) { lk = new(Link) var l int buf := pool.BufPoolMax.Get().([]byte) defer pool.PutBufPoolMax(buf) if l, err = s.GetLen(); err != nil { return } else if _, err = s.ReadLen(l, buf); err != nil { return } else { arr := strings.Split(string(buf[:l]), common.CONN_DATA_SEQ) lk.ConnType = arr[0] lk.Host = arr[1] lk.Compress = common.GetBoolByStr(arr[2]) lk.Crypt = common.GetBoolByStr(arr[3]) lk.RemoteAddr = arr[4] } return } //send info for link func (s *Conn) SendHealthInfo(info, status string) (int, error) { raw := bytes.NewBuffer([]byte{}) common.BinaryWrite(raw, info, status) return s.Write(raw.Bytes()) } //get health info from conn func (s *Conn) GetHealthInfo() (info string, status bool, err error) { var l int buf := pool.BufPoolMax.Get().([]byte) defer pool.PutBufPoolMax(buf) if l, err = s.GetLen(); err != nil { return } else if _, err = s.ReadLen(l, buf); err != nil { return } else { arr := strings.Split(string(buf[:l]), common.CONN_DATA_SEQ) if len(arr) >= 2 { return arr[0], common.GetBoolByStr(arr[1]), nil } } return "", false, errors.New("receive health info error") } //send host info func (s *Conn) SendHostInfo(h *file.Host) (int, error) { /* The task info is formed as follows: +----+-----+---------+ |type| len | content | +----+---------------+ | 4 | 4 | ... | +----+---------------+ */ raw := bytes.NewBuffer([]byte{}) binary.Write(raw, binary.LittleEndian, []byte(common.NEW_HOST)) common.BinaryWrite(raw, h.Host, h.Target.TargetStr, h.HeaderChange, h.HostChange, h.Remark, h.Location, h.Scheme) return s.Write(raw.Bytes()) } //get task or host result of add func (s *Conn) GetAddStatus() (b bool) { binary.Read(s.Conn, binary.LittleEndian, &b) return } func (s *Conn) WriteAddOk() error { return binary.Write(s.Conn, binary.LittleEndian, true) } func (s *Conn) WriteAddFail() error { defer s.Close() return binary.Write(s.Conn, binary.LittleEndian, false) } //get task info func (s *Conn) GetHostInfo() (h *file.Host, err error) { var l int buf := pool.BufPoolMax.Get().([]byte) defer pool.PutBufPoolMax(buf) if l, err = s.GetLen(); err != nil { return } else if _, err = s.ReadLen(l, buf); err != nil { return } else { arr := strings.Split(string(buf[:l]), common.CONN_DATA_SEQ) h = new(file.Host) h.Target = new(file.Target) h.Id = int(file.GetCsvDb().GetHostId()) h.Host = arr[0] h.Target.TargetStr = arr[1] h.HeaderChange = arr[2] h.HostChange = arr[3] h.Remark = arr[4] h.Location = arr[5] h.Scheme = arr[6] if h.Scheme == "" { h.Scheme = "all" } h.Flow = new(file.Flow) h.NoStore = true } return } //send task info func (s *Conn) SendConfigInfo(c *config.CommonConfig) (int, error) { /* The task info is formed as follows: +----+-----+---------+ |type| len | content | +----+---------------+ | 4 | 4 | ... | +----+---------------+ */ raw := bytes.NewBuffer([]byte{}) binary.Write(raw, binary.LittleEndian, []byte(common.NEW_CONF)) common.BinaryWrite(raw, c.Cnf.U, c.Cnf.P, common.GetStrByBool(c.Cnf.Crypt), common.GetStrByBool(c.Cnf.Compress), strconv.Itoa(c.Client.RateLimit), strconv.Itoa(int(c.Client.Flow.FlowLimit)), strconv.Itoa(c.Client.MaxConn), c.Client.Remark) return s.Write(raw.Bytes()) } //get task info func (s *Conn) GetConfigInfo() (c *file.Client, err error) { var l int buf := pool.BufPoolMax.Get().([]byte) defer pool.PutBufPoolMax(buf) if l, err = s.GetLen(); err != nil { return } else if _, err = s.ReadLen(l, buf); err != nil { return } else { arr := strings.Split(string(buf[:l]), common.CONN_DATA_SEQ) c = file.NewClient("", true, false) c.Cnf.U = arr[0] c.Cnf.P = arr[1] c.Cnf.Crypt = common.GetBoolByStr(arr[2]) c.Cnf.Compress = common.GetBoolByStr(arr[3]) c.RateLimit = common.GetIntNoErrByStr(arr[4]) c.Flow.FlowLimit = int64(common.GetIntNoErrByStr(arr[5])) c.MaxConn = common.GetIntNoErrByStr(arr[6]) c.Remark = arr[7] } return } //send task info func (s *Conn) SendTaskInfo(t *file.Tunnel) (int, error) { /* The task info is formed as follows: +----+-----+---------+ |type| len | content | +----+---------------+ | 4 | 4 | ... | +----+---------------+ */ raw := bytes.NewBuffer([]byte{}) binary.Write(raw, binary.LittleEndian, []byte(common.NEW_TASK)) common.BinaryWrite(raw, t.Mode, t.Ports, t.Target.TargetStr, t.Remark, t.TargetAddr, t.Password, t.LocalPath, t.StripPre, t.ServerIp) return s.Write(raw.Bytes()) } //get task info func (s *Conn) GetTaskInfo() (t *file.Tunnel, err error) { var l int buf := pool.BufPoolMax.Get().([]byte) defer pool.PutBufPoolMax(buf) if l, err = s.GetLen(); err != nil { return } else if _, err = s.ReadLen(l, buf); err != nil { return } else { arr := strings.Split(string(buf[:l]), common.CONN_DATA_SEQ) t = new(file.Tunnel) t.Target = new(file.Target) t.Mode = arr[0] t.Ports = arr[1] t.Target.TargetStr = arr[2] t.Id = int(file.GetCsvDb().GetTaskId()) t.Status = true t.Flow = new(file.Flow) t.Remark = arr[3] t.TargetAddr = arr[4] t.Password = arr[5] t.LocalPath = arr[6] t.StripPre = arr[7] if len(arr) > 8 { t.ServerIp = arr[8] } t.NoStore = true } return } //close func (s *Conn) Close() error { return s.Conn.Close() } //write func (s *Conn) Write(b []byte) (int, error) { return s.Conn.Write(b) } //read func (s *Conn) Read(b []byte) (int, error) { return s.Conn.Read(b) } //write sign flag func (s *Conn) WriteClose() (int, error) { return s.Write([]byte(common.RES_CLOSE)) } //write main func (s *Conn) WriteMain() (int, error) { return s.Write([]byte(common.WORK_MAIN)) } //write main func (s *Conn) WriteConfig() (int, error) { return s.Write([]byte(common.WORK_CONFIG)) } //write chan func (s *Conn) WriteChan() (int, error) { return s.Write([]byte(common.WORK_CHAN)) } //get the assembled amount data(len 4 and content) func GetLenBytes(buf []byte) (b []byte, err error) { raw := bytes.NewBuffer([]byte{}) if err = binary.Write(raw, binary.LittleEndian, int32(len(buf))); err != nil { return } if err = binary.Write(raw, binary.LittleEndian, buf); err != nil { return } b = raw.Bytes() return } //udp connection setting func SetUdpSession(sess *kcp.UDPSession) { sess.SetStreamMode(true) sess.SetWindowSize(1024, 1024) sess.SetReadBuffer(64 * 1024) sess.SetWriteBuffer(64 * 1024) sess.SetNoDelay(1, 10, 2, 1) sess.SetMtu(1600) sess.SetACKNoDelay(true) sess.SetWriteDelay(false) } //conn1 mux conn func CopyWaitGroup(conn1, conn2 net.Conn, crypt bool, snappy bool, rate *rate.Rate, flow *file.Flow, isServer bool, rb []byte) { var in, out int64 var wg sync.WaitGroup connHandle := GetConn(conn1, crypt, snappy, rate, isServer) if rb != nil { connHandle.Write(rb) } go func(in *int64) { wg.Add(1) *in, _ = common.CopyBuffer(connHandle, conn2) connHandle.Close() conn2.Close() wg.Done() }(&in) out, _ = common.CopyBuffer(conn2, connHandle) connHandle.Close() conn2.Close() wg.Wait() if flow != nil { flow.Add(in, out) } } //get crypt or snappy conn func GetConn(conn net.Conn, cpt, snappy bool, rt *rate.Rate, isServer bool) (io.ReadWriteCloser) { if cpt { if isServer { return rate.NewRateConn(crypt.NewTlsServerConn(conn), rt) } return rate.NewRateConn(crypt.NewTlsClientConn(conn), rt) } else if snappy { return rate.NewRateConn(NewSnappyConn(conn), rt) } return rate.NewRateConn(conn, rt) }