|
@@ -15,9 +15,12 @@
|
|
package net
|
|
package net
|
|
|
|
|
|
import (
|
|
import (
|
|
|
|
+ "bytes"
|
|
|
|
+ "errors"
|
|
"fmt"
|
|
"fmt"
|
|
"io"
|
|
"io"
|
|
"net"
|
|
"net"
|
|
|
|
+ "sync"
|
|
"time"
|
|
"time"
|
|
|
|
|
|
"github.com/fatedier/frp/utils/log"
|
|
"github.com/fatedier/frp/utils/log"
|
|
@@ -48,6 +51,13 @@ type WrapReadWriteCloserConn struct {
|
|
log.Logger
|
|
log.Logger
|
|
}
|
|
}
|
|
|
|
|
|
|
|
+func WrapReadWriteCloserToConn(rwc io.ReadWriteCloser) Conn {
|
|
|
|
+ return &WrapReadWriteCloserConn{
|
|
|
|
+ ReadWriteCloser: rwc,
|
|
|
|
+ Logger: log.NewPrefixLogger(""),
|
|
|
|
+ }
|
|
|
|
+}
|
|
|
|
+
|
|
func (conn *WrapReadWriteCloserConn) LocalAddr() net.Addr {
|
|
func (conn *WrapReadWriteCloserConn) LocalAddr() net.Addr {
|
|
return (*net.TCPAddr)(nil)
|
|
return (*net.TCPAddr)(nil)
|
|
}
|
|
}
|
|
@@ -57,47 +67,15 @@ func (conn *WrapReadWriteCloserConn) RemoteAddr() net.Addr {
|
|
}
|
|
}
|
|
|
|
|
|
func (conn *WrapReadWriteCloserConn) SetDeadline(t time.Time) error {
|
|
func (conn *WrapReadWriteCloserConn) SetDeadline(t time.Time) error {
|
|
- return nil
|
|
|
|
|
|
+ return &net.OpError{Op: "set", Net: "wrap", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
|
|
}
|
|
}
|
|
|
|
|
|
func (conn *WrapReadWriteCloserConn) SetReadDeadline(t time.Time) error {
|
|
func (conn *WrapReadWriteCloserConn) SetReadDeadline(t time.Time) error {
|
|
- return nil
|
|
|
|
|
|
+ return &net.OpError{Op: "set", Net: "wrap", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
|
|
}
|
|
}
|
|
|
|
|
|
func (conn *WrapReadWriteCloserConn) SetWriteDeadline(t time.Time) error {
|
|
func (conn *WrapReadWriteCloserConn) SetWriteDeadline(t time.Time) error {
|
|
- return nil
|
|
|
|
-}
|
|
|
|
-
|
|
|
|
-func WrapReadWriteCloserToConn(rwc io.ReadWriteCloser) Conn {
|
|
|
|
- return &WrapReadWriteCloserConn{
|
|
|
|
- ReadWriteCloser: rwc,
|
|
|
|
- Logger: log.NewPrefixLogger(""),
|
|
|
|
- }
|
|
|
|
-}
|
|
|
|
-
|
|
|
|
-type Listener interface {
|
|
|
|
- Accept() (Conn, error)
|
|
|
|
- Close() error
|
|
|
|
- log.Logger
|
|
|
|
-}
|
|
|
|
-
|
|
|
|
-type LogListener struct {
|
|
|
|
- l net.Listener
|
|
|
|
- net.Listener
|
|
|
|
- log.Logger
|
|
|
|
-}
|
|
|
|
-
|
|
|
|
-func WrapLogListener(l net.Listener) Listener {
|
|
|
|
- return &LogListener{
|
|
|
|
- l: l,
|
|
|
|
- Listener: l,
|
|
|
|
- Logger: log.NewPrefixLogger(""),
|
|
|
|
- }
|
|
|
|
-}
|
|
|
|
-
|
|
|
|
-func (logL *LogListener) Accept() (Conn, error) {
|
|
|
|
- c, err := logL.l.Accept()
|
|
|
|
- return WrapConn(c), err
|
|
|
|
|
|
+ return &net.OpError{Op: "set", Net: "wrap", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
|
|
}
|
|
}
|
|
|
|
|
|
func ConnectServer(protocol string, addr string) (c Conn, err error) {
|
|
func ConnectServer(protocol string, addr string) (c Conn, err error) {
|
|
@@ -136,3 +114,45 @@ func ConnectServerByHttpProxy(httpProxy string, protocol string, addr string) (c
|
|
return nil, fmt.Errorf("unsupport protocol: %s", protocol)
|
|
return nil, fmt.Errorf("unsupport protocol: %s", protocol)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
+
|
|
|
|
+type SharedConn struct {
|
|
|
|
+ Conn
|
|
|
|
+ sync.Mutex
|
|
|
|
+ buf *bytes.Buffer
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// the bytes you read in io.Reader, will be reserved in SharedConn
|
|
|
|
+func NewShareConn(conn Conn) (*SharedConn, io.Reader) {
|
|
|
|
+ sc := &SharedConn{
|
|
|
|
+ Conn: conn,
|
|
|
|
+ buf: bytes.NewBuffer(make([]byte, 0, 1024)),
|
|
|
|
+ }
|
|
|
|
+ return sc, io.TeeReader(conn, sc.buf)
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+func (sc *SharedConn) Read(p []byte) (n int, err error) {
|
|
|
|
+ sc.Lock()
|
|
|
|
+ if sc.buf == nil {
|
|
|
|
+ sc.Unlock()
|
|
|
|
+ return sc.Conn.Read(p)
|
|
|
|
+ }
|
|
|
|
+ sc.Unlock()
|
|
|
|
+ n, err = sc.buf.Read(p)
|
|
|
|
+
|
|
|
|
+ if err == io.EOF {
|
|
|
|
+ sc.Lock()
|
|
|
|
+ sc.buf = nil
|
|
|
|
+ sc.Unlock()
|
|
|
|
+ var n2 int
|
|
|
|
+ n2, err = sc.Conn.Read(p[n:])
|
|
|
|
+
|
|
|
|
+ n += n2
|
|
|
|
+ }
|
|
|
|
+ return
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+func (sc *SharedConn) WriteBuff(buffer []byte) (err error) {
|
|
|
|
+ sc.buf.Reset()
|
|
|
|
+ _, err = sc.buf.Write(buffer)
|
|
|
|
+ return err
|
|
|
|
+}
|